Big restructure of UI code