UI fixes and cleanup