Refactoring and model improvements