#!/usr/bin/env python3
"""Build the PGx CN target-region BED (spec PCN-001) from UCSC ncbiRefSeq (GRCh38).

    python scripts/pgx/build_cn_regions.py ncbiRefSeq.hg38.txt.gz OUT.bed [--version N]

Regions (0-based half-open BED: chrom start0 end region gene), from the MANE/first-listed transcript:
  CYP2D6.p5        2 kb 5' flank of CYP2D6 (upstream of exon 1; the gene is on the minus strand)
  CYP2D6.intron.2  CYP2D6 intron 2
  CYP2D6.exon.9    CYP2D6 exon 9
  CYP2A6.exon.1    CYP2A6 exon 1 (± 500 bp)
  CYP2A6.intron.7  CYP2A6 intron 7
  GSTM1, GSTT1 (on chr22_KI270879v1_alt, IN-062), UGT2B17, SULT1A1, CYP2E1: whole gene.
Exons and introns are numbered in transcript order (strand-aware). [DB: NCBI RefSeq via UCSC ncbiRefSeq]
"""

from __future__ import annotations

import argparse
import gzip
from pathlib import Path

TRANSCRIPT = {"CYP2D6": "NM_000106", "CYP2A6": "NM_000762", "GSTM1": "NM_000561", "GSTT1": "NM_000853",
              "UGT2B17": "NM_001077", "SULT1A1": "NM_001055", "CYP2E1": "NM_000773"}  # fmt: skip
FLANK = 2000
PAD = 500


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("refseq", type=Path)
    ap.add_argument("out", type=Path)
    ap.add_argument("--version", type=int, default=1)
    a = ap.parse_args()
    tx: dict[str, tuple[str, str, list[int], list[int]]] = {}
    with gzip.open(a.refseq, "rt") as fh:
        for line in fh:
            f = line.rstrip("\n").split("\t")
            name, chrom, strand, gene = f[1], f[2], f[3], f[12]
            if gene not in TRANSCRIPT or not name.startswith(TRANSCRIPT[gene] + "."):
                continue
            if "_" in chrom and not (gene == "GSTT1" and chrom == "chr22_KI270879v1_alt"):
                continue  # primary assembly, except GSTT1 which is only on the alt contig
            starts = [int(x) for x in f[9].rstrip(",").split(",")]
            ends = [int(x) for x in f[10].rstrip(",").split(",")]
            tx.setdefault(gene, (chrom, strand, starts, ends))
    missing = set(TRANSCRIPT) - set(tx)
    if missing:
        raise SystemExit(f"transcripts not found: {sorted(missing)}")
    rows: list[tuple[str, int, int, str, str]] = []

    def exon(gene: str, k: int) -> tuple[int, int]:
        chrom, strand, s, e = tx[gene]
        i = k - 1 if strand == "+" else len(s) - k
        return s[i], e[i]

    def intron(gene: str, k: int) -> tuple[int, int]:
        chrom, strand, s, e = tx[gene]
        a_, b_ = exon(gene, k), exon(gene, k + 1)
        return (a_[1], b_[0]) if strand == "+" else (b_[1], a_[0])

    for gene in ("CYP2D6", "CYP2A6"):
        chrom, strand, s, e = tx[gene]
        if gene == "CYP2D6":
            p5 = (e[-1], e[-1] + FLANK) if strand == "-" else (s[0] - FLANK, s[0])
            rows += [(chrom, *p5, "CYP2D6.p5", gene), (chrom, *intron(gene, 2), "CYP2D6.intron.2", gene),
                     (chrom, *exon(gene, 9), "CYP2D6.exon.9", gene)]  # fmt: skip
        else:
            x1 = exon(gene, 1)
            rows += [(chrom, x1[0] - PAD, x1[1] + PAD, "CYP2A6.exon.1", gene),
                     (chrom, *intron(gene, 7), "CYP2A6.intron.7", gene)]  # fmt: skip
    for gene in ("GSTM1", "GSTT1", "UGT2B17", "SULT1A1", "CYP2E1"):
        chrom, strand, s, e = tx[gene]
        rows.append((chrom, s[0], e[-1], gene, gene))
    order = {f"chr{i}": i for i in range(1, 23)}
    rows.sort(key=lambda r: (order.get(r[0].split("_")[0], 99), r[0], r[1]))
    with a.out.open("w") as fh:
        fh.write(f"# simurg PGx CN regions v{a.version}, GRCh38, from UCSC ncbiRefSeq (PCN-001)\n")
        for r in rows:
            fh.write("\t".join(map(str, r)) + "\n")
    print(a.out.read_text())
    return 0


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