Fixes and clean up of UI