Added models
Fixed fucked up code!
Refazctored