Bug fixes and small ui updates