#!/usr/bin/env python3
"""
Demo script showing DAG visualization integration with QA_LLM_Module
"""

import json
from dag_visualizer import DAGVisualizer

def create_demo_task_data():
    """Create sample task data for demonstration"""
    return {
        "tasks": [
            {
                "task": "excavate_soil_from_pile",
                "instruction_function": {
                    "name": "excavate_soil_from_pile",
                    "robot_ids": ["robot_excavator_01"],
                    "dependencies": [],
                    "object_keywords": ["soil_pile"]
                }
            },
            {
                "task": "transport_soil_to_pit",
                "instruction_function": {
                    "name": "transport_soil_to_pit",
                    "robot_ids": ["robot_dump_truck_01"],
                    "dependencies": ["excavate_soil_from_pile"],
                    "object_keywords": ["pit"]
                }
            },
            {
                "task": "avoid_puddle_areas",
                "instruction_function": {
                    "name": "avoid_areas_for_all_robots",
                    "robot_ids": ["robot_excavator_01", "robot_dump_truck_01"],
                    "dependencies": [],
                    "object_keywords": ["puddle1", "puddle2"]
                }
            },
            {
                "task": "coordinate_dump_operation",
                "instruction_function": {
                    "name": "coordinate_dump_operation",
                    "robot_ids": ["robot_dump_truck_01", "robot_excavator_01"],
                    "dependencies": ["transport_soil_to_pit", "avoid_puddle_areas"],
                    "object_keywords": ["pit", "dumping_zone"]
                }
            }
        ]
    }

def main():
    print("🎯 DAG Visualization Demo for QA_LLM_Module")
    print("=" * 50)
    
    # Create sample task data
    task_data = create_demo_task_data()
    print("Sample task data:")
    print(json.dumps(task_data, indent=2))
    print()
    
    # Initialize visualizer
    visualizer = DAGVisualizer()
    
    # Create detailed visualization
    print("Creating detailed DAG visualization...")
    detailed_image = visualizer.create_dag_visualization(
        task_data, 
        title="DART-LLM: Robot Task Dependency Graph"
    )
    
    if detailed_image:
        print(f"✅ Detailed visualization saved: {detailed_image}")
    else:
        print("❌ Failed to create detailed visualization")
    
    # Create simplified visualization
    print("\nCreating simplified DAG visualization...")
    simple_image = visualizer.create_simplified_dag_visualization(
        task_data,
        title="DART-LLM: Task Dependencies"
    )
    
    if simple_image:
        print(f"✅ Simplified visualization saved: {simple_image}")
    else:
        print("❌ Failed to create simplified visualization")
    
    print("\n" + "=" * 50)
    print("🚀 Integration Instructions:")
    print("1. Install requirements: pip install -r requirements.txt")
    print("2. Run the main application: python3 main.py")
    print("3. Enter a query that generates task JSON")
    print("4. The DAG visualization will appear automatically!")
    print("\n📝 Example queries:")
    print("- 'Move half of the soil from the soil pile to the pit.'")
    print("- 'Excavate material and transport it while avoiding obstacles.'")

if __name__ == "__main__":
    main()