Added UI class, minor bugs fixed.