UI fixes & improved functionality