Cleaned up code and added required functionality.