Adding a class diagram, fixing a few bugs