ViewModel refactor