refactored and cleanted up code