refactored a little bit of code :)