changed the code to oop