refactored and cleaned up the code, thanks @alexnederlof