Cleaned up code and updated UI