Refactor & added model code