Cleaned up code
Fixes #1