Bugfix
Tidied the code