Cleaned up code, closer to squashing bug