Refactored and moved code around

Refactor