UI Cleanup bug fixes