some clean up of code