# /// script
# requires-python = ">=3.11"
# dependencies = ["numpy==2.3.3", "pandas==2.3.3", "openpyxl==3.1.5"]
# ///
"""Reproduce a BRCA1 comparison from stored scores; no model or GPU needed.

Run: uv run --script reproduce_brca1.py --fetch
Requires the `tabix` executable from HTSlib for the first CADD regional fetch.
Read README.md for the data terms and limits of this retrospective exercise.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import subprocess
import time
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import pandas as pd

STUDY_REVISION = "f8ae11c9a1c45505487c978b4cb7e064c4e59905"
STUDY = f"https://raw.githubusercontent.com/rewire-bio/evo2-cloud-operations/{STUDY_REVISION}"
FILES = {
    "labels.xlsx": (
        "https://raw.githubusercontent.com/ArcInstitute/evo2/53f195997257c56c00e5ef8d33a54f5baad143a6/notebooks/brca1/41586_2018_461_MOESM3_ESM.xlsx",
        "0100fc651b6ec6015f02a5813660afdcf9cf3969790992104e99e74f7348d0fd",
    ),
    "evo2_scores.tsv": (
        f"{STUDY}/evidence/runs/recorded-b75c67aa/platforms/P2/attempt-1/w1_brca1_scores.tsv",
        "c76d845efd76543aa5237133c215f9c137436aa92caa736946d163ae50a4e942",
    ),
    "published_results.json": (
        f"{STUDY}/evidence/current-predictors/results.json",
        "403437941f5a0e08b3730fa7bbfcba0b8e616b0e87b2f8e3f227b2f86a342318",
    ),
}
CADD_URL = "https://kircherlab.bihealth.org/download/CADD/v1.7/GRCh37/whole_genome_SNVs.tsv.gz"
CADD_REGION = "17:41197000-41277000"
CADD_SHA256 = "3b8d4e83076e8fb1b2dcd0498119dc22cd148ebcf4fa0d16304d7c1cfc7f3168"
CADD_INDEX_MD5 = "6bc06b08fd9a521d05d970a06939e52d"
CODING = {"Missense", "Nonsense", "Synonymous"}
BOOTSTRAPS = 2000
SEED = 7


def checksum(path: Path, algorithm: str = "sha256") -> str:
    with path.open("rb") as handle:
        return hashlib.file_digest(handle, algorithm).hexdigest()


def acquire(cache: Path, fetch: bool) -> None:
    """Verify every input, fetching only declared files when explicitly requested."""
    cache.mkdir(parents=True, exist_ok=True)
    for name, (url, expected) in FILES.items():
        path = cache / name
        if not path.exists():
            if not fetch:
                raise RuntimeError(f"Missing {path}; run with --fetch after reading the data terms")
            temporary = path.with_suffix(path.suffix + ".download")
            try:
                with urllib.request.urlopen(url, timeout=60) as source, temporary.open("wb") as output:
                    shutil.copyfileobj(source, output)
                if checksum(temporary) != expected:
                    raise RuntimeError(f"Unexpected downloaded content: {url}")
                temporary.replace(path)
            finally:
                temporary.unlink(missing_ok=True)
        if checksum(path) != expected:
            raise RuntimeError(f"Checksum mismatch: {path}; do not substitute different inputs")

    cadd_path = cache / "cadd_v1.7.tsv"
    if not cadd_path.exists():
        if not fetch:
            raise RuntimeError("Missing CADD extract; run with --fetch")
        tabix = shutil.which("tabix")
        if tabix is None:
            raise RuntimeError("Install the tabix executable from HTSlib for the CADD regional query")
        # This downloads an index and a small genomic region, not the 79 GB full table.
        header = subprocess.check_output([tabix, "-H", CADD_URL], cwd=cache, text=True, timeout=300)
        columns = header.splitlines()[-1].lstrip("#").split("\t")
        names = ["Chrom", "Pos", "Ref", "Alt", "RawScore", "PHRED"]
        keep = [columns.index(name) for name in names]
        body = subprocess.check_output([tabix, CADD_URL, CADD_REGION], cwd=cache, text=True, timeout=300)
        temporary = cadd_path.with_suffix(".download")
        try:
            with temporary.open("w", newline="\n") as output:
                output.write("\t".join(names) + "\n")
                for line in body.splitlines():
                    values = line.split("\t")
                    output.write("\t".join(values[i] for i in keep) + "\n")
            if checksum(temporary) != CADD_SHA256:
                raise RuntimeError("CADD regional extract differs from the published study input")
            index = cache / "whole_genome_SNVs.tsv.gz.tbi"
            if index.exists() and checksum(index, "md5") != CADD_INDEX_MD5:
                raise RuntimeError("CADD tabix index differs from the published study index")
            temporary.replace(cadd_path)
        finally:
            temporary.unlink(missing_ok=True)
    if checksum(cadd_path) != CADD_SHA256:
        raise RuntimeError("Cached CADD extract checksum mismatch")


def joined_scores(cache: Path) -> pd.DataFrame:
    """Match genome build, position and both alleles; refuse silent row loss."""
    raw = pd.read_excel(cache / "labels.xlsx", header=2)
    if not set(raw["func.class"]).issubset({"LOF", "FUNC", "INT"}):
        raise RuntimeError("Unexpected functional label")
    variants = pd.DataFrame({
        "id": [f"chr17:{p}:{r}>{a}" for p, r, a in zip(raw["position (hg19)"], raw["reference"], raw["alt"])],
        "pos": raw["position (hg19)"].astype(int),
        "consequence": raw["consequence"],
        "lof": (raw["func.class"] == "LOF").astype(int),
        "cadd_v1.3": raw["CADD.score"],
    }).set_index("id")
    variants["region"] = np.where(variants["consequence"].isin(CODING), "coding", "noncoding")
    measured = pd.read_csv(cache / "evo2_scores.tsv", sep="\t").set_index("id")
    cadd = pd.read_csv(cache / "cadd_v1.7.tsv", sep="\t", dtype={"Chrom": str})
    cadd["id"] = "chr" + cadd["Chrom"] + ":" + cadd["Pos"].astype(str) + ":" + cadd["Ref"] + ">" + cadd["Alt"]
    cadd = cadd.set_index("id")
    if any(frame.index.has_duplicates for frame in (variants, measured, cadd)):
        raise RuntimeError("Duplicate variant keys; inspect rather than silently collapsing them")
    if set(variants.index) != set(measured.index):
        raise RuntimeError("Stored Evo 2 variants and label variants do not match")
    # Negative likelihood delta is used so larger values indicate loss of function.
    variants["evo2_7b"] = -measured["delta"].reindex(variants.index)
    variants["cadd_v1.7"] = cadd["RawScore"].reindex(variants.index)
    columns = ["evo2_7b", "cadd_v1.3", "cadd_v1.7"]
    if variants[columns].isna().any().any() or not np.isfinite(variants[columns].to_numpy()).all():
        raise RuntimeError("Missing/nonfinite scores; do not compare different subsets")
    if len(variants) != 3893 or int(variants["lof"].sum()) != 823:
        raise RuntimeError("Unexpected BRCA1 cohort")
    return variants


def auroc(labels: np.ndarray, scores: np.ndarray) -> float:
    """Mann–Whitney AUROC, counting tied positive/negative scores as half."""
    labels = labels.astype(bool)
    positives = int(labels.sum())
    negatives = len(labels) - positives
    if not positives or not negatives:
        return float("nan")
    order = np.argsort(scores, kind="mergesort")
    ordered = scores[order]
    starts = np.r_[0, np.flatnonzero(ordered[1:] != ordered[:-1]) + 1]
    ends = np.r_[starts[1:], len(scores)]
    ranks = np.empty(len(scores), dtype=float)
    ranks[order] = np.repeat((starts + ends + 1) / 2, ends - starts)
    return float((ranks[labels].sum() - positives * (positives + 1) / 2) / (positives * negatives))


def paired_bootstrap(frame: pd.DataFrame) -> dict:
    """Resample positions, preserving alternate alleles at the same position together."""
    labels = frame["lof"].to_numpy()
    positions = frame["pos"].to_numpy()
    members = [np.flatnonzero(positions == p) for p in np.unique(positions)]
    names = ["evo2_7b", "cadd_v1.3", "cadd_v1.7"]
    scores = {name: frame[name].to_numpy() for name in names}
    point = {name: auroc(labels, scores[name]) for name in names}
    draws = {name: np.empty(BOOTSTRAPS) for name in names}
    rng = np.random.default_rng(SEED)
    for draw in range(BOOTSTRAPS):
        indices = np.concatenate([members[i] for i in rng.integers(0, len(members), len(members))])
        for name in names:
            draws[name][draw] = auroc(labels[indices], scores[name][indices])
    result = {}
    for name in names:
        result[name] = {"auroc": point[name], "ci95": np.nanpercentile(draws[name], [2.5, 97.5]).tolist()}
        if name != "evo2_7b":
            interval = np.nanpercentile(draws["evo2_7b"] - draws[name], [2.5, 97.5]).tolist()
            result[name].update({
                "difference_vs_evo2_7b": point["evo2_7b"] - point[name],
                "difference_vs_evo2_7b_ci95": interval,
                "evo2_advantage": "holds" if interval[0] > 0 else "reverses" if interval[1] < 0 else "not detected",
            })
    return {"n": len(frame), "lof": int(labels.sum()), "positions": len(members), "auroc": result}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--fetch", action="store_true", help="Download declared inputs after reading their terms")
    parser.add_argument("--cache", type=Path, default=Path(".brca1-cache"))
    parser.add_argument("--output", type=Path, default=Path("brca1-reproduced.json"))
    args = parser.parse_args()
    start = time.monotonic()
    acquire(args.cache.resolve(), args.fetch)
    variants = joined_scores(args.cache)
    published = json.loads((args.cache / "published_results.json").read_text())
    results = {}
    for region in ("all", "coding", "noncoding"):
        frame = variants if region == "all" else variants[variants["region"] == region]
        results[region] = paired_bootstrap(frame)
        for name, actual in results[region]["auroc"].items():
            expected = published["main"][region]["auroc"][name]
            for key in ("auroc", "ci95", "difference_vs_evo2_7b", "difference_vs_evo2_7b_ci95"):
                if key in actual and not np.allclose(actual[key], expected[key], atol=1e-12, rtol=0):
                    raise RuntimeError(f"Published-result mismatch: {region}/{name}/{key}")
    output = {
        "retrieved_or_checked_utc": datetime.now(timezone.utc).isoformat(),
        "study_revision": STUDY_REVISION,
        "exercise": "Recalculation from stored predictions and published experimental labels; no new inference or clinical validation",
        "score_direction": "Evo 2 uses -delta; higher means more likely experimental loss of function",
        "label": "LOF versus FUNC and INT in the Findlay BRCA1 saturation-genome-editing assay",
        "bootstrap": {"group": "genomic position", "draws": BOOTSTRAPS, "seed": SEED},
        "coverage": {"variants": len(variants), "missing_scores": 0},
        "sources": {name: {"url": url, "sha256": digest} for name, (url, digest) in FILES.items()},
        "cadd": {"url": CADD_URL, "region": CADD_REGION, "extract_sha256": CADD_SHA256},
        "results": results,
        "published_result_check": "passed; all AUROCs and paired intervals reproduce to 1e-12",
        "elapsed_seconds": round(time.monotonic() - start, 2),
    }
    args.output.write_text(json.dumps(output, indent=2) + "\n")
    print("Region     n     Evo 2   CADD1.3 CADD1.7  Evo2-CADD1.7 [95% CI]")
    for region, result in results.items():
        values = result["auroc"]
        delta = values["cadd_v1.7"]
        low, high = delta["difference_vs_evo2_7b_ci95"]
        print(f"{region:10} {result['n']:4}  {values['evo2_7b']['auroc']:.3f}   {values['cadd_v1.3']['auroc']:.3f}   {values['cadd_v1.7']['auroc']:.3f}   {delta['difference_vs_evo2_7b']:+.3f} [{low:+.3f}, {high:+.3f}]")
    print(f"Published values reproduced. Wrote {args.output} in {output['elapsed_seconds']:.1f} seconds.")


if __name__ == "__main__":
    main()
