"""Evaluate a single AF2-predicted PDB against a case's locked eval regions.

For a known-state case (KaiB, GA/GB):
- pLDDT overall and per-region (common_core, switch_region, switch_region_2A)
- Cα RMSD to state_A and state_B on common_core, switch_region, switch_region_2A
- TM-score to state_A and state_B on common_core
- Primary hit criterion (protocol §9.1, BINDING 3-condition): RMSD <= 3.0 Å on
  common_core AND mean pLDDT >= 70 overall AND switch_region pLDDT >= 70.
  NaN switch_region pLDDT → §9.1 NOT evaluable → hit_primary=False (issue #7).

For the discovery case (Mpt53):
- pLDDT overall
- RMSD and TM-score to state_A_reference (1LU4) on the full 134-aa common_core

CLI:
    python src/eval/evaluate_prediction.py PRED.pdb --case KaiB [--json OUT.json]
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Iterable

import numpy as np
import yaml
from Bio.PDB import PDBParser, Superimposer

# ROOT holds data/, configs/ and results/ (all paths below are ROOT-relative).
# Default = repo layout (this file at <ROOT>/src/eval/); override with the
# SF_BENCH_ROOT env var to point at a relocated benchmark bundle root.
_ENV_ROOT = os.environ.get("SF_BENCH_ROOT")
ROOT = Path(_ENV_ROOT).resolve() if _ENV_ROOT else Path(__file__).resolve().parents[2]
MANIFEST_REGIONS = ROOT / "data" / "manifests" / "eval_regions.tsv"
STRUCT_ROOT = ROOT / "data" / "processed" / "structures"
CASES_YAML = ROOT / "configs" / "cases.yaml"


def parse_csv_list(s: str) -> list[int]:
    if not s or s in ("NA", ""):
        return []
    return [int(x) for x in s.split(",") if x.strip()]


def load_construct_start(case: str) -> int:
    """Return construct_start for the case (used to offset predicted PDB numbering).

    ColabFold emits predicted residues numbered 1..L; our eval_regions are in the
    state PDB's numbering (e.g. KaiB 5..95, Mpt53 38..173 mapped from UniProt).
    Predicted residue i maps to region residue (i + construct_start - 1).
    """
    cases = yaml.safe_load(CASES_YAML.read_text())["cases"]
    for c in cases:
        if c["case_name"] == case:
            return int(c["construct_start"])
    raise KeyError(f"no cases.yaml entry for {case}")


def load_regions(case: str) -> dict:
    """Return common_core / switch_region_3A / switch_region_2A residue lists.

    For GA/GB the stored row is `GA_GB_GA98_vs_GB98`.
    For Mpt53 the row is `Mpt53` with switch cols NA.
    """
    row_key = {
        "KaiB": "KaiB",
        "GA_GB": "GA_GB_GA98_vs_GB98",
        "Mpt53": "Mpt53",
    }.get(case)
    if row_key is None:
        raise ValueError(f"unknown case {case}")
    with MANIFEST_REGIONS.open() as f:
        reader = csv.DictReader(f, delimiter="\t")
        for row in reader:
            if row["case"] == row_key:
                return {
                    "case": row_key,
                    "common_core": parse_csv_list(row.get("common_core_residues", "")),
                    "switch_3A": parse_csv_list(row.get("switch_region_3A_residues", "")),
                    "switch_2A": parse_csv_list(row.get("switch_region_2A_residues", "")),
                }
    raise KeyError(f"no eval_regions row for {row_key}")


def load_state_paths(case: str) -> dict[str, Path]:
    """Map state name -> PDB path for the case."""
    if case == "KaiB":
        return {
            "state_A_2QKE": STRUCT_ROOT / "KaiB" / "state_A.pdb",
            "state_B_5JYT": STRUCT_ROOT / "KaiB" / "state_B.pdb",
        }
    if case == "GA_GB":
        # Two fold canonical references
        return {
            "GA98_2LHC": STRUCT_ROOT / "GA_GB" / "GA98_2LHC.pdb",
            "GB98_2LHD": STRUCT_ROOT / "GA_GB" / "GB98_2LHD.pdb",
        }
    if case == "Mpt53":
        return {
            "state_A_1LU4": STRUCT_ROOT / "Mpt53" / "state_A.pdb",
        }
    raise ValueError(f"unknown case {case}")


def read_ca_per_resid(pdb: Path, chain: str | None = None,
                      resid_offset: int = 0) -> dict[int, np.ndarray]:
    """Return {resid: Cα coord} for the first model, first (or named) chain.

    resid_offset is added to each residue id (used to map 1LU4 crystal 1001..
    into UniProt 38..171 frame when needed).
    """
    s = PDBParser(QUIET=True).get_structure("x", str(pdb))
    m = next(iter(s))
    for c in m:
        if chain is not None and c.id != chain:
            continue
        out: dict[int, np.ndarray] = {}
        for r in c:
            if r.id[0] != " ":
                continue
            if "CA" not in r:
                continue
            out[r.id[1] + resid_offset] = r["CA"].coord
        if out:
            return out
    raise RuntimeError(f"no Cα atoms found in {pdb} chain={chain}")


def read_ca_bfactors(pdb: Path, chain: str | None = None,
                     resid_offset: int = 0) -> dict[int, float]:
    """Return {resid: Cα B-factor} — for AF2 outputs, B-factor == pLDDT."""
    s = PDBParser(QUIET=True).get_structure("x", str(pdb))
    m = next(iter(s))
    for c in m:
        if chain is not None and c.id != chain:
            continue
        out: dict[int, float] = {}
        for r in c:
            if r.id[0] != " ":
                continue
            if "CA" not in r:
                continue
            out[r.id[1] + resid_offset] = float(r["CA"].bfactor)
        if out:
            return out
    raise RuntimeError(f"no Cα B-factors in {pdb} chain={chain}")


def rmsd_on_residues(pred_ca: dict[int, np.ndarray],
                     ref_ca: dict[int, np.ndarray],
                     residues: Iterable[int]) -> tuple[float, int]:
    """Superpose pred onto ref using the given residue set; return RMSD + N."""
    residues = [r for r in residues if r in pred_ca and r in ref_ca]
    if len(residues) < 3:
        return float("nan"), len(residues)
    from Bio.PDB.Atom import Atom
    pred_atoms = [Atom("CA", pred_ca[r], 1.0, 1.0, " ", "CA", 1, "C") for r in residues]
    ref_atoms = [Atom("CA", ref_ca[r], 1.0, 1.0, " ", "CA", 1, "C") for r in residues]
    sup = Superimposer()
    sup.set_atoms(ref_atoms, pred_atoms)
    return float(sup.rms), len(residues)


def tmalign_score(pdb_a: Path, pdb_b: Path) -> tuple[float, float, int]:
    """Run TMalign, return (TM-score normalized by chain_1, chain_2, aligned_len).

    TMalign is OPTIONAL: it only fills the tmalign_* columns and is NOT part of
    the hit_primary criterion. If TMalign is not on PATH we skip it cleanly and
    return NaNs.
    """
    import shutil as _sh
    tm = _sh.which("TMalign")
    if tm is None:
        return float("nan"), float("nan"), 0
    try:
        proc = subprocess.run(
            [tm, str(pdb_a), str(pdb_b)],
            capture_output=True, text=True, timeout=120,
        )
    except Exception as e:
        print(f"WARN: TMalign failed for {pdb_a} vs {pdb_b}: {e}", file=sys.stderr)
        return float("nan"), float("nan"), 0
    tm1 = tm2 = float("nan")
    aln = 0
    for line in proc.stdout.splitlines():
        if line.startswith("TM-score=") and "normalized by length of Chain_1" in line:
            tm1 = float(line.split()[1])
        elif line.startswith("TM-score=") and "normalized by length of Chain_2" in line:
            tm2 = float(line.split()[1])
        elif line.startswith("Aligned length="):
            # "Aligned length= 66, RMSD= 3.86, Seq_id=n_identical/n_aligned= 0.000"
            parts = line.replace(",", " ").split()
            for i, p in enumerate(parts):
                if p == "length=":
                    aln = int(parts[i + 1])
    return tm1, tm2, aln


def evaluate(pdb: Path, case: str) -> dict:
    # §9.1 binding 3-condition hit criterion (CLAUDE.md §Evaluation Protocol):
    # ALL of (1) Cα RMSD ≤ 3.0 Å on common_core, (2) mean pLDDT ≥ 70 overall,
    # (3) mean pLDDT ≥ 70 in switch_region must hold. NaN switch_region pLDDT
    # means switch_region is undefined (or no overlap) → §9.1 is NOT evaluable
    # → hit_primary MUST be False (issue #7 fix).
    regions = load_regions(case)
    states = load_state_paths(case)

    # Predicted structure Cα + pLDDT (B-factors). Our predicted PDBs have chain A
    # and residues numbered 1..L. Shift to the state PDB / eval_region numbering.
    construct_start = load_construct_start(case)
    pred_offset = construct_start - 1
    pred_ca = read_ca_per_resid(pdb, chain="A", resid_offset=pred_offset)
    pred_plddt = read_ca_bfactors(pdb, chain="A", resid_offset=pred_offset)

    mean_plddt_overall = float(np.mean(list(pred_plddt.values())))
    mean_plddt_core = float(np.mean([pred_plddt[r] for r in regions["common_core"]
                                      if r in pred_plddt]))
    mean_plddt_switch3 = (float(np.mean([pred_plddt[r] for r in regions["switch_3A"]
                                          if r in pred_plddt]))
                          if regions["switch_3A"] else float("nan"))
    mean_plddt_switch2 = (float(np.mean([pred_plddt[r] for r in regions["switch_2A"]
                                          if r in pred_plddt]))
                          if regions["switch_2A"] else float("nan"))

    result = {
        "pdb": str(pdb.relative_to(ROOT)) if pdb.is_relative_to(ROOT) else str(pdb),
        "case": case,
        "pred_len": len(pred_ca),
        "mean_plddt_overall": mean_plddt_overall,
        "mean_plddt_core": mean_plddt_core,
        "mean_plddt_switch_3A": mean_plddt_switch3,
        "mean_plddt_switch_2A": mean_plddt_switch2,
        "states": {},
    }

    for state_name, state_pdb in states.items():
        # Determine state chain + offset
        if state_name.startswith("state_A_1LU4"):
            ref_chain = "A"
            ref_offset = -963  # crystal 1001..1134 -> UniProt 38..171
        elif state_name.startswith("state_A_2QKE"):
            ref_chain = "B"
            ref_offset = 0
        elif state_name.startswith("state_B_5JYT"):
            ref_chain = "A"
            ref_offset = 0
        else:
            # GA/GB canonical refs (GA98, GB98)
            ref_chain = "A"
            ref_offset = 0

        ref_ca = read_ca_per_resid(state_pdb, chain=ref_chain, resid_offset=ref_offset)

        rms_core, n_core = rmsd_on_residues(pred_ca, ref_ca, regions["common_core"])
        rms_s3, n_s3 = (rmsd_on_residues(pred_ca, ref_ca, regions["switch_3A"])
                        if regions["switch_3A"] else (float("nan"), 0))
        rms_s2, n_s2 = (rmsd_on_residues(pred_ca, ref_ca, regions["switch_2A"])
                        if regions["switch_2A"] else (float("nan"), 0))
        tm1, tm2, aln = tmalign_score(pdb, state_pdb)

        # §9.1 binding 3-condition: all three thresholds must hold. If
        # switch_region pLDDT is NaN, §9.1 is NOT evaluable → hit_primary=False.
        no_switch_region_defined = bool(np.isnan(mean_plddt_switch3))
        hit_primary = (
            not np.isnan(rms_core)
            and rms_core <= 3.0
            and mean_plddt_overall >= 70.0
            and (not no_switch_region_defined)
            and mean_plddt_switch3 >= 70.0
        )
        if no_switch_region_defined:
            print(
                f"WARN: switch_region pLDDT is NaN for {pdb} (case={case}, "
                f"state={state_name}); §9.1 not evaluable → hit_primary=False",
                file=sys.stderr,
            )

        result["states"][state_name] = {
            "rmsd_common_core_A": rms_core,
            "rmsd_switch_3A": rms_s3,
            "rmsd_switch_2A": rms_s2,
            "n_common_core_aligned": n_core,
            "n_switch_3A_aligned": n_s3,
            "n_switch_2A_aligned": n_s2,
            "tmalign_tm1": tm1,
            "tmalign_tm2": tm2,
            "tmalign_aligned_len": aln,
            "hit_primary": hit_primary,
            "no_switch_region_defined": no_switch_region_defined,
        }

    return result


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("pdb", type=Path, help="Path to AF2-predicted PDB")
    ap.add_argument("--case", required=True, choices=["KaiB", "GA_GB", "Mpt53"])
    ap.add_argument("--json", type=Path, default=None,
                    help="Optional path to dump JSON result")
    args = ap.parse_args(argv)

    result = evaluate(args.pdb, args.case)
    if args.json:
        args.json.parent.mkdir(parents=True, exist_ok=True)
        args.json.write_text(json.dumps(result, indent=2, default=float))
    print(json.dumps(result, indent=2, default=float))
    return 0


if __name__ == "__main__":
    sys.exit(main())