Bug fix and changed UI