"""Rho-Surgery — Surgical Plan generation and verification.

Wraps existing diagnostics (audit, layer heatmap, probe landscape) into a
machine-readable SurgicalPlan that prescribes:
  1. Which categories are at risk of collateral damage
  2. Recommended HybridConfig (compress, SAE, rho, gamma)
  3. Protection strategy (which behaviors/categories to protect)

And after execution, verifies the outcome:
  - Pass/fail per category (within acceptable collateral tolerance)
  - Suggested adjustments if verification fails

Usage:
    from rho_eval.surgical_planner import generate_surgical_plan, verify_surgical_outcome

    # Generate plan from audit
    plan = generate_surgical_plan("Qwen/Qwen2.5-7B-Instruct")
    plan.to_json("surgical_plan.json")

    # After running hybrid pipeline with plan's config:
    verification = verify_surgical_outcome(plan, result)
    print(verification["summary"])
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Any, Optional


# ── Risk classification thresholds ──────────────────────────────────────

HIGH_RISK_ACCURACY = 0.85   # Categories with baseline > 85% have most to lose
MEDIUM_RISK_ACCURACY = 0.65  # 65-85% = medium risk
MIN_PROBES_FOR_PROTECTION = 10  # Need enough probes for reliable signal

# γ weight defaults based on risk profile
GAMMA_DEFAULTS = {
    "conservative": 0.2,   # Protect aggressively (may limit target improvement)
    "balanced": 0.1,       # Default balance between target and protection
    "aggressive": 0.05,    # Prioritize target improvement, accept some damage
}


@dataclass
class CategoryRisk:
    """Risk profile for a single bias category."""
    name: str
    accuracy: float
    n_probes: int
    biased_rate: float = 0.0
    risk_level: str = "low"  # "high", "medium", "low"

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class SurgicalPlan:
    """Machine-readable plan for a Rho-Surgery intervention.

    Generated by ``generate_surgical_plan()`` from audit data and
    optional layer heatmap / probe landscape analysis. Consumed by
    ``experiments/rho_surgery.py`` to configure the hybrid pipeline.
    """

    # ── Model info ──────────────────────────────────────────────────────
    model_name: str
    target_behaviors: list[str] = field(default_factory=lambda: ["sycophancy"])

    # ── Baseline audit snapshot ──────────────────────────────────────────
    baseline_scores: dict[str, float] = field(default_factory=dict)
    """Per-behavior rho scores from the baseline audit."""

    # ── Category risk profile ────────────────────────────────────────────
    category_risk: dict[str, dict] = field(default_factory=dict)
    """category_name → {accuracy, n_probes, biased_rate, risk_level}"""

    high_risk_categories: list[str] = field(default_factory=list)
    """Categories with risk_level == "high" AND sufficient probes."""

    # ── Layer analysis (optional) ─────────────────────────────────────────
    recommended_sae_layer: Optional[int] = None
    layer_sensitivity: dict[str, dict] = field(default_factory=dict)
    """layer_idx (as str) → {behavior: delta_rho}"""

    # ── Recommended config ────────────────────────────────────────────────
    compress_ratio: float = 0.7
    freeze_fraction: float = 0.75
    sae_layer: Optional[int] = None
    rho_weight: float = 0.2
    gamma_weight: float = 0.1
    protection_behaviors: list[str] = field(default_factory=lambda: ["bias"])
    protection_categories: list[str] = field(default_factory=list)
    """Categories to protect with γ loss (auto-selected from high-risk)."""

    # ── Probe landscape (optional) ────────────────────────────────────────
    probe_redundancy: dict[str, float] = field(default_factory=dict)
    """behavior → mean pairwise similarity (higher = more redundant probes)."""

    cross_behavior_clusters: list[dict] = field(default_factory=list)
    """Clusters spanning multiple behaviors (potential interference zones)."""

    # ── Metadata ──────────────────────────────────────────────────────────
    strategy: str = "balanced"
    """One of: conservative, balanced, aggressive."""

    notes: list[str] = field(default_factory=list)
    """Human-readable notes about the plan's reasoning."""

    # ── Serialization ─────────────────────────────────────────────────────

    def to_dict(self) -> dict:
        return asdict(self)

    def to_json(self, path: str | Path) -> Path:
        """Save plan to JSON file."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(self.to_dict(), indent=2, default=str))
        return path

    @classmethod
    def from_json(cls, path: str | Path) -> "SurgicalPlan":
        """Load plan from JSON file."""
        path = Path(path)
        data = json.loads(path.read_text())
        return cls(**data)

    def to_hybrid_config(self):
        """Convert plan to a HybridConfig for the hybrid pipeline."""
        from .hybrid.schema import HybridConfig
        return HybridConfig(
            compress_ratio=self.compress_ratio,
            freeze_fraction=self.freeze_fraction,
            sae_layer=self.sae_layer,
            rho_weight=self.rho_weight,
            target_behaviors=tuple(self.target_behaviors),
            gamma_weight=self.gamma_weight,
            protection_behaviors=tuple(self.protection_behaviors),
            protection_categories=tuple(self.protection_categories),
        )

    def summary(self) -> str:
        """Human-readable summary of the surgical plan."""
        lines = [
            f"Surgical Plan for {self.model_name}",
            f"  Target:     {', '.join(self.target_behaviors)}",
            f"  Strategy:   {self.strategy}",
            f"  SVD:        ratio={self.compress_ratio}, freeze={self.freeze_fraction}",
            f"  SAE:        layer={self.sae_layer}",
            f"  SFT:        ρ={self.rho_weight}, γ={self.gamma_weight}",
            f"  Protection: {', '.join(self.protection_categories) or 'none'}",
        ]
        if self.high_risk_categories:
            lines.append(f"  High-risk:  {', '.join(self.high_risk_categories)}")
        for note in self.notes:
            lines.append(f"  Note: {note}")
        return "\n".join(lines)


# ── Plan generation ─────────────────────────────────────────────────────


def _classify_category_risk(
    category_metrics: dict[str, dict],
) -> tuple[dict[str, CategoryRisk], list[str]]:
    """Classify each category's risk level from audit metrics.

    Returns:
        Tuple of (all_risks dict, high_risk_category_names list).
    """
    risks: dict[str, CategoryRisk] = {}
    high_risk: list[str] = []

    for cat, metrics in category_metrics.items():
        accuracy = metrics.get("accuracy", 0.0)
        n = metrics.get("n", 0)
        biased_rate = metrics.get("biased_rate", 0.0)

        if accuracy >= HIGH_RISK_ACCURACY:
            risk_level = "high"
        elif accuracy >= MEDIUM_RISK_ACCURACY:
            risk_level = "medium"
        else:
            risk_level = "low"

        risk = CategoryRisk(
            name=cat,
            accuracy=accuracy,
            n_probes=n,
            biased_rate=biased_rate,
            risk_level=risk_level,
        )
        risks[cat] = risk

        # Only protect categories with enough probes for reliable signal
        if risk_level == "high" and n >= MIN_PROBES_FOR_PROTECTION:
            high_risk.append(cat)

    return risks, high_risk


def _recommend_sae_layer(
    heatmap_data: dict,
    target_behavior: str = "sycophancy",
) -> tuple[Optional[int], dict]:
    """Recommend SAE layer from layer heatmap data.

    Picks the layer with largest positive Δρ for the target behavior
    while minimizing negative Δρ on other behaviors.

    Args:
        heatmap_data: Layer heatmap JSON data (from scripts/steering analysis).
        target_behavior: Behavior to optimize for.

    Returns:
        Tuple of (recommended_layer, sensitivity_dict).
    """
    if not heatmap_data:
        return None, {}

    sensitivity: dict[str, dict] = {}
    best_layer = None
    best_score = float("-inf")

    for layer_str, layer_data in heatmap_data.items():
        try:
            layer_idx = int(layer_str)
        except ValueError:
            continue

        layer_deltas = {}
        for behavior, delta in layer_data.items():
            if isinstance(delta, (int, float)):
                layer_deltas[behavior] = delta

        sensitivity[layer_str] = layer_deltas

        # Score: target improvement minus worst collateral damage
        target_delta = layer_deltas.get(target_behavior, 0.0)
        non_target_damage = min(
            (v for k, v in layer_deltas.items() if k != target_behavior),
            default=0.0,
        )
        score = target_delta - abs(min(non_target_damage, 0.0))

        if score > best_score:
            best_score = score
            best_layer = layer_idx

    return best_layer, sensitivity


def _select_gamma_weight(
    high_risk_categories: list[str],
    strategy: str = "balanced",
) -> float:
    """Auto-select γ weight based on risk profile and strategy.

    More high-risk categories → higher γ weight (more protection needed).
    """
    base = GAMMA_DEFAULTS.get(strategy, GAMMA_DEFAULTS["balanced"])

    # Scale up slightly if many categories at risk
    n_risk = len(high_risk_categories)
    if n_risk >= 4:
        return round(min(base * 1.5, 0.3), 2)
    elif n_risk >= 2:
        return base
    elif n_risk == 1:
        return max(base * 0.75, 0.05)
    else:
        return 0.0  # Nothing to protect


def generate_surgical_plan(
    model_name: str,
    target_behavior: str = "sycophancy",
    *,
    audit_report=None,
    audit_json_path: Optional[str] = None,
    layer_heatmap_path: Optional[str] = None,
    probe_landscape_path: Optional[str] = None,
    strategy: str = "balanced",
    compress_ratio: float = 0.7,
    freeze_fraction: float = 0.75,
) -> SurgicalPlan:
    """Generate a surgical plan from diagnostic data.

    Analyzes audit results to identify at-risk categories, recommends
    intervention parameters, and produces a machine-readable plan.

    Args:
        model_name: HuggingFace model ID.
        target_behavior: Primary behavior to improve.
        audit_report: Pre-computed AuditReport (optional).
        audit_json_path: Path to saved audit JSON (alternative to report).
        layer_heatmap_path: Path to layer heatmap JSON (optional).
        probe_landscape_path: Path to probe landscape JSON (optional).
        strategy: "conservative", "balanced", or "aggressive".
        compress_ratio: SVD compression ratio (default 0.7).
        freeze_fraction: Layer freeze fraction (default 0.75).

    Returns:
        SurgicalPlan with recommended config and protection strategy.
    """
    notes: list[str] = []

    # ── 1. Load or use audit data ─────────────────────────────────────
    baseline_scores: dict[str, float] = {}
    category_metrics: dict[str, dict] = {}

    if audit_report is not None:
        # Use pre-computed report
        baseline_scores = {
            name: r.rho for name, r in audit_report.behaviors.items()
        }
        category_metrics = audit_report.category_report("bias")
        notes.append("Plan generated from live audit report")

    elif audit_json_path:
        # Load from saved JSON — supports both AuditReport and HybridResult formats
        audit_data = json.loads(Path(audit_json_path).read_text())

        if "model" in audit_data and "behaviors" in audit_data:
            # Standard AuditReport format
            from .output.schema import AuditReport
            report = AuditReport.load(audit_json_path)
            baseline_scores = {
                name: r.rho for name, r in report.behaviors.items()
            }
            category_metrics = report.category_report("bias")
            notes.append(f"Plan generated from audit report: {audit_json_path}")

        elif "phases" in audit_data:
            # HybridResult format — extract baseline phase data
            phases = audit_data.get("phases", [])
            for phase in phases:
                if phase.get("phase") == "baseline":
                    phase_report = phase.get("details", {}).get("report", {})
                    behaviors = phase_report.get("behaviors", {})
                    baseline_scores = {
                        name: bdata.get("rho", 0.0)
                        for name, bdata in behaviors.items()
                    }
                    bias_data = behaviors.get("bias", {})
                    category_metrics = bias_data.get(
                        "metadata", {},
                    ).get("category_metrics", {})
                    break
            notes.append(f"Plan generated from hybrid result: {audit_json_path}")

        else:
            notes.append(f"Unknown audit format at {audit_json_path}")

    else:
        notes.append("No audit data provided — using default protection")

    # ── 2. Category risk classification ──────────────────────────────
    category_risk_objs, high_risk = _classify_category_risk(category_metrics)
    category_risk_dicts = {k: v.to_dict() for k, v in category_risk_objs.items()}

    if high_risk:
        notes.append(
            f"High-risk categories (baseline >85%, n≥{MIN_PROBES_FOR_PROTECTION}): "
            f"{', '.join(high_risk)}"
        )
    else:
        notes.append("No high-risk categories identified — γ protection disabled")

    # ── 3. Layer heatmap analysis (optional) ─────────────────────────
    recommended_sae_layer = None
    layer_sensitivity: dict[str, dict] = {}

    if layer_heatmap_path:
        path = Path(layer_heatmap_path)
        if path.exists():
            heatmap_data = json.loads(path.read_text())
            recommended_sae_layer, layer_sensitivity = _recommend_sae_layer(
                heatmap_data, target_behavior,
            )
            if recommended_sae_layer is not None:
                notes.append(f"SAE layer {recommended_sae_layer} recommended from heatmap")
        else:
            notes.append(f"Layer heatmap not found: {layer_heatmap_path}")

    # ── 4. Probe landscape analysis (optional) ───────────────────────
    probe_redundancy: dict[str, float] = {}
    cross_behavior_clusters: list[dict] = []

    if probe_landscape_path:
        path = Path(probe_landscape_path)
        if path.exists():
            landscape = json.loads(path.read_text())
            probe_redundancy = landscape.get("behavior_redundancy", {})
            cross_behavior_clusters = landscape.get("cross_behavior_clusters", [])
            if cross_behavior_clusters:
                notes.append(
                    f"{len(cross_behavior_clusters)} cross-behavior clusters found "
                    f"(potential interference zones)"
                )
        else:
            notes.append(f"Probe landscape not found: {probe_landscape_path}")

    # ── 5. Auto-select γ weight ──────────────────────────────────────
    gamma_weight = _select_gamma_weight(high_risk, strategy)

    # ── 6. Build protection strategy ─────────────────────────────────
    protection_categories = list(high_risk)
    protection_behaviors = ["bias"] if high_risk else []

    # ── 7. Assemble plan ─────────────────────────────────────────────
    plan = SurgicalPlan(
        model_name=model_name,
        target_behaviors=[target_behavior],
        baseline_scores=baseline_scores,
        category_risk=category_risk_dicts,
        high_risk_categories=high_risk,
        recommended_sae_layer=recommended_sae_layer,
        layer_sensitivity=layer_sensitivity,
        compress_ratio=compress_ratio,
        freeze_fraction=freeze_fraction,
        sae_layer=recommended_sae_layer,
        rho_weight=0.2,
        gamma_weight=gamma_weight,
        protection_behaviors=protection_behaviors,
        protection_categories=protection_categories,
        probe_redundancy=probe_redundancy,
        cross_behavior_clusters=cross_behavior_clusters,
        strategy=strategy,
        notes=notes,
    )

    return plan


# ── Verification ────────────────────────────────────────────────────────


def verify_surgical_outcome(
    plan: SurgicalPlan,
    result,
    *,
    max_collateral_per_category: float = 0.05,
    max_collateral_overall: float = 0.10,
    min_target_improvement: float = 0.05,
) -> dict:
    """Verify a surgical outcome against the plan's expectations.

    Checks:
    1. Target behavior improved by at least min_target_improvement
    2. No protected category dropped by more than max_collateral_per_category
    3. Overall non-target regression within max_collateral_overall

    Args:
        plan: The SurgicalPlan that guided the intervention.
        result: HybridResult from the hybrid pipeline.
        max_collateral_per_category: Max acceptable per-category accuracy drop.
        max_collateral_overall: Max acceptable overall non-target regression.
        min_target_improvement: Min acceptable target behavior improvement.

    Returns:
        Dict with:
          - passed: bool — overall pass/fail
          - target_check: {passed, improvement, threshold}
          - category_checks: {category: {passed, before, after, delta}}
          - overall_check: {passed, regression, threshold}
          - suggestions: list[str] — adjustments if failed
          - summary: str — one-line summary
    """
    suggestions: list[str] = []

    # ── 1. Target improvement check ──────────────────────────────────
    target_improvement = result.target_improvement
    target_passed = target_improvement >= min_target_improvement
    target_check = {
        "passed": target_passed,
        "improvement": round(target_improvement, 4),
        "threshold": min_target_improvement,
    }
    if not target_passed:
        suggestions.append(
            f"Target improvement ({target_improvement:.4f}) below threshold "
            f"({min_target_improvement}). Try: increase rho_weight or add more "
            f"sycophancy contrast pairs."
        )

    # ── 2. Per-category checks (from final audit) ────────────────────
    category_checks: dict[str, dict] = {}

    # Get final audit's category metrics
    final_phase = result.phases[-1] if result.phases else None
    final_report = final_phase.details.get("report", {}) if final_phase else {}
    final_behaviors = final_report.get("behaviors", {})
    final_bias = final_behaviors.get("bias", {})
    final_cat_metrics = final_bias.get("metadata", {}).get("category_metrics", {})

    # Get baseline category metrics
    baseline_phase = result.phases[0] if result.phases else None
    baseline_report = baseline_phase.details.get("report", {}) if baseline_phase else {}
    baseline_behaviors = baseline_report.get("behaviors", {})
    baseline_bias = baseline_behaviors.get("bias", {})
    baseline_cat_metrics = baseline_bias.get("metadata", {}).get("category_metrics", {})

    n_category_failures = 0
    for cat in plan.protection_categories:
        before = baseline_cat_metrics.get(cat, {}).get("accuracy", 0.0)
        after = final_cat_metrics.get(cat, {}).get("accuracy", 0.0)
        delta = after - before
        passed = delta >= -max_collateral_per_category

        category_checks[cat] = {
            "passed": passed,
            "before": round(before, 4),
            "after": round(after, 4),
            "delta": round(delta, 4),
            "threshold": -max_collateral_per_category,
        }

        if not passed:
            n_category_failures += 1
            suggestions.append(
                f"Category '{cat}' dropped {delta:+.1%} (threshold: "
                f"{-max_collateral_per_category:+.1%}). Try: increase gamma_weight "
                f"or add '{cat}' to protection_categories."
            )

    # ── 3. Overall non-target regression check ───────────────────────
    overall_regression = result.non_target_regression
    overall_passed = overall_regression >= -max_collateral_overall
    overall_check = {
        "passed": overall_passed,
        "regression": round(overall_regression, 4),
        "threshold": -max_collateral_overall,
    }
    if not overall_passed:
        suggestions.append(
            f"Overall non-target regression ({overall_regression:.4f}) exceeds "
            f"threshold ({-max_collateral_overall}). Consider conservative strategy."
        )

    # ── 4. Generate suggestions for next iteration ────────────────────
    if n_category_failures > 0 and plan.gamma_weight < 0.3:
        suggestions.append(
            f"Suggestion: increase gamma_weight from {plan.gamma_weight} to "
            f"{min(plan.gamma_weight * 1.5, 0.3):.2f}"
        )

    if n_category_failures > 2:
        suggestions.append(
            "Multiple category failures — consider 'conservative' strategy "
            "or reducing rho_weight."
        )

    # ── 5. Overall verdict ────────────────────────────────────────────
    all_passed = target_passed and overall_passed and n_category_failures == 0

    n_cats_checked = len(category_checks)
    n_cats_passed = sum(1 for c in category_checks.values() if c["passed"])

    summary = (
        f"{'PASS' if all_passed else 'FAIL'}: "
        f"target {'✓' if target_passed else '✗'} ({target_improvement:+.4f}), "
        f"categories {n_cats_passed}/{n_cats_checked} passed, "
        f"overall {'✓' if overall_passed else '✗'} ({overall_regression:+.4f})"
    )

    return {
        "passed": all_passed,
        "target_check": target_check,
        "category_checks": category_checks,
        "overall_check": overall_check,
        "suggestions": suggestions,
        "summary": summary,
    }
