Tidy up jQuery script