"""Count CpG sites in human reference genome FASTA files.

Supports the ProteinIQ guide "How many CpG sites are in the human genome?"
(content/guides/how-many-cpg-sites-in-human-genome.mdx) and writes
cpg-counts.csv next to this script.

Inputs (UCSC soft-masked FASTA, downloaded 2026-10-23; md5 from each
directory's md5sum.txt):

  hg38  https://hgdownload.soe.ucsc.edu/goldenPath/hg38/bigZips/hg38.fa.gz
        GRCh38 initial release, no patches (Last-Modified 2014-01-16),
        983,659,424 bytes, md5 1c9dcaddfa41027f17cd8f7a82c7293b
  hs1   https://hgdownload.soe.ucsc.edu/goldenPath/hs1/bigZips/hs1.fa.gz
        T2T-CHM13v2.0, CHM13 plus HG002 chrY (Last-Modified 2022-07-18),
        974,842,401 bytes, md5 a493d5402cc86ecc3f54f6346d980036
  hg19  https://hgdownload.soe.ucsc.edu/goldenPath/hg19/bigZips/hg19.fa.gz
        GRCh37 (Last-Modified 2018-08-21), 948,731,419 bytes,
        md5 806c02398f5ac5da8ffd6da2d1d5d1a9; used only to reconcile the
        28,217,448 CpGs of the Loyfer et al. (2023) methylation atlas

Reproduce (Python 3.10+, standard library only, about 10 minutes):

  curl -LO https://hgdownload.soe.ucsc.edu/goldenPath/hg38/bigZips/hg38.fa.gz
  curl -LO https://hgdownload.soe.ucsc.edu/goldenPath/hs1/bigZips/hs1.fa.gz
  curl -LO https://hgdownload.soe.ucsc.edu/goldenPath/hg19/bigZips/hg19.fa.gz
  python count_cpg_sites.py --hg38 hg38.fa.gz --hs1 hs1.fa.gz --hg19 hg19.fa.gz

Definitions:

  cpg_dinucleotides   positions where C is followed by G on the given strand,
                      case-insensitive. CG is its own reverse complement, so
                      this equals the number of CpG sites on the double helix.
  cpg_cytosines       2 x cpg_dinucleotides: one C per strand per site, the
                      unit that strand-specific methylation calls report.
  cpg_c_softmasked    CpGs whose C is lowercase (RepeatMasker and TRF repeats
                      in UCSC FASTA). Counted, not removed.
  cpg_in_chrY_par     CpGs inside the chrY pseudoautosomal regions, which
                      repeat chrX sequence. Subtract to count each PAR site once.
  n_bp                N bases (gaps). A CpG cannot include an N, so no CpG is
                      lost to masking; gaps only lower the sequenced length.
  cpg_obs_exp         CpG / (C x G / acgt_bp), with acgt_bp the non-N length.

Sequence categories: primary (chr1-22, X, Y, M), unlocalized (*_random),
unplaced (chrUn_*), alt (*_alt in hg38, *_hap* in hg19).
"""

from __future__ import annotations

import argparse
import csv
import gzip
import re
from collections import defaultdict
from pathlib import Path

PRIMARY = [f"chr{i}" for i in range(1, 23)] + ["chrX", "chrY", "chrM"]

# chrY pseudoautosomal regions, 0-based half-open.
# GRCh38: GRC (https://www.ncbi.nlm.nih.gov/grc/human), PAR1 10,001-2,781,479
# and PAR2 56,887,903-57,217,415 (1-based). hg19 uses the GRCh37 values.
# CHM13v2.0: chm13v2.0_PAR.bed from the T2T analysis set on human-pangenomics S3.
CHRY_PAR = {
    "hg38": [(10_000, 2_781_479), (56_887_902, 57_217_415)],
    "hg19": [(10_000, 2_649_520), (59_034_049, 59_373_566)],
    "hs1": [(0, 2_458_320), (62_122_809, 62_460_029)],
}

FIELDS = [
    "assembly",
    "sequence",
    "category",
    "sequences",
    "length_bp",
    "n_bp",
    "acgt_bp",
    "c_bp",
    "g_bp",
    "softmasked_bp",
    "cpg_dinucleotides",
    "cpg_cytosines",
    "cpg_c_softmasked",
    "cpg_in_chrY_par",
    "cpg_obs_exp",
]

COUNTED = FIELDS[3:-1]


def category(name: str) -> str:
    if name in PRIMARY:
        return "primary"
    if name.endswith("_random"):
        return "unlocalized"
    if name.startswith("chrUn"):
        return "unplaced"
    if name.endswith("_alt") or "_hap" in name:
        return "alt"
    return "other"


def read_fasta(path: Path):
    name, chunks = None, []
    with gzip.open(path, "rb") as handle:
        for line in handle:
            if line.startswith(b">"):
                if name is not None:
                    yield name, b"".join(chunks)
                name, chunks = line[1:].split()[0].decode(), []
            else:
                chunks.append(line.rstrip(b"\r\n"))
    if name is not None:
        yield name, b"".join(chunks)


def count_sequence(assembly: str, name: str, seq: bytes) -> dict:
    upper = seq.upper()
    a, c, g, t = (upper.count(base) for base in (b"A", b"C", b"G", b"T"))
    cpg = upper.count(b"CG")
    par = 0
    if name == "chrY":
        par = sum(upper[s : e + 1].count(b"CG") for s, e in CHRY_PAR[assembly])
    lower = sum(seq.count(base) for base in (b"a", b"c", b"g", b"t", b"n"))
    return {
        "sequences": 1,
        "length_bp": len(seq),
        "n_bp": upper.count(b"N"),
        "acgt_bp": a + c + g + t,
        "c_bp": c,
        "g_bp": g,
        "softmasked_bp": lower,
        "cpg_dinucleotides": cpg,
        "cpg_cytosines": 2 * cpg,
        "cpg_c_softmasked": seq.count(b"cg") + seq.count(b"cG"),
        "cpg_in_chrY_par": par,
    }


def with_ratio(row: dict) -> dict:
    expected = row["c_bp"] * row["g_bp"] / row["acgt_bp"] if row["acgt_bp"] else 0
    row["cpg_obs_exp"] = f"{row['cpg_dinucleotides'] / expected:.4f}" if expected else ""
    return row


def summarize(assembly: str, path: Path) -> list[dict]:
    per_chrom: dict[str, dict] = {}
    groups: dict[str, dict] = defaultdict(lambda: dict.fromkeys(COUNTED, 0))
    for name, seq in read_fasta(path):
        counts = count_sequence(assembly, name, seq)
        cat = category(name)
        if cat == "primary":
            per_chrom[name] = counts
        for key in (cat, "primary_autosomes_xy" if cat == "primary" and name != "chrM" else None, "all"):
            if key:
                for field in COUNTED:
                    groups[key][field] += counts[field]
        print(f"{assembly} {name} {counts['cpg_dinucleotides']:,}", flush=True)

    rows = [
        with_ratio({"assembly": assembly, "sequence": name, "category": "primary", **per_chrom[name]})
        for name in PRIMARY
        if name in per_chrom
    ]
    labels = {
        "primary_autosomes_xy": ("chr1-22+X+Y", "total"),
        "primary": ("chr1-22+X+Y+M", "total"),
        "unlocalized": ("*_random", "unlocalized"),
        "unplaced": ("chrUn_*", "unplaced"),
        "alt": ("alt/haplotype", "alt"),
        "other": ("other", "other"),
        "all": ("all sequences", "total"),
    }
    for key, (label, cat) in labels.items():
        if key in groups:
            rows.append(with_ratio({"assembly": assembly, "sequence": label, "category": cat, **groups[key]}))
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    for assembly in ("hg38", "hs1", "hg19"):
        parser.add_argument(f"--{assembly}", type=Path, help=f"UCSC {assembly}.fa.gz")
    parser.add_argument("--out", type=Path, default=Path(__file__).with_name("cpg-counts.csv"))
    args = parser.parse_args()

    rows = []
    for assembly in ("hg38", "hs1", "hg19"):
        path = getattr(args, assembly)
        if path:
            rows.extend(summarize(assembly, path))
    with args.out.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=FIELDS, lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)
    print(f"wrote {len(rows)} rows to {args.out}")


if __name__ == "__main__":
    main()
