Did some more work on class diagram