Fixed messy logic in model class