"""Judge Recalibration Diagnostic — Path 4.

Standalone diagnostic script that re-runs pairwise comparisons from
Run 3 (4 conditions) using a REORDERED criteria prompt.

Purpose:  Detect whether the current judge's forced-choice preference
is dominated by *surface fluency* vs *insight quality*.

Pre-registration (documented before running):
  PREDICTION: If we reorder pairwise criteria to put "Insight quality"
  first (instead of "Structural depth"), SCAFFOLD_FULL's Elo will
  increase by >= 30 points relative to VOCAB_ONLY.
  NULL:  Criteria order has no effect (< 15 Elo point shift).

Usage:
    python -m evaluation.diagnostic_judge_recal [--run-dir RUN_DIR]

Requires: EVAL_JUDGE_API_KEY environment variable.
"""
from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import logging
import random
import re
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Tuple

from .corpus.loader import TestText, load_corpus
from .judges.llm_client import JudgeLLM
from .judges.pairwise import PairwiseResult, _parse_pairwise
from .sanitize import sanitize_for_judge
from .stats.bradley_terry import fit_bradley_terry, strengths_to_elo

logger = logging.getLogger(__name__)

# ── Reordered pairwise prompt ────────────────────────────────────
# Original order: 1-Structural depth, 2-Multi-perspective, 3-Integration,
#                 4-Insight quality, 5-Error absence
# Diagnostic order: 1-Insight quality, 2-Structural depth, 3-Integration,
#                   4-Multi-perspective, 5-Error absence

REORDERED_PAIRWISE_SYSTEM_PROMPT = """\
You are a blind evaluator comparing two analyses of the same text.
You do NOT know which analysis was produced by which method.

EVALUATION CRITERIA (in order of importance):
1. Insight quality: Are the insights non-obvious and well-grounded?
2. Structural depth: Does the analysis reveal hidden structure in the text?
3. Integration: Do the perspectives connect to each other meaningfully?
4. Multi-perspective coverage: How many distinct analytical angles are applied?
5. Error absence: Is the analysis free from logical errors and unsupported claims?

INSTRUCTIONS:
- Read both analyses carefully
- Compare on each criterion above
- You MUST pick a winner. There is NO tie option.
- Even if both analyses are close in quality, identify the one that is \
marginally better. There is always a difference.
- If criteria votes are split, use this tie-breaking hierarchy: \
Insight quality > Integration > Structural depth > Multi-perspective > Error absence

OUTPUT FORMAT:
Winner: [A or B]
Reasoning: [2-3 sentences explaining the decision]
Per-criterion:
  Insight quality: [A or B]
  Structural depth: [A or B]
  Integration: [A or B]
  Multi-perspective: [A or B]
  Error absence: [A or B]
"""

REORDERED_PAIRWISE_USER_PROMPT = """\
ORIGINAL TEXT:
{input_text}

ANALYSIS A:
{analysis_a}

ANALYSIS B:
{analysis_b}
"""

# Conditions from Run 3 (4-condition pipeline)
RUN3_CONDITIONS = ["NAKED", "SCAFFOLD_FULL", "SCAFFOLD_PARTIAL", "VOCAB_ONLY"]


def _load_run3_outputs(run_dir: Path) -> Dict[str, Dict[str, str]]:
    """Load raw outputs from a completed run's raw/ directory.

    Returns {text_id: {condition_id: raw_output}}.
    Falls back to the global cache if raw/ doesn't exist.
    """
    outputs: Dict[str, Dict[str, str]] = {}
    raw_dir = run_dir / "raw"

    if raw_dir.exists():
        for cond_dir in raw_dir.iterdir():
            if not cond_dir.is_dir() or cond_dir.name not in RUN3_CONDITIONS:
                continue
            for fp in cond_dir.glob("*.json"):
                data = json.loads(fp.read_text(encoding="utf-8"))
                tid = data["text_id"]
                outputs.setdefault(tid, {})[cond_dir.name] = data["raw_output"]
        return outputs

    # Fallback: reconstruct from global cache
    logger.info("raw/ not found — reconstructing from global cache")
    cache_dir = Path("evaluation/results/cache")
    corpus = load_corpus()
    for text in corpus:
        outputs[text.id] = {}
        for cond in RUN3_CONDITIONS:
            key = f"{text.id}_{cond}"
            h = hashlib.sha256(key.encode()).hexdigest()[:12]
            cache_path = cache_dir / f"{h}.json"
            if cache_path.exists():
                data = json.loads(cache_path.read_text(encoding="utf-8"))
                outputs[text.id][cond] = data["raw_output"]
            else:
                logger.warning("Missing cache for %s × %s", text.id, cond)

    # Remove texts that don't have all 4 conditions
    complete = {
        tid: conds for tid, conds in outputs.items()
        if len(conds) == len(RUN3_CONDITIONS)
    }
    logger.info("Loaded %d complete texts from cache", len(complete))
    return complete


async def _reordered_compare_pair(
    judge: JudgeLLM,
    text_id: str,
    input_text: str,
    output_a: str,
    condition_a: str,
    output_b: str,
    condition_b: str,
) -> PairwiseResult:
    """Run one pairwise comparison with the REORDERED criteria prompt."""
    # Anti-bias: randomise presentation order
    if random.random() < 0.5:
        shown_a, shown_b = output_a, output_b
        map_a, map_b = condition_a, condition_b
    else:
        shown_a, shown_b = output_b, output_a
        map_a, map_b = condition_b, condition_a

    response = await judge.judge(
        REORDERED_PAIRWISE_SYSTEM_PROMPT,
        REORDERED_PAIRWISE_USER_PROMPT.format(
            input_text=input_text[:2000],
            analysis_a=shown_a[:4000],
            analysis_b=shown_b[:4000],
        ),
    )

    winner_label, reasoning, per_criterion = _parse_pairwise(response)

    if winner_label == "A":
        true_winner = map_a
    elif winner_label == "B":
        true_winner = map_b
    else:
        true_winner = "TIE"

    return PairwiseResult(
        text_id=text_id,
        condition_a=map_a,
        condition_b=map_b,
        winner=winner_label,
        true_winner=true_winner,
        reasoning=reasoning,
        per_criterion=per_criterion,
    )


async def run_diagnostic(run_dir: Path) -> Dict[str, Any]:
    """Main diagnostic: rerun pairwise with reordered criteria, compare Elo."""
    judge = JudgeLLM()
    corpus = load_corpus()
    corpus_map = {t.id: t for t in corpus}

    # Load original Elo from Run 3
    elo_path = run_dir / "elo_ratings.json"
    original_elo: Dict[str, float] = {}
    if elo_path.exists():
        original_elo = json.loads(elo_path.read_text(encoding="utf-8"))
    logger.info("Original Elo: %s", original_elo)

    # Load raw outputs
    outputs = _load_run3_outputs(run_dir)
    text_ids = sorted(outputs.keys())
    logger.info("Running reordered pairwise on %d texts × C(%d,2)=%d pairs",
                len(text_ids), len(RUN3_CONDITIONS),
                len(RUN3_CONDITIONS) * (len(RUN3_CONDITIONS) - 1) // 2)

    # Run all pairwise comparisons with reordered prompt
    all_pairwise: List[PairwiseResult] = []
    for idx, tid in enumerate(text_ids, 1):
        text = corpus_map.get(tid)
        if not text:
            logger.warning("Text %s not in corpus, skipping", tid)
            continue

        # Sanitize outputs
        cond_outputs = {
            cid: sanitize_for_judge(outputs[tid][cid], cid)
            for cid in RUN3_CONDITIONS
        }

        # All C(4,2)=6 pairs
        conditions = list(cond_outputs.keys())
        for i in range(len(conditions)):
            for j in range(i + 1, len(conditions)):
                result = await _reordered_compare_pair(
                    judge, tid, text.text,
                    cond_outputs[conditions[i]], conditions[i],
                    cond_outputs[conditions[j]], conditions[j],
                )
                all_pairwise.append(result)
        if idx % 10 == 0:
            logger.info("  Progress: %d/%d texts done", idx, len(text_ids))

    # Convert to BT records
    bt_records: List[Dict] = []
    for r in all_pairwise:
        if r.true_winner == "TIE":
            bt_records.append({
                "winner": r.condition_a,
                "loser": r.condition_b,
                "cond_a": r.condition_a,
                "cond_b": r.condition_b,
            })
            bt_records.append({
                "winner": r.condition_b,
                "loser": r.condition_a,
                "cond_a": r.condition_a,
                "cond_b": r.condition_b,
            })
        else:
            loser = r.condition_b if r.true_winner == r.condition_a else r.condition_a
            bt_records.append({
                "winner": r.true_winner,
                "loser": loser,
                "cond_a": r.condition_a,
                "cond_b": r.condition_b,
            })

    # Fit BT model
    strengths = fit_bradley_terry(bt_records, RUN3_CONDITIONS)
    new_elo = strengths_to_elo(strengths)

    # Build comparison report
    report = {
        "diagnostic": "judge_recalibration",
        "timestamp": datetime.now().isoformat(),
        "source_run": str(run_dir),
        "prompt_change": "Criteria reordered: Insight quality moved from #4 to #1",
        "pre_registration": {
            "prediction": "SCAFFOLD_FULL Elo increases >= 30 points relative to VOCAB_ONLY",
            "null_hypothesis": "Criteria order has no effect (< 15 Elo point shift)",
        },
        "original_elo": original_elo,
        "reordered_elo": new_elo,
        "elo_deltas": {},
        "relative_shifts": {},
        "per_criterion_wins": {},
        "total_comparisons": len(all_pairwise),
    }

    # Compute deltas
    for cond in RUN3_CONDITIONS:
        orig = original_elo.get(cond, 1500)
        new = new_elo.get(cond, 1500)
        report["elo_deltas"][cond] = round(new - orig, 1)

    # Relative shift: SCAFFOLD_FULL vs VOCAB_ONLY gap change
    if "SCAFFOLD_FULL" in original_elo and "VOCAB_ONLY" in original_elo:
        orig_gap = original_elo["VOCAB_ONLY"] - original_elo["SCAFFOLD_FULL"]
        new_gap = new_elo["VOCAB_ONLY"] - new_elo["SCAFFOLD_FULL"]
        report["relative_shifts"]["vocab_vs_full_gap_original"] = round(orig_gap, 1)
        report["relative_shifts"]["vocab_vs_full_gap_reordered"] = round(new_gap, 1)
        report["relative_shifts"]["gap_reduction"] = round(orig_gap - new_gap, 1)

    # Per-criterion win counts
    criterion_wins: Dict[str, Dict[str, int]] = {
        cond: {} for cond in RUN3_CONDITIONS
    }
    for r in all_pairwise:
        for crit, winner_label in r.per_criterion.items():
            if winner_label == "A":
                cond_winner = r.condition_a
            elif winner_label == "B":
                cond_winner = r.condition_b
            else:
                continue
            criterion_wins[cond_winner].setdefault(crit, 0)
            criterion_wins[cond_winner][crit] = criterion_wins[cond_winner].get(crit, 0) + 1
    report["per_criterion_wins"] = criterion_wins

    # Decision
    gap_reduction = report["relative_shifts"].get("gap_reduction", 0)
    if abs(gap_reduction) >= 30:
        report["verdict"] = (
            f"SIGNIFICANT: Criteria reordering shifted VOCAB-FULL gap by "
            f"{gap_reduction:.0f} points. Judge IS sensitive to criteria order. "
            f"Interpret Elo rankings with caution."
        )
    elif abs(gap_reduction) >= 15:
        report["verdict"] = (
            f"MARGINAL: Criteria reordering shifted VOCAB-FULL gap by "
            f"{gap_reduction:.0f} points. Borderline sensitivity detected."
        )
    else:
        report["verdict"] = (
            f"NULL CONFIRMED: Criteria reordering shifted VOCAB-FULL gap by only "
            f"{gap_reduction:.0f} points. Judge rankings are robust to criteria order."
        )

    return report


def _format_report(results: Dict[str, Any]) -> str:
    """Format diagnostic results as a readable markdown report."""
    lines = [
        "# Judge Recalibration Diagnostic",
        "",
        f"**Timestamp:** {results['timestamp']}",
        f"**Source run:** {results['source_run']}",
        f"**Prompt change:** {results['prompt_change']}",
        f"**Total comparisons:** {results['total_comparisons']}",
        "",
        "## Pre-registration",
        f"- **Prediction:** {results['pre_registration']['prediction']}",
        f"- **Null:** {results['pre_registration']['null_hypothesis']}",
        "",
        "## Elo Comparison",
        "",
        "| Condition | Original Elo | Reordered Elo | Delta |",
        "|-----------|-------------|---------------|-------|",
    ]

    for cond in RUN3_CONDITIONS:
        orig = results["original_elo"].get(cond, 1500)
        new = results["reordered_elo"].get(cond, 1500)
        delta = results["elo_deltas"].get(cond, 0)
        sign = "+" if delta > 0 else ""
        lines.append(f"| {cond} | {orig:.0f} | {new:.0f} | {sign}{delta:.0f} |")

    lines.extend([
        "",
        "## Relative Shift Analysis",
        "",
    ])
    shifts = results.get("relative_shifts", {})
    if shifts:
        lines.append(f"- VOCAB_ONLY − SCAFFOLD_FULL gap (original): **{shifts.get('vocab_vs_full_gap_original', 'N/A'):.0f}**")
        lines.append(f"- VOCAB_ONLY − SCAFFOLD_FULL gap (reordered): **{shifts.get('vocab_vs_full_gap_reordered', 'N/A'):.0f}**")
        lines.append(f"- Gap reduction: **{shifts.get('gap_reduction', 'N/A'):.0f}** points")

    lines.extend([
        "",
        "## Per-Criterion Win Counts",
        "",
    ])
    # Collect all criteria
    all_criteria = set()
    for cond_wins in results.get("per_criterion_wins", {}).values():
        all_criteria.update(cond_wins.keys())
    if all_criteria:
        header = "| Condition | " + " | ".join(sorted(all_criteria)) + " |"
        sep = "|-----------|" + "|".join(["------"] * len(all_criteria)) + "|"
        lines.append(header)
        lines.append(sep)
        for cond in RUN3_CONDITIONS:
            wins = results.get("per_criterion_wins", {}).get(cond, {})
            row = f"| {cond} | " + " | ".join(
                str(wins.get(c, 0)) for c in sorted(all_criteria)
            ) + " |"
            lines.append(row)

    lines.extend([
        "",
        "## Verdict",
        "",
        f"**{results['verdict']}**",
        "",
    ])
    return "\n".join(lines)


async def main(run_dir: str | None = None):
    """Entry point for the diagnostic."""
    logging.basicConfig(
        level=logging.INFO,
        format="[%(asctime)s] %(levelname)-8s %(message)s",
        datefmt="%m/%d/%y %H:%M:%S",
    )

    # Find the run directory
    if run_dir:
        rd = Path(run_dir)
    else:
        # Auto-detect: find the latest 4-condition run (Run 3)
        results_dir = Path("evaluation/results")
        candidates = sorted(
            [d for d in results_dir.iterdir()
             if d.is_dir() and d.name.startswith("run_")
             and (d / "elo_ratings.json").exists()],
            key=lambda d: d.name,
        )
        # Pick the latest run that has exactly 4 conditions in elo
        rd = None
        for c in reversed(candidates):
            elo = json.loads((c / "elo_ratings.json").read_text(encoding="utf-8"))
            if set(elo.keys()) == set(RUN3_CONDITIONS):
                rd = c
                break
        if rd is None:
            # Fall back to latest run with elo_ratings
            rd = candidates[-1] if candidates else None
        if rd is None:
            logger.error("No completed run found in evaluation/results/")
            return

    logger.info("Using run: %s", rd)
    results = await run_diagnostic(rd)

    # Save results
    out_dir = Path("evaluation/results") / "diagnostic_judge_recal"
    out_dir.mkdir(parents=True, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")

    json_path = out_dir / f"results_{ts}.json"
    json_path.write_text(
        json.dumps(results, indent=2, ensure_ascii=False, default=str),
        encoding="utf-8",
    )
    logger.info("Results saved to %s", json_path)

    report_path = out_dir / f"report_{ts}.md"
    report_text = _format_report(results)
    report_path.write_text(report_text, encoding="utf-8")
    logger.info("Report saved to %s", report_path)

    # Print summary (use ascii replacement for Windows console encoding)
    safe_text = report_text.replace("\u2212", "-")
    print("\n" + "=" * 60)
    print(safe_text)
    print("=" * 60)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Judge Recalibration Diagnostic")
    parser.add_argument("--run-dir", type=str, default=None,
                        help="Path to the source run directory (default: auto-detect Run 3)")
    args = parser.parse_args()
    asyncio.run(main(args.run_dir))
