Cleaned up code, fixed a bug.