Tidied up, added comments