#!/usr/bin/env python3
"""Extract GeT-RM consensus diplotypes (CDC Genetic Testing Reference Material program, consolidated PGx/HLA table,
public) for the individuals of an identity table.

    python scripts/pgx/getrm_truth.py --xlsx Consolidated_PGx-HLA_table_1-22-25-V4.xlsx \
        --identity scripts/pgx/demo36_identity.tsv --out getrm_truth.tsv
Writes ``individual, gene, diplotype`` (the table's text, trimmed) for the star-allele genes.
"""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

GENES = ("CYP2B6", "CYP2C9", "CYP2C19", "CYP2D6", "CYP3A4", "CYP3A5", "CYP4F2", "NAT2", "NUDT15", "SLCO1B1", "TPMT",
         "UGT1A1", "DPYD", "CYP2A6", "CYP1A2", "CYP2C8", "GSTM1", "GSTT1", "UGT2B17")  # fmt: skip


def main() -> int:
    import openpyxl

    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--xlsx", 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:
        inds = {r["individual"] for r in csv.DictReader((ln for ln in fh if not ln.startswith("#")), delimiter="\t")}
    wb = openpyxl.load_workbook(a.xlsx, read_only=True)
    rows = list(wb[wb.sheetnames[0]].iter_rows(values_only=True))
    head = [str(h or "").split("\n")[0].strip() for h in rows[1]]
    col = {g: head.index(g) for g in GENES if g in head}
    vk = next((i for i, h in enumerate(rows[1]) if str(h or "").startswith("VKORC1") and "rs9923231" in str(h)
               and "Reference" not in str(h)), None)  # fmt: skip
    if vk is not None:
        col["VKORC1"] = vk  # c.-1639G>A genotype, e.g. "GA"
    out = []
    for r in rows[2:]:
        if not r or not r[1] or str(r[1]).strip() not in inds:
            continue
        for g, i in col.items():
            v = r[i]
            if v is not None and str(v).strip():
                out.append((str(r[1]).strip(), g, " ".join(str(v).split())))
    with a.out.open("w") as fh:
        fh.write("individual\tgene\tdiplotype\n")
        for t in sorted(out):
            fh.write("\t".join(t) + "\n")
    print(f"{len({t[0] for t in out})} individual(s), {len(out)} gene call(s) -> {a.out}")
    return 0


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