#!/usr/bin/env python3
"""Reference analysis — 1000 Genomes super-population allele-frequency comparison.

Question shape: "Does variant <rsID> (or the variants in region <R>) differ in
frequency across the five 1000G super-populations (AFR/AMR/EAS/EUR/SAS)?"
Textbook use: population-genetics differentiation at a locus (e.g. LCT/lactase,
APOE, SLC24A5 pigmentation).

Input:
  --kg   CSV from fetch_1000g.py (chrom,pos,id,ref,alt,af,eas_af,amr_af,afr_af,eur_af,sas_af)
Output:
  figure.png  — grouped bar chart of per-super-population AF for the top variant(s)
  stats.json  — chi-square test of allele-count homogeneity across populations,
                plus a simple pairwise Fst (Weir-Cockerham-style) for the lead SNP

Real statistic: for the most differentiated variant, a chi-square test on the
2x5 contingency table of (alt allele count, ref allele count) across the five
super-populations (nominal sample sizes assumed from 1000G phase-3 super-pop N),
and a pairwise Fst matrix. A null result (no population differs) is reported.

Deterministic: no randomness in the statistic; --seed only affects nothing here
but is recorded for provenance parity across templates.

Usage:
    python3 population_allele_freq.py --kg work/kg_apoe.csv \
        --label "APOE region" --seed 1234 --outdir work

Exit: 0 computed, 1 could not compute, 2 usage.
"""

from __future__ import annotations

import argparse
import json
import sys

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy import stats

# 1000G phase-3 super-population sample sizes (2N haplotypes). These are fixed
# public constants of the panel, used to turn AF back into integer allele counts
# for the chi-square contingency table.
SUPERPOP_2N = {"afr": 1322, "amr": 694, "eas": 1008, "eur": 1006, "sas": 978}
POPS = ["afr", "amr", "eas", "eur", "sas"]


def fst_pair(p1, p2, n1, n2):
    """Hudson-style Fst for a biallelic locus between two populations."""
    if p1 in (0, 1) and p1 == p2:
        return 0.0
    num = (p1 - p2) ** 2 - p1 * (1 - p1) / (n1 - 1) - p2 * (1 - p2) / (n2 - 1)
    den = p1 * (1 - p2) + p2 * (1 - p1)
    return float(num / den) if den > 0 else 0.0


def main() -> int:
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--kg", required=True)
    ap.add_argument("--label", default="LCT/MCM6 lactase-persistence region")
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--outdir", default=".")
    args = ap.parse_args()

    df = pd.read_csv(args.kg, dtype=str).fillna("")
    for p in POPS:
        df[f"{p}_af"] = pd.to_numeric(df.get(f"{p}_af"), errors="coerce")
    df = df.dropna(subset=[f"{p}_af" for p in POPS])
    if df.empty:
        _fail(
            args.outdir, args.label, "no variants with all-population AF in the slice"
        )
        return 1

    # Lead variant = largest spread across populations (max-min AF).
    afmat = df[[f"{p}_af" for p in POPS]].to_numpy(dtype=float)
    spread = afmat.max(axis=1) - afmat.min(axis=1)
    lead_i = int(np.argmax(spread))
    lead = df.iloc[lead_i]
    lead_af = {p: float(lead[f"{p}_af"]) for p in POPS}

    # Chi-square on the 2x5 (alt, ref) allele-count table for the lead variant.
    alt_counts = [round(lead_af[p] * SUPERPOP_2N[p]) for p in POPS]
    ref_counts = [SUPERPOP_2N[p] - a for a in alt_counts]
    table = np.array([alt_counts, ref_counts])
    result = {
        "label": args.label,
        "seed": args.seed,
        "reference_template": "population_allele_freq",
        "n_variants_in_slice": int(len(df)),
        "lead_variant": lead.get("id") or f"{lead['chrom']}:{lead['pos']}",
        "lead_variant_af_by_pop": lead_af,
    }
    if table.min() < 0 or table.sum() == 0:
        _fail(
            args.outdir,
            args.label,
            "degenerate allele-count table for the lead variant",
        )
        return 1

    chi2, p_value, dof, _ = stats.chi2_contingency(table)
    fst = {
        f"{a}-{b}": round(
            fst_pair(lead_af[a], lead_af[b], SUPERPOP_2N[a], SUPERPOP_2N[b]), 4
        )
        for i, a in enumerate(POPS)
        for b in POPS[i + 1 :]
    }
    max_fst_pair = max(fst, key=fst.get)

    result.update(
        {
            "outcome": "success" if p_value < 0.05 else "null_result",
            "test": "chi-square of allele-count homogeneity across 5 super-populations",
            "chi2": float(chi2),
            "dof": int(dof),
            "p_value": float(p_value),
            "pairwise_fst": fst,
            "max_fst_pair": max_fst_pair,
            "max_fst": fst[max_fst_pair],
            "significant_at_0.05": bool(p_value < 0.05),
            "headline_statistic": f"chi2 p = {p_value:.2e}, max Fst = {fst[max_fst_pair]:.3f} ({max_fst_pair})",
        }
    )
    with open(f"{args.outdir}/stats.json", "w") as fh:
        json.dump(result, fh, indent=2)

    _render(args.outdir, lead_af, args.label, result)
    print(json.dumps(result, indent=2))
    return 0


def _fail(outdir, label, reason):
    result = {
        "label": label,
        "reference_template": "population_allele_freq",
        "outcome": "failed",
        "reason": reason,
    }
    with open(f"{outdir}/stats.json", "w") as fh:
        json.dump(result, fh, indent=2)
    print(json.dumps(result, indent=2))


def _render(outdir, lead_af, label, result):
    fig, ax = plt.subplots(figsize=(6, 4.5))
    colors = ["#8e44ad", "#e67e22", "#16a085", "#2980b9", "#c0392b"]
    ax.bar(
        [p.upper() for p in POPS], [lead_af[p] for p in POPS], color=colors, alpha=0.85
    )
    ax.set_ylabel("Alt allele frequency")
    ax.set_title(
        f"{label}: {result['lead_variant']} across 1000G super-populations\n"
        f"{result['headline_statistic']}",
        fontsize=10,
    )
    ax.set_ylim(0, 1)
    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())