Sprint 2: Refactored all code