Refactor. More code