Cleaned up code a bit, and creating pull request