refactoring and clean up code :)