Cleaned up code and readme