Fixed #6 and did some code cleanup