Repair to suit new model
Merge pull request #23 from ilovearsenal93/repair

refactor latter