cleaned the code up