Fixing UI bugs