refactored the class