Fixes and UI improvements