Bug fix and UI