Refactor model

refactor model, add unit tests