Refactoring some code

Refactoring code more to follow