Refactor for modularity