#!/usr/bin/env python3
"""ANCHOR reference analysis — ClinVar x gnomAD x Ensembl cross-annotation.

*** This is Gregor's guaranteed demo path. Keep it bulletproof and fast. ***

Question shape it answers: "For gene <G>, how do population allele frequencies
(gnomAD) differ between variants ClinVar calls pathogenic/likely-pathogenic vs
benign/likely-benign?" — the textbook expectation is that pathogenic variants
are RARER. We test that with a real statistic and render a real figure.

Inputs (produced by the fetchers, all zero-auth small API calls):
  --clinvar   CSV from fetch_clinvar.py  (rsid, clinical_significance, ...)
  --gnomad    CSV from fetch_gnomad.py   (rsid, af, consequence, ...)
Output:
  figure.png  — log10(AF) distribution, pathogenic vs benign (box + strip)
  stats.json  — Mann-Whitney U p-value, medians, group sizes, effect direction

Real statistic: a Mann-Whitney U test on log10(allele frequency) between the
ClinVar pathogenic group and the benign group (non-parametric — AF is heavily
skewed). A null/degenerate result (one group empty, or no significant
difference) is REPORTED HONESTLY, never massaged.

Deterministic: fixed seed pins the strip-plot jitter so the figure is
reproducible. No randomness enters the statistic itself.

Usage:
    python3 anchor_clinvar_gnomad_ensembl.py \
        --clinvar work/clinvar_apoe.csv --gnomad work/gnomad_apoe.csv \
        --gene APOE --seed 1234 --outdir work

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
    # A ClinVar record can carry a compound significance; take the dominant call.
    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 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("--gene", required=True)
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--outdir", default=".")
    args = ap.parse_args()

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

    clinvar = pd.read_csv(args.clinvar, dtype=str).fillna("")
    gnomad = pd.read_csv(args.gnomad, dtype=str).fillna("")
    gnomad["af"] = pd.to_numeric(gnomad["af"], errors="coerce")

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

    # Join key: prefer rsID when BOTH sources carry it, else fall back to the
    # genomic locus chrom:pos. Real-world reason (verified against live data):
    # ClinVar's esummary JSON reliably gives chrom+pos but often omits ref/alt
    # and the dbSNP rsID, while gnomAD always has chrom:pos:ref:alt (+rsID). So
    # a coordinate join is what actually connects the two sources. We collapse
    # gnomAD to the MAX-AF variant per locus (a locus can be multiallelic) and
    # disclose the join granularity in the paper — never silently.
    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)

    result = {
        "gene": args.gene,
        "seed": args.seed,
        "reference_template": "clinvar_gnomad_ensembl",
        "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)),
    }

    if len(path_af) < 3 or len(ben_af) < 3:
        # Honest null: not enough overlapping annotated variants to test.
        result["outcome"] = "null_result"
        result["reason"] = (
            "fewer than 3 variants in one group after the ClinVar x gnomAD join; "
            "no test performed (reported honestly, not fabricated)"
        )
        result["median_af_pathogenic"] = _safe_median(path_af)
        result["median_af_benign"] = _safe_median(ben_af)
        _write_stats(args.outdir, result)
        # Still render whatever we have so the paper has a figure.
        _render(args.outdir, path_af, ben_af, args.gene, result, rng)
        print(json.dumps(result, indent=2))
        return 0

    # Mann-Whitney U on the raw AFs (rank-based; equivalent on log scale).
    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))
    result.update({
        "outcome": "success",
        "test": "Mann-Whitney U (two-sided) on allele frequency",
        "u_statistic": float(u_stat),
        "p_value": float(p_value),
        "median_af_pathogenic": med_path,
        "median_af_benign": med_ben,
        "direction": ("pathogenic rarer" if med_path < med_ben
                      else "pathogenic not rarer"),
        "significant_at_0.05": bool(p_value < 0.05),
        "headline_statistic": f"p = {p_value:.2e} (Mann-Whitney U)",
    })
    _write_stats(args.outdir, result)
    _render(args.outdir, path_af, ben_af, 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, path_af, ben_af, gene, result, rng):
    fig, ax = plt.subplots(figsize=(6, 5))
    groups, labels, colors = [], [], []
    if len(path_af):
        groups.append(np.log10(path_af)); labels.append(f"Pathogenic\n(n={len(path_af)})"); colors.append("#c0392b")
    if len(ben_af):
        groups.append(np.log10(ben_af)); labels.append(f"Benign\n(n={len(ben_af)})"); colors.append("#2980b9")

    if groups:
        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, (g, color) in enumerate(zip(groups, colors), start=1):
            jitter = rng.uniform(-0.12, 0.12, size=len(g))  # seeded jitter
            ax.scatter(np.full(len(g), i) + jitter, g, s=18, color=color,
                       alpha=0.7, edgecolors="none", zorder=3)
        ax.set_xticks(range(1, len(labels) + 1)); ax.set_xticklabels(labels)

    ax.set_ylabel("log10(gnomAD allele frequency)")
    subtitle = result.get("headline_statistic") or result.get("reason", "")
    ax.set_title(f"{gene}: ClinVar significance vs gnomAD frequency\n{subtitle}",
                 fontsize=10)
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    fig.savefig(f"{outdir}/figure.png", dpi=130)
    plt.close(fig)


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