changed code structure to class definition