Cleaned up code a bit...