Refactoring code

Some changes were made in code structure