Refactor Code

Make the code a little bit nicer