Tidied up the code

Tidied up the code