clean up code
altering to a nicer structure