Made some small changes to the class. Will add comments next.