Bug fixes based on UI