#!/usr/bin/env python3
"""Replicate concordance of `pgx star-allele call`: the top genotype of every gene compared between the replicate
arrays of one individual (PGX-029 reproducibility on real replicates).

    python scripts/pgx/replicate_concordance.py --calls STAR_DIR/pgx_calls.json --identity identity.tsv --out rep.json
"""

from __future__ import annotations

import argparse
import csv
import itertools
import json
from collections import defaultdict
from pathlib import Path


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--calls", required=True, type=Path)
    ap.add_argument("--identity", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    a = ap.parse_args()
    with a.identity.open() as fh:
        ident = {r["sample_id"]: r["individual"] for r in csv.DictReader((ln for ln in fh if not ln.startswith("#")),
                                                                           delimiter="\t")}  # fmt: skip
    calls: dict[str, dict[str, dict[str, str]]] = defaultdict(dict)
    for s in json.loads(a.calls.read_text())["samples"]:
        ind = ident.get(s["sample_id"])
        if ind:
            calls[ind][s["sample_id"]] = {g["gene"]: (g["candidates"][0]["genotype"] if g["candidates"] else "")
                                          for g in s["genes"]}  # fmt: skip
    per_gene: dict[str, list[int]] = defaultdict(lambda: [0, 0])
    discord = []
    for ind, reps in sorted(calls.items()):
        for (_s1, c1), (_s2, c2) in itertools.combinations(sorted(reps.items()), 2):
            for gene in sorted(set(c1) & set(c2)):
                ok = c1[gene] == c2[gene]
                per_gene[gene][0] += ok
                per_gene[gene][1] += 1
                if not ok:
                    discord.append({"individual": ind, "gene": gene, "a": c1[gene], "b": c2[gene]})
    tot = [sum(v[0] for v in per_gene.values()), sum(v[1] for v in per_gene.values())]
    out = {"pairs_by_gene": {g: {"concordant": v[0], "pairs": v[1]} for g, v in sorted(per_gene.items())},
           "total": {"concordant": tot[0], "pairs": tot[1], "rate": round(tot[0] / tot[1], 4) if tot[1] else None},
           "discordant": discord}  # fmt: skip
    a.out.write_text(json.dumps(out, indent=1))
    print(json.dumps(out["total"]), f"{len(discord)} discordant gene pair(s)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())