Clean up of code