UI/UX fixes