Tidied some code