Adapted code to new UI. Major progress on #4.