Code cleanup
Adding more to the UI