UI bug fix and improvements