Bug fix and models

Some bugs and made model class