Ui code cleanup