Refactor

Moved methods to appropriate classes and simplified code