Major clean up of code