Refactor and cleanup
Clean up