Cleaned up code. Made it simple and work