Bug fixes and minor UI improvements