Major refactoring to generate_graph method