Cleaned up and improved UI