Cleanup. Refactoring codebase