Fix ui bugs