Bug fix and UI update