Changed classes and code, added functionality