Tidied up code, fixed a string comparison issue.