"""
Reasoning Graph Visualization

Generates visual representations of reasoning graphs from:
- CoT traces (extracted as sequential/self-correcting graphs)
- AGoT graph structures
- Comparison charts across methods and models

Requires: matplotlib, networkx
"""

from __future__ import annotations
import io
import json
import math
import os
import sys
from pathlib import Path
from typing import Optional

# Add project root to path
_project_root = Path(__file__).resolve().parent.parent.parent
if str(_project_root) not in sys.path:
    sys.path.insert(0, str(_project_root))

try:
    import matplotlib
    matplotlib.use("Agg")  # non-interactive backend
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    from matplotlib.lines import Line2D
    HAS_MPL = True
except ImportError:
    HAS_MPL = False

try:
    import networkx as nx
    HAS_NX = True
except ImportError:
    HAS_NX = False

import numpy as np

from src.graph_analysis.reasoning_graph import (
    build_reasoning_graph,
    build_graph_from_agot,
    compute_properties,
    compute_small_world_index,
)


def plot_reasoning_graph(
    text_or_graph,
    title: str = "Reasoning Graph",
    method: str = "self_correct",
    save_path: Optional[str] = None,
    figsize=(10, 8),
):
    """Plot a reasoning graph (from text or AGoT graph object)."""
    if not HAS_MPL or not HAS_NX:
        print("matplotlib and networkx required for visualization.")
        return

    if isinstance(text_or_graph, str):
        G = build_reasoning_graph(text_or_graph, method=method)
    else:
        # Assume AGoT graph
        G = build_graph_from_agot(text_or_graph)

    if G.number_of_nodes() == 0:
        print("Empty graph — nothing to visualize.")
        return

    plt.figure(figsize=figsize)

    pos = nx.spring_layout(G, k=2, iterations=50, seed=42)

    # Draw nodes
    node_colors = []
    for node in G.nodes():
        data = G.nodes[node]
        if isinstance(node, int) or "_sub" not in str(node):
            node_colors.append("skyblue")
        else:
            node_colors.append("lightgreen")

    nx.draw_networkx_nodes(G, pos, node_color=node_colors,
                           node_size=800, alpha=0.9)
    nx.draw_networkx_labels(G, pos, font_size=9, font_weight="bold")

    # Draw edges with different styles
    forward_edges = [(u, v) for u, v, d in G.edges(data=True)
                     if d.get("label") == "forward" or not d]
    back_edges = [(u, v) for u, v, d in G.edges(data=True)
                  if d.get("label") in ("backtrack", "re-evaluate")]

    nx.draw_networkx_edges(G, pos, edgelist=forward_edges,
                           edge_color="gray", arrows=True, alpha=0.6)
    nx.draw_networkx_edges(G, pos, edgelist=back_edges,
                           edge_color="red", style="dashed",
                           arrows=True, alpha=0.7)

    plt.title(title, fontsize=14, pad=20)
    plt.axis("off")

    # Legend
    legend_elements = [
        mpatches.Patch(color="skyblue", label="Reasoning step"),
        mpatches.Patch(color="lightgreen", label="Sub-graph node"),
        Line2D([0], [0], color="gray", label="Forward reasoning"),
        Line2D([0], [0], color="red", linestyle="--", label="Backtrack/revise"),
    ]
    plt.legend(handles=legend_elements, loc="upper right")

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Saved: {save_path}")
    else:
        plt.show()
    plt.close()


def plot_comparison_chart(
    results: dict,
    title: str = "Method Comparison",
    save_path: Optional[str] = None,
):
    """Plot a bar chart comparing accuracy across methods."""
    if not HAS_MPL:
        print("matplotlib required for visualization.")
        return

    methods = list(results.keys())
    accuracies = [results[m].get("accuracy", 0) for m in methods]
    times = [results[m].get("avg_time_s", 0) for m in methods]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    # Accuracy bar chart
    colors = ["#4CAF50", "#2196F3", "#FF9800", "#9C27B0"]
    bars1 = ax1.bar(methods, accuracies, color=colors[:len(methods)])
    ax1.set_ylabel("Accuracy")
    ax1.set_title(f"{title} — Accuracy", fontsize=12)
    ax1.set_ylim(0, 1.0)
    ax1.axhline(y=0.5, color="gray", linestyle="--", alpha=0.3)

    for bar, acc in zip(bars1, accuracies):
        ax1.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
                 f"{acc:.1%}", ha="center", va="bottom", fontsize=10)

    # Time comparison
    bars2 = ax2.bar(methods, times, color=colors[:len(methods)])
    ax2.set_ylabel("Avg Time (seconds)")
    ax2.set_title(f"{title} — Time per Sample", fontsize=12)

    for bar, t in zip(bars2, times):
        ax2.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.5,
                 f"{t:.1f}s", ha="center", va="bottom", fontsize=10)

    plt.tight_layout()
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Saved: {save_path}")
    else:
        plt.show()
    plt.close()


def plot_scale_comparison(
    model_results: dict,
    metric: str = "accuracy",
    title: str = "Model Scale Comparison",
    save_path: Optional[str] = None,
):
    """Plot accuracy vs model size for different methods."""
    if not HAS_MPL:
        return

    fig, ax = plt.subplots(figsize=(10, 6))

    markers = {"IO": "o", "CoT": "s", "CoT-SC": "^", "ToT": "D", "AGoT": "*"}
    colors = {"IO": "#4CAF50", "CoT": "#2196F3", "CoT-SC": "#FF9800",
              "ToT": "#9C27B0", "AGoT": "#E91E63"}

    # Group by method
    method_data = {}
    for model_name, tasks in model_results.items():
        # Extract model size from name
        size = 14 if "14b" in model_name else 32
        for task, methods in tasks.items():
            if isinstance(methods, dict):
                for method, res in methods.items():
                    key = method.split("(")[0].strip()  # handle "CoT-SC (n=5)"
                    if key not in method_data:
                        method_data[key] = {"sizes": [], "values": []}
                    val = res.get(metric, 0) if isinstance(res, dict) else 0
                    method_data[key]["sizes"].append(size)
                    method_data[key]["values"].append(val)

    for method, data in method_data.items():
        if data["sizes"]:
            ax.scatter(data["sizes"], data["values"],
                       marker=markers.get(method, "o"),
                       color=colors.get(method, "gray"),
                       s=100, label=method, zorder=5)
            # Connect points
            if len(data["sizes"]) > 1:
                order = np.argsort(data["sizes"])
                ax.plot(np.array(data["sizes"])[order],
                        np.array(data["values"])[order],
                        color=colors.get(method, "gray"),
                        linestyle="--", alpha=0.3)

    ax.set_xlabel("Model Size (B parameters)")
    ax.set_ylabel(metric.replace("_", " ").title())
    ax.set_title(title)
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Saved: {save_path}")
    plt.close()


def plot_graph_property_comparison(
    analysis_results: list,
    save_path: Optional[str] = None,
):
    """Plot graph property differences between correct and incorrect answers."""
    if not HAS_MPL:
        return

    correct = [r for r in analysis_results if r.get("correct")]
    incorrect = [r for r in analysis_results if not r.get("correct")]

    fig, axes = plt.subplots(2, 3, figsize=(14, 8))

    properties = [
        ("properties", "n_nodes", "Node Count"),
        ("properties", "n_cycles", "Cycle Count"),
        ("properties", "diameter", "Graph Diameter"),
        ("small_world", "small_world_index", "Small-World Index"),
        ("properties", "avg_clustering", "Avg Clustering"),
        ("properties", "avg_path_length", "Avg Path Length"),
    ]

    for idx, (section, key, label) in enumerate(properties):
        ax = axes[idx // 3][idx % 3]

        c_vals = [r.get(section, {}).get(key, 0) for r in correct
                  if isinstance(r.get(section, {}).get(key, (int, float)), (int, float))]
        i_vals = [r.get(section, {}).get(key, 0) for r in incorrect
                  if isinstance(r.get(section, {}).get(key, (int, float)), (int, float))]

        if c_vals:
            ax.boxplot([c_vals, i_vals] if i_vals else [c_vals],
                       labels=["Correct", "Incorrect"][:len([c_vals, i_vals])]
                       if i_vals else ["Correct"])
        elif i_vals:
            ax.boxplot([i_vals], labels=["Incorrect"])

        ax.set_ylabel(label)
        ax.set_title(f"{label} by Correctness")
        ax.grid(True, alpha=0.3)

    plt.tight_layout()
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Saved: {save_path}")
    plt.close()


# ── CLI ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Visualize reasoning graphs")
    parser.add_argument("--results", type=str, help="JSON results file to analyze & plot")
    parser.add_argument("--output-dir", type=str, default="data/results/figures",
                        help="Output directory for figures")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    if args.results:
        # Load results and plot comparison
        with open(args.results) as f:
            data = json.load(f)

        summary = data.get("summary", {})
        details = data.get("details", [])

        # Plot graph properties vs correctness
        from src.graph_analysis.reasoning_graph import analyze_reasoning_trace
        traces = []
        for d in details:
            if "full_response" in d:
                analysis = analyze_reasoning_trace(
                    d["full_response"], label=str(d.get("question", "")[:40])
                )
                analysis["correct"] = d.get("correct", False)
                traces.append(analysis)

        if traces:
            fig_path = output_dir / f"{summary.get('method', 'analysis')}_properties.png"
            plot_graph_property_comparison(traces, save_path=str(fig_path))

        # Plot a sample reasoning graph
        for d in details[:3]:
            if "full_response" in d:
                q_short = d.get("question", "sample")[:30].replace(" ", "_")
                fig_path = output_dir / f"graph_{q_short}.png"
                plot_reasoning_graph(
                    d["full_response"],
                    title=f"Reasoning: {d.get('question', '')[:40]}...",
                    save_path=str(fig_path),
                )

        print(f"\nFigures saved to: {output_dir}/")
    else:
        # Demo
        print("Generating demo visualization...")
        sample_text = """First, calculate the total: 16 eggs per day.
        She eats 3 for breakfast and uses 4 for muffins.
        So 3+4=7 eggs used. Remaining = 16-7=9 eggs.
        9 eggs = 9/12 = 0.75 dozen. At $2/dozen = $1.50.
        Wait, let me double-check: 16-7 is indeed 9. 9/12=0.75. 0.75*2=1.50.
        That seems right. But what if she also uses eggs for something else?
        No, the problem only mentions breakfast and muffins.
        Final answer: $1.50."""
        plot_reasoning_graph(
            sample_text,
            title="Demo Reasoning Graph",
            save_path=str(output_dir / "demo_graph.png"),
        )
        print(f"Demo saved to: {output_dir}/")
