Changed to follow MVVM and some refactoring