Reorganize the new UI code a bit