Refectored my code

Did small changes to and tested my changes it all works