Improved UI and corrected bugs