Refactored code. Must test
Refactored code works