Cleaned some code after initial code review.