"""
Reasoning Graph Extraction and Analysis

Implements the methodology from "Topology of Reasoning" (NeurIPS 2025):
1. Extract hidden states or key reasoning steps from model output
2. Build a reasoning graph (nodes = reasoning steps, edges = transitions)
3. Compute graph properties:
   - Cyclicity (loop detection rate, number of cycles per sample)
   - Graph diameter (longest shortest path)
   - Small-world index S = (C/C_rand) / (L/L_rand)

This works at the semantic level on the output text rather than
hidden states (which require model-internal access).
"""

from __future__ import annotations
import json
import math
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Optional

import networkx as nx
import numpy as np

# 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))


# ── Reasoning Step Extraction ───────────────────────────────────────

def extract_reasoning_steps(text: str) -> list[str]:
    """
    Split a reasoning trace into individual steps.
    Heuristics: look for sentence boundaries, numbered lists, step markers.
    """
    # Split on common step markers
    steps = []

    # Try numbered list pattern
    numbered = re.findall(r'(?:^|\n)\s*(?:\d+[\.\)]|[-\*])\s*([^\n]+)', text)
    if len(numbered) >= 2:
        return [s.strip() for s in numbered]

    # Try explicit step markers
    for marker in ["Step ", "Step ", "First,", "Second,", "Third,",
                    "Then,", "Next,", "Finally,", "Therefore,",
                    "首先", "其次", "然后", "最后"]:
        pass  # handled below

    # Split on sentence boundaries (period, newline)
    sentences = re.split(r'(?:\.|\!|\?)\s+', text)
    sentences = [s.strip() for s in sentences if len(s.strip()) > 20]

    if len(sentences) >= 2:
        return sentences

    # Fall back to newline splitting
    lines = [l.strip() for l in text.split("\n") if len(l.strip()) > 20]
    if lines:
        return lines

    return [text]


def extract_thinking_segments(text: str) -> list[str]:
    """
    Extract segments between "wait", "but", "however", "alternatively",
    "let me check", "let me re-think", etc. — these indicate back-tracking
    or alternative reasoning paths that form graph cycles.
    """
    segments = []
    # Split at self-correction markers
    markers = r'(?:Wait|However|But|Let me check|Let me re(?:-|)think|'
    markers += r'Alternatively|Hmm|Actually|No,|重新检查|等等|不对)'
    segments = re.split(markers, text, flags=re.IGNORECASE)
    return [s.strip() for s in segments if len(s.strip()) > 30]


# ── Reasoning Graph Building ────────────────────────────────────────

def build_reasoning_graph(
    text: str,
    method: str = "sequential",
) -> nx.DiGraph:
    """
    Build a directed graph from a reasoning trace.

    Args:
        text: The model output text
        method: "sequential" (default) — each step connects to the next
                "self_correct" — adds back-edges for correction segments

    Returns:
        networkx.DiGraph with node attributes: content, index
    """
    G = nx.DiGraph()

    if method == "sequential":
        steps = extract_reasoning_steps(text)
        for i, step in enumerate(steps):
            G.add_node(i, content=step[:100], index=i)
            if i > 0:
                G.add_edge(i - 1, i, label="forward")

    elif method == "self_correct":
        segments = extract_thinking_segments(text)
        for i, seg in enumerate(segments):
            G.add_node(i, content=seg[:100], index=i)
            if i > 0:
                # Forward edge (always present)
                G.add_edge(i - 1, i, label="forward")
                # Back-edge: correction segments loop back to earlier reasoning
                if any(marker in seg.lower() for marker in
                       ["wait", "however", "but", "actually", "check",
                        "rethink", "re-think", "hmm", "no,"]):
                    # Connect back to the first segment (re-evaluating initial approach)
                    if i >= 2:
                        G.add_edge(i, 0, label="backtrack")
                        G.add_edge(0, i, label="re-evaluate")

    return G


def build_graph_from_agot(agot_graph) -> nx.DiGraph:
    """
    Convert an AGoT graph into a NetworkX DiGraph for property analysis.
    """
    G = nx.DiGraph()
    for node in agot_graph.nodes:
        G.add_node(node["id"], content=node["content"][:100],
                   score=node.get("score", 0), depth=node.get("depth", 0))
    for edge in agot_graph.edges:
        G.add_edge(edge["from"], edge["to"], label=edge.get("label", ""))
    # Add nodes/edges from sub-graphs
    for sg in getattr(agot_graph, "sub_graphs", {}).values():
        for node in sg.nodes:
            G.add_node(f"{node['id']}_sub", content=node["content"][:100],
                       score=node.get("score", 0), depth=node.get("depth", 0))
        for edge in sg.edges:
            G.add_edge(f"{edge['from']}_sub", f"{edge['to']}_sub",
                       label=edge.get("label", ""))
    return G


# ── Graph Properties ────────────────────────────────────────────────

def compute_properties(G: nx.DiGraph) -> dict:
    """
    Compute all reasoning graph properties.
    """
    props = {}

    # Basic stats
    props["n_nodes"] = G.number_of_nodes()
    props["n_edges"] = G.number_of_edges()
    props["density"] = nx.density(G) if G.number_of_nodes() > 1 else 0.0

    # Average degree
    degrees = [d for _, d in G.degree()]
    props["avg_degree"] = float(np.mean(degrees)) if degrees else 0.0
    props["max_degree"] = max(degrees) if degrees else 0

    # ── Cyclicity ──────────────────────────────────────────────
    # Count simple cycles (limited to avoid excessive computation)
    props["has_cycle"] = False
    props["n_cycles"] = 0

    try:
        cycles = list(nx.simple_cycles(G))
        props["has_cycle"] = len(cycles) > 0
        props["n_cycles"] = len(cycles)
    except (nx.NetworkXNoCycle, nx.NetworkXNotImplemented):
        props["has_cycle"] = False
        props["n_cycles"] = 0

    # ── Diameter ───────────────────────────────────────────────
    # For weakly connected component only
    try:
        if G.number_of_nodes() > 1:
            # Use weakly connected component with most nodes
            wccs = list(nx.weakly_connected_components(G))
            largest = max(wccs, key=len)
            subG = G.subgraph(largest)

            if subG.number_of_nodes() >= 2:
                # Shortest path lengths in underlying undirected graph
                UG = subG.to_undirected()
                if nx.is_connected(UG):
                    props["diameter"] = nx.diameter(UG)
                    props["avg_path_length"] = nx.average_shortest_path_length(UG)
                    props["radius"] = nx.radius(UG)
                else:
                    # Compute within the largest connected component
                    components = list(nx.connected_components(UG))
                    main = UG.subgraph(max(components, key=len))
                    if main.number_of_nodes() >= 2:
                        props["diameter"] = nx.diameter(main)
                        props["avg_path_length"] = nx.average_shortest_path_length(main)
                    else:
                        props["diameter"] = 0
                        props["avg_path_length"] = 0.0

                    # For disconnected graphs, report effective diameter
                    if not nx.is_connected(UG):
                        props["disconnected"] = True
                        props["n_components"] = len(components)
            else:
                props["diameter"] = 0
                props["avg_path_length"] = 0.0
        else:
            props["diameter"] = 0
            props["avg_path_length"] = 0.0
    except Exception as e:
        props["diameter_error"] = str(e)
        props["diameter"] = -1
        props["avg_path_length"] = -1.0

    # ── Clustering (local transitivity) ────────────────────────
    try:
        UG = G.to_undirected()
        props["avg_clustering"] = nx.average_clustering(UG)
        props["global_clustering"] = nx.transitivity(UG)
    except Exception:
        props["avg_clustering"] = 0.0
        props["global_clustering"] = 0.0

    return props


def compute_small_world_index(G: nx.DiGraph, n_rand: int = 100) -> dict:
    """
    Compute the small-world index S = (C/C_rand) / (L/L_rand)

    S > 1 indicates small-world structure (high clustering + short paths).
    Higher S = more efficient reasoning structure per Topology of Reasoning paper.
    """
    props = {}
    UG = G.to_undirected()

    # Observed clustering and path length
    try:
        if UG.number_of_nodes() < 3:
            return {"small_world_index": 0, "reason": "too_few_nodes"}

        C = nx.average_clustering(UG)

        if nx.is_connected(UG) and UG.number_of_nodes() >= 2:
            L = nx.average_shortest_path_length(UG)
        else:
            # Use largest component
            components = list(nx.connected_components(UG))
            if not components:
                return {"small_world_index": 0, "reason": "no_components"}
            main = UG.subgraph(max(components, key=len))
            if main.number_of_nodes() < 2:
                return {"small_world_index": 0, "reason": "component_too_small"}
            L = nx.average_shortest_path_length(main)
    except Exception as e:
        return {"small_world_index": -1, "error": str(e)}

    # Random graph properties (average over n_rand ER graphs)
    n_nodes = G.number_of_nodes()
    p_edge = nx.density(UG)

    C_rands = []
    L_rands = []

    for _ in range(n_rand):
        ER = nx.erdos_renyi_graph(n=n_nodes, p=p_edge, seed=_)
        if ER.number_of_edges() == 0:
            continue
        try:
            C_rands.append(nx.average_clustering(ER))
            if nx.is_connected(ER) and ER.number_of_nodes() >= 2:
                L_rands.append(nx.average_shortest_path_length(ER))
        except Exception:
            continue

    if not C_rands or not L_rands:
        return {"small_world_index": -1, "error": "random_graph_generation_failed"}

    C_rand = np.mean(C_rands)
    L_rand = np.mean(L_rands)

    if C_rand == 0 or L_rand == 0:
        return {"small_world_index": -1, "error": "division_by_zero"}

    S = (C / C_rand) / (L / L_rand)

    props["small_world_index"] = S
    props["C_observed"] = C
    props["L_observed"] = L
    props["C_random"] = C_rand
    props["L_random"] = L_rand
    props["n_random_samples"] = len(C_rands)

    return props


# ── Full Analysis Pipeline ──────────────────────────────────────────

def analyze_reasoning_trace(
    text: str,
    method: str = "self_correct",
    label: str = "",
) -> dict:
    """
    Full pipeline: text → graph → properties → small-world index.
    """
    G = build_reasoning_graph(text, method=method)
    props = compute_properties(G)
    sw = compute_small_world_index(G)

    result = {
        "label": label or "unnamed",
        "text_length": len(text),
        "graph_method": method,
        "graph": {
            "n_nodes": G.number_of_nodes(),
            "n_edges": G.number_of_edges(),
            "edges": [f"{u}→{v}" for u, v in G.edges()],
        },
        "properties": props,
        "small_world": sw,
    }
    return result


def analyze_agot_graph(agot_graph, label: str = "") -> dict:
    """Analyze an AGoT graph using the same property pipeline."""
    G = build_graph_from_agot(agot_graph)
    props = compute_properties(G)
    sw = compute_small_world_index(G)

    return {
        "label": label or "agot_graph",
        "graph": {
            "n_nodes": G.number_of_nodes(),
            "n_edges": G.number_of_edges(),
        },
        "properties": props,
        "small_world": sw,
    }


# ── Batch Analysis ──────────────────────────────────────────────────

def analyze_results_file(json_path: str) -> dict:
    """Analyze all reasoning traces in a results file."""
    with open(json_path) as f:
        data = json.load(f)

    traces = []
    for detail in data.get("details", []):
        full_response = detail.get("full_response", "")
        if full_response:
            analysis = analyze_reasoning_trace(
                full_response,
                label=detail.get("question", "")[:50]
            )
            analysis["correct"] = detail.get("correct", False)
            traces.append(analysis)

    # Aggregate statistics
    if traces:
        agg = {
            "file": json_path,
            "n_traces": len(traces),
            "n_correct": sum(1 for t in traces if t.get("correct")),
            "accuracy": sum(1 for t in traces if t.get("correct")) / len(traces),
            "avg_nodes": np.mean([t["graph"]["n_nodes"] for t in traces]),
            "avg_cycles": np.mean([t["properties"]["n_cycles"] for t in traces]),
            "cycle_rate": sum(1 for t in traces if t["properties"]["has_cycle"]) / len(traces),
            "avg_diameter": np.mean([t["properties"].get("diameter", 0) for t in traces]),
            "avg_sw_index": np.mean(
                [t["small_world"].get("small_world_index", 0)
                 for t in traces if t["small_world"].get("small_world_index", -1) >= 0]
            ),
        }
    else:
        agg = {"file": json_path, "n_traces": 0}

    return {"aggregate": agg, "traces": traces}


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

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Analyze reasoning graphs")
    parser.add_argument("--file", type=str, help="JSON results file to analyze")
    parser.add_argument("--text", type=str, help="Direct text input for analysis")
    parser.add_argument("--verbose", action="store_true", help="Show detailed output")
    args = parser.parse_args()

    if args.file:
        result = analyze_results_file(args.file)
        print(json.dumps(result["aggregate"], indent=2))
        if args.verbose:
            for t in result["traces"]:
                print(f"\n  Trace: {t['label']}")
                print(f"    Nodes: {t['graph']['n_nodes']}, Edges: {t['graph']['n_edges']}")
                print(f"    Cycles: {t['properties']['n_cycles']}")
                print(f"    Diameter: {t['properties'].get('diameter', 'N/A')}")
                print(f"    SW Index: {t['small_world'].get('small_world_index', 'N/A')}")

    elif args.text:
        result = analyze_reasoning_trace(args.text, label="direct_input")
        print(json.dumps(result, indent=2, default=str))
    else:
        # Demo
        sample = """First, I need to find the total number of eggs: 16 per day.
        She eats 3 for breakfast, and uses 4 for muffins. That's 3+4=7 eggs used.
        So remaining eggs = 16-7 = 9 eggs per day.
        She sells them at $2 per dozen. 9 eggs = 9/12 = 0.75 dozen.
        Wait, let me check: 9/12 = 3/4 = 0.75. That's correct.
        So money = 0.75 × $2 = $1.50.
        But actually, is there any other use of eggs? No, only breakfast and muffins.
        Final answer: $1.50 per day."""
        result = analyze_reasoning_trace(sample, method="self_correct", label="demo")
        print("=== Demo Analysis ===")
        print(f"Nodes: {result['graph']['n_nodes']}")
        print(f"Edges: {result['graph']['n_edges']}")
        print(f"Edges: {result['graph']['edges']}")
        print(f"Cycles: {result['properties']['n_cycles']}")
        print(f"Has cycle: {result['properties']['has_cycle']}")
        print(f"Diameter: {result['properties'].get('diameter', 'N/A')}")
        print(f"SW Index: {result['small_world'].get('small_world_index', 'N/A')}")
