UI code cleanup