#!/usr/bin/env python3
"""Adapted ANCHOR reference analysis — two-gene ClinVar x gnomAD gap comparison.

Question it answers: "For ABCA4, is the ClinVar pathogenic-vs-benign gnomAD
allele-frequency gap larger than for a comparable housekeeping gene (ACTB)?"

For each of two genes, the script runs the anchor clinvar_gnomad_ensembl
within-gene analysis: a Mann-Whitney U test (two-sided) on gnomAD population
allele frequency (AF) between ClinVar pathogenic/likely-pathogenic and
benign/likely-benign variants. It then defines each gene's "gap" as

    gap = log10(median benign AF) - log10(median pathogenic AF)

i.e. the log10-fold separation between the two groups' median population AFs.
The cross-gene comparison is the observed gap difference
(ABCA4 gap - ACTB gap). A fixed-seed bootstrap resampling of each gene's two AF
groups estimates the one-sided p-value for the directional hypothesis that
ABCA4's gap exceeds ACTB's (p_boot = P(diff_boot <= 0)); this bootstrap
statistic is reproducible with the pinned --seed.

Real statistics: the per-gene Mann-Whitney U p-values and the cross-gene gap
difference are fully deterministic (no randomness). The bootstrap p-value is the
only seeded/sampled number and is deterministic given the fixed seed.

Inputs (produced by the fetchers, all zero-auth small API calls):
  --clinvar   CSV from fetch_clinvar.py for the primary gene (ABCA4)
  --gnomad    CSV from fetch_gnomad.py  for the primary gene (ABCA4)
  --clinvar2  CSV from fetch_clinvar.py for the housekeeping gene (ACTB)
  --gnomad2   CSV from fetch_gnomad.py  for the housekeeping gene (ACTB)
  --gene      primary gene symbol (ABCA4)
  --gene2     housekeeping gene symbol (ACTB)
Output:
  figure.png  — log10(AF) distributions, pathogenic vs benign, one panel per gene
  stats.json  — per-gene Mann-Whitney U stats, medians, gaps, bootstrap p-value

Exit: 0 computed a stat (significant OR null — both are success),
      1 could not compute (disclosed failure), 2 usage.
"""
from __future__ import annotations

import argparse
import json
import sys

import matplotlib
matplotlib.use("Agg")  # headless PNG rendering — no display, server-safe
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy import stats

PATHOGENIC = ("pathogenic", "likely pathogenic", "likely_pathogenic")
BENIGN = ("benign", "likely benign", "likely_benign")


def classify(sig: str) -> str | None:
    s = (sig or "").strip().lower()
    if not s:
        return None
    if any(p in s for p in PATHOGENIC) and "conflicting" not in s:
        return "pathogenic"
    if any(b in s for b in BENIGN) and "conflicting" not in s:
        return "benign"
    return None


def analyze_gene(clinvar_path: str, gnomad_path: str, gene: str) -> dict:
    """Run the anchor per-gene analysis and return the measured arrays + stats."""
    clinvar = pd.read_csv(clinvar_path, dtype=str).fillna("")
    gnomad = pd.read_csv(gnomad_path, dtype=str).fillna("")
    gnomad["af"] = pd.to_numeric(gnomad["af"], errors="coerce")

    clinvar["call"] = clinvar["clinical_significance"].map(classify)
    clinvar = clinvar.dropna(subset=["call"])

    both_have_rs = (clinvar["rsid"].str.startswith("rs").any()
                    and gnomad["rsid"].str.startswith("rs").any())
    gnomad = gnomad.dropna(subset=["af"])
    gnomad = gnomad[gnomad["af"] > 0].copy()

    if both_have_rs:
        join_key, join_desc = "rsid", "dbSNP rsID"
        gnomad_j = (gnomad[gnomad["rsid"].str.startswith("rs")]
                    .sort_values("af", ascending=False)
                    .drop_duplicates(subset=["rsid"]))
        clinvar_j = clinvar[clinvar["rsid"].str.startswith("rs")].copy()
    else:
        join_key, join_desc = "locus", "genomic locus chrom:pos (ref/alt-agnostic)"
        gnomad["locus"] = gnomad["chrom"].astype(str) + ":" + gnomad["pos"].astype(str)
        clinvar["locus"] = clinvar["chrom"].astype(str) + ":" + clinvar["pos"].astype(str)
        gnomad_j = gnomad.sort_values("af", ascending=False).drop_duplicates(subset=["locus"])
        clinvar_j = clinvar.copy()

    merged = clinvar_j.merge(
        gnomad_j[[join_key, "af", "consequence"]], on=join_key, how="inner",
    ).drop_duplicates(subset=[join_key])

    path_af = merged.loc[merged["call"] == "pathogenic", "af"].to_numpy(dtype=float)
    ben_af = merged.loc[merged["call"] == "benign", "af"].to_numpy(dtype=float)

    out = {
        "gene": gene,
        "join_key": join_desc,
        "n_clinvar_classified": int(len(clinvar)),
        "n_merged_with_gnomad_af": int(len(merged)),
        "n_pathogenic": int(len(path_af)),
        "n_benign": int(len(ben_af)),
        "median_af_pathogenic": _safe_median(path_af),
        "median_af_benign": _safe_median(ben_af),
        "path_af": path_af,
        "ben_af": ben_af,
    }
    if len(path_af) >= 3 and len(ben_af) >= 3:
        u_stat, p_value = stats.mannwhitneyu(path_af, ben_af, alternative="two-sided")
        med_path = float(np.median(path_af))
        med_ben = float(np.median(ben_af))
        gap = float(np.log10(med_ben) - np.log10(med_path))
        out.update({
            "mw_u_statistic": float(u_stat),
            "mw_p_value": float(p_value),
            "direction": ("pathogenic rarer" if med_path < med_ben
                          else "pathogenic not rarer"),
            "gap_log10": gap,
            "test_ran": True,
        })
    else:
        out.update({"test_ran": False, "gap_log10": None})
    return out


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--clinvar", required=True)
    ap.add_argument("--gnomad", required=True)
    ap.add_argument("--clinvar2", required=True)
    ap.add_argument("--gnomad2", required=True)
    ap.add_argument("--gene", required=True)
    ap.add_argument("--gene2", required=True)
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--outdir", default=".")
    ap.add_argument("--bootstrap-reps", type=int, default=10000)
    args = ap.parse_args()

    rng = np.random.default_rng(args.seed)  # pins strip-plot jitter + bootstrap

    g1 = analyze_gene(args.clinvar, args.gnomad, args.gene)
    g2 = analyze_gene(args.clinvar2, args.gnomad2, args.gene2)

    result = {
        "gene": args.gene,
        "housekeeping_gene": args.gene2,
        "seed": args.seed,
        "bootstrap_reps": args.bootstrap_reps,
        "reference_template": "clinvar_gnomad_ensembl",
        "join_key": g1["join_key"],
        "genes": {
            g1["gene"]: {k: g1[k] for k in g1 if k not in ("path_af", "ben_af")},
            g2["gene"]: {k: g2[k] for k in g2 if k not in ("path_af", "ben_af")},
        },
    }

    if not (g1["test_ran"] and g2["test_ran"]):
        result["outcome"] = "null_result"
        result["reason"] = (
            "fewer than 3 variants in at least one group for one gene after the "
            "ClinVar x gnomAD join; the within-gene test and the cross-gene "
            "comparison were not performed (reported honestly, not fabricated)"
        )
        _write_stats(args.outdir, result)
        _render(args.outdir, g1, g2, args.gene, result, rng)
        print(json.dumps(result, indent=2))
        return 0

    # Deterministic cross-gene gap difference (ABCA4 gap - housekeeping gap).
    gap_diff = float(g1["gap_log10"] - g2["gap_log10"])

    # Seeded bootstrap resampling: re-estimate each gene's gap and compare.
    diffs = np.empty(args.bootstrap_reps, dtype=float)
    for i in range(args.bootstrap_reps):
        p1 = rng.choice(g1["path_af"], size=len(g1["path_af"]), replace=True)
        b1 = rng.choice(g1["ben_af"], size=len(g1["ben_af"]), replace=True)
        p2 = rng.choice(g2["path_af"], size=len(g2["path_af"]), replace=True)
        b2 = rng.choice(g2["ben_af"], size=len(g2["ben_af"]), replace=True)
        gap1 = float(np.log10(np.median(b1)) - np.log10(np.median(p1)))
        gap2 = float(np.log10(np.median(b2)) - np.log10(np.median(p2)))
        diffs[i] = gap1 - gap2
    p_boot = float((np.sum(diffs <= 0.0) + 1.0) / (args.bootstrap_reps + 1.0))

    result.update({
        "test": "per-gene Mann-Whitney U (two-sided) on allele frequency + "
                "seeded bootstrap one-sided p for the cross-gene gap difference",
        "p_value": float(p_boot),
        "significant_at_0.05": bool(p_boot < 0.05),
        "gap_abca4_log10": float(g1["gap_log10"]),
        "gap_housekeeping_log10": float(g2["gap_log10"]),
        "gap_difference_log10": gap_diff,
        "bootstrap_p_one_sided": p_boot,
        "gap_difference_direction": ("ABCA4 gap larger" if gap_diff > 0
                                     else "ABCA4 gap not larger"),
        "outcome": "success" if p_boot < 0.05 else "null_result",
        "headline_statistic": (
            f"gap difference {gap_diff:.4g} log10 AF (ABCA4 - ACTB); "
            f"bootstrap one-sided p = {p_boot:.3g}"
        ),
    })
    _write_stats(args.outdir, result)
    _render(args.outdir, g1, g2, args.gene, result, rng)
    print(json.dumps(result, indent=2))
    return 0


def _safe_median(a):
    return float(np.median(a)) if len(a) else None


def _write_stats(outdir, result):
    with open(f"{outdir}/stats.json", "w") as fh:
        json.dump(result, fh, indent=2)


def _render(outdir, g1, g2, gene, result, rng):
    fig, axes = plt.subplots(1, 2, figsize=(11, 5), sharey=True)
    for ax, g, col in ((axes[0], g1, "#c0392b"), (axes[1], g2, "#8e44ad")):
        groups, labels, colors = [], [], []
        if len(g["path_af"]):
            groups.append(np.log10(g["path_af"])); labels.append(f"Pathogenic\n(n={len(g['path_af'])})"); colors.append("#c0392b")
        if len(g["ben_af"]):
            groups.append(np.log10(g["ben_af"])); labels.append(f"Benign\n(n={len(g['ben_af'])})"); colors.append("#2980b9")
        bp = ax.boxplot(groups, patch_artist=True, widths=0.5, showfliers=False)
        for patch, color in zip(bp["boxes"], colors):
            patch.set_facecolor(color); patch.set_alpha(0.35)
        for i, (gr, color) in enumerate(zip(groups, colors), start=1):
            jitter = rng.uniform(-0.12, 0.12, size=len(gr))  # seeded jitter
            ax.scatter(np.full(len(gr), i) + jitter, gr, s=16, color=color,
                       alpha=0.7, edgecolors="none", zorder=3)
        ax.set_xticks(range(1, len(labels) + 1)); ax.set_xticklabels(labels)
        subtitle = ""
        if g.get("gap_log10") is not None:
            subtitle = f"gap = {g['gap_log10']:.3f} log10 AF"
        ax.set_title(f"{g['gene']}\n{subtitle}", fontsize=10)
        ax.set_ylabel("log10(gnomAD allele frequency)")
        ax.grid(axis="y", alpha=0.3)
    headline = result.get("headline_statistic") or result.get("reason", "")
    fig.suptitle(f"ClinVar significance vs gnomAD frequency: {gene} vs {g2['gene']}\n{headline}",
                 fontsize=9)
    fig.tight_layout(rect=[0, 0, 1, 0.94])
    fig.savefig(f"{outdir}/figure.png", dpi=130)
    plt.close(fig)


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