Improve highlighting

Hopefully addresses #7
I only changed the regular highlighting, I haven't tested the other versions.

There are a couple of whitespace changes in there as well, if that's OK, let me know