Refactored the "UI" code a bit.