This example demonstrates the hierarchical visualization capabilities of Paroto, showing how primary design parameters flow through models to constraints.

Features demonstrated: - Automatic hierarchy extraction from OpenMDAO problems - Multiple visualization formats (Mermaid, NetworkX, Cytoscape) - Parameter level identification - Graph analysis and statistics

import openmdao.api as om

from paroto.viz import (
    analyze_hierarchy,
    export_graphml,
    export_json,
    extract_hierarchy,
    generate_hierarchical_graph,
    generate_simple_graph,
    to_networkx,
)


def create_example_problem():
    """Create example optimization problem with parameter hierarchy.

    This creates a simple problem with:
    - 3 primary parameters (design variables)
    - 2 intermediate models
    - 2 constraints

    Returns
    -------
    om.Problem
        Configured OpenMDAO problem
    """
    prob = om.Problem()

    # Design variables (primary parameters)
    indeps = prob.model.add_subsystem("design", om.IndepVarComp(), promotes=["*"])
    indeps.add_output("length", val=10.0, units="m")
    indeps.add_output("width", val=5.0, units="m")
    indeps.add_output("height", val=3.0, units="m")

    # Intermediate model: compute volume
    prob.model.add_subsystem(
        "volume_calc",
        om.ExecComp(
            "volume = length * width * height",
            length={"units": "m"},
            width={"units": "m"},
            height={"units": "m"},
            volume={"units": "m**3"},
        ),
        promotes=["*"],
    )

    # Intermediate model: compute surface area
    prob.model.add_subsystem(
        "area_calc",
        om.ExecComp(
            "surface_area = 2 * (length*width + length*height + width*height)",
            length={"units": "m"},
            width={"units": "m"},
            height={"units": "m"},
            surface_area={"units": "m**2"},
        ),
        promotes=["*"],
    )

    # Constraint 1: Volume must be at least 100 m³
    prob.model.add_subsystem(
        "volume_constraint",
        om.ExecComp("volume_margin = volume - 100.0", units="m**3"),
        promotes=["*"],
    )

    # Constraint 2: Surface area must not exceed 200 m²
    prob.model.add_subsystem(
        "area_constraint",
        om.ExecComp("area_margin = 200.0 - surface_area", units="m**2"),
        promotes=["*"],
    )

    prob.setup()
    prob.run_model()

    return prob


def main():
    """Run the hierarchical visualization example."""
    print("=" * 80)
    print("Hierarchical Parameter Visualization - Gallery Example")
    print("=" * 80)
    print()

    # Step 1: Create problem
    print("Step 1: Creating example optimization problem...")
    prob = create_example_problem()
    print("  [OK] Problem with 3 parameters, 2 models, 2 constraints")
    print()

    # Step 2: Extract hierarchy
    print("Step 2: Extracting parameter hierarchy...")
    primary_params = ["length", "width", "height"]
    hierarchy = extract_hierarchy(prob, primary_params)

    print(f"  [OK] Found {len(hierarchy['nodes'])} nodes")
    print(f"  [OK] Found {len(hierarchy['edges'])} edges")
    print(f"  [OK] Identified {len(hierarchy['primary_params'])} primary parameters")
    print()

    # Step 3: Analyze hierarchy
    print("Step 3: Analyzing hierarchy structure...")
    graph = to_networkx(hierarchy)
    stats = analyze_hierarchy(graph)

    print(f"  Levels in hierarchy: {stats['num_levels']}")
    print(f"  Nodes by level: {dict(sorted(stats['nodes_by_level'].items()))}")
    print(f"  Nodes by type: {stats['nodes_by_type']}")
    print()

    # Step 4: Generate visualizations
    print("Step 4: Generating visualizations...")

    # Simple all-to-all diagram
    simple_file = generate_simple_graph(
        parameters=["length", "width", "height"],
        models=["Volume Calculator", "Surface Area Calculator"],
        constraints=["Volume >= 100 m³", "Area <= 200 m²"],
        output_path="gallery_simple_diagram.mmd",
    )
    print(f"  [OK] Simple diagram: {simple_file}")

    # Hierarchical diagram showing actual levels
    hier_file = generate_hierarchical_graph(
        hierarchy, "gallery_hierarchical_diagram.mmd", orientation="TD"
    )
    print(f"  [OK] Hierarchical diagram: {hier_file}")

    # NetworkX exports for external tools
    graphml_file = export_graphml(graph, "gallery_graph.graphml")
    json_file = export_json(graph, "gallery_graph.json")
    print(f"  [OK] GraphML export: {graphml_file}")
    print(f"  [OK] JSON export: {json_file}")
    print()

    print("=" * 80)
    print("Gallery example complete!")
    print()
    print("Visualization files created:")
    print(f"  - {simple_file} (simple all-to-all)")
    print(f"  - {hier_file} (hierarchical levels)")
    print(f"  - {graphml_file} (for yEd, Gephi)")
    print(f"  - {json_file} (for D3.js, web apps)")
    print()
    print("Next steps:")
    print("  1. Open .mmd files at https://mermaid.live/")
    print("  2. Open .graphml in yEd for advanced layout editing")
    print("  3. Use .json for custom web visualizations")
    print("=" * 80)


if __name__ == "__main__":
    main()