Fixing bugs on new UI structure