Code cleanup

Just some simple code clean up for better overview