Refactor model classes and fix bug.