Clean up code
Forgot to add new files to last commit