Small improvements and UI fixes