cleaned up the code

cleaned up the code