Changed/added some more UI code