Cleaned up code
ditto