Code clean-up and UI enhancements