Added ViewModels and updated controller