Cleaned up your code a little bit.