Fixed few bugs regarding the new UI