UI code clean up