Update project structure and clean up code