SimurgArray/src/simurg/formats/bcf.py

"""BCF 2.2 writer and CSI index (spec OUT-040, OUT-038).

Records are built from typed values; the file is BGZF-compressed and indexed with a CSI index (min_shift 14,
depth 5, the pseudo-bin carries record counts). [DOC: VCF/BCF specification v4.3 §6 (BCF2), CSI index format —
layouts only]
"""

from __future__ import annotations

import struct
from collections.abc import Sequence

from simurg.formats.bgzf import PSEUDO_BIN, BgzfWriter, bgzf_compress, reg2bin

INT8, INT16, INT32, FLOAT, CHAR = 1, 2, 3, 5, 7
FLOAT_MISSING = struct.pack("<I", 0x7F800001)
INT8_VECTOR_END = -127
INT8_MISSING = -128


def _type_byte(n: int, t: int) -> bytes:
    if n < 15:
        return bytes([(n << 4) | t])
    return bytes([(15 << 4) | t]) + typed_ints([n])


def typed_ints(values: Sequence[int]) -> bytes:
    """Smallest integer type that holds every value (missing/vector-end sentinels excluded from the range)."""
    lo, hi = (min(values), max(values)) if values else (0, 0)
    if lo >= -120 and hi <= 127:
        return _type_byte(len(values), INT8) + struct.pack(f"<{len(values)}b", *values)
    if lo >= -32760 and hi <= 32767:
        return _type_byte(len(values), INT16) + struct.pack(f"<{len(values)}h", *values)
    return _type_byte(len(values), INT32) + struct.pack(f"<{len(values)}i", *values)


def typed_string(s: str) -> bytes:
    b = s.encode()
    return _type_byte(len(b), CHAR) + b


def float_bytes(values: Sequence[float | None]) -> bytes:
    out = bytearray()
    for v in values:
        out += FLOAT_MISSING if v is None or v != v else struct.pack("<f", v)
    return bytes(out)


def gt_bytes(genotypes: Sequence[tuple[int, ...] | None], ploidies: Sequence[int]) -> tuple[int, bytes]:
    """FORMAT/GT for all samples: width = max ploidy; unphased (a+1)<<1; missing 0; haploid padded with END."""
    width = max(1, max(ploidies))
    vals: list[int] = []
    for g, p in zip(genotypes, ploidies, strict=True):
        p = max(p, 1)
        cells = [0] * p if g is None else [(a + 1) << 1 for a in g]
        cells += [INT8_VECTOR_END] * (width - len(cells))
        vals += cells
    return width, bytes(vals_to_int8(vals))


def vals_to_int8(vals: Sequence[int]) -> bytes:
    return struct.pack(f"<{len(vals)}b", *vals)


def encode_record(
    contig_id: int,
    pos0: int,
    rlen: int,
    vid: str,
    alleles: Sequence[str],
    filter_id: int,
    n_sample: int,
    fmt: Sequence[tuple[int, bytes, bytes]],
    info: Sequence[tuple[int, bytes]] = (),
) -> bytes:
    """``fmt``: (format dictionary index, type-descriptor bytes, per-sample payload) in FORMAT order;
    ``info``: (dictionary index, typed value) pairs."""
    shared = bytearray()
    shared += struct.pack("<iii", contig_id, pos0, rlen)
    shared += FLOAT_MISSING  # QUAL
    shared += struct.pack("<I", (len(alleles) << 16) | len(info))  # n_allele << 16 | n_info
    shared += struct.pack("<I", (len(fmt) << 24) | n_sample)
    shared += typed_string(vid if vid else ".")
    for a in alleles:
        shared += typed_string(a)
    shared += typed_ints([filter_id])
    for key, value in info:
        shared += typed_ints([key]) + value
    indiv = bytearray()
    for key, desc, payload in fmt:
        indiv += typed_ints([key]) + desc + payload
    return struct.pack("<II", len(shared), len(indiv)) + bytes(shared) + bytes(indiv)


class CsiBuilder:
    """CSI index (min_shift 14, depth 5) with per-bin loffset and the pseudo-bin."""

    def __init__(self, n_ref: int) -> None:
        self.refs: list[dict[str, object]] = [{} for _ in range(n_ref)]

    def add(self, rid: int, beg: int, end: int, vstart: int, vend: int) -> None:
        r = self.refs[rid]
        if not r:
            r.update(bins={}, n=0, first=vstart, last=vend)
        bins: dict[int, list[list[int]]] = r["bins"]  # type: ignore[assignment]
        b = reg2bin(beg, end)
        chunks = bins.setdefault(b, [])
        if chunks and chunks[-1][1] == vstart:
            chunks[-1][1] = vend
        else:
            chunks.append([vstart, vend])
        r["n"] = int(r["n"]) + 1  # type: ignore[call-overload]
        r["last"] = vend

    def encode(self) -> bytes:
        out = bytearray(b"CSI\x01") + struct.pack("<iii", 14, 5, 0) + struct.pack("<i", len(self.refs))
        for r in self.refs:
            if not r:
                out += struct.pack("<i", 0)
                continue
            bins: dict[int, list[list[int]]] = r["bins"]  # type: ignore[assignment]
            out += struct.pack("<i", len(bins) + 1)
            for b in sorted(bins):
                ch = bins[b]
                out += struct.pack("<IQi", b, ch[0][0], len(ch))
                for s, e in ch:
                    out += struct.pack("<QQ", s, e)
            out += struct.pack("<IQi", PSEUDO_BIN, 0, 2) + struct.pack("<QQQQ", r["first"], r["last"], r["n"], 0)
        out += struct.pack("<Q", 0)
        return bgzf_compress(bytes(out))


class BcfWriter:
    """Stream BCF records into BGZF; header text must follow VCF 4.3 with PASS as the first FILTER line."""

    def __init__(self, fh: object, header_text: str) -> None:
        self.w = BgzfWriter(fh)  # type: ignore[arg-type]
        text = header_text.encode() + b"\0"
        self.w.write_record(b"BCF\x02\x02" + struct.pack("<I", len(text)) + text)

    def write(self, record: bytes) -> tuple[int, int]:
        return self.w.write_record(record)

    def close(self) -> None:
        self.w.close()


def header_dictionaries(header_text: str) -> tuple[dict[str, int], dict[str, int]]:
    """(string dictionary: FILTER/INFO/FORMAT IDs → index, contig dictionary: name → index) as htslib builds them."""
    strings: dict[str, int] = {"PASS": 0}
    contigs: dict[str, int] = {}
    for line in header_text.splitlines():
        for kind in ("FILTER", "INFO", "FORMAT"):
            if line.startswith(f"##{kind}=<ID="):
                key = line[len(kind) + 7 :].split(",", 1)[0].split(">", 1)[0]
                strings.setdefault(key, len(strings))
        if line.startswith("##contig=<ID="):
            name = line[13:].split(",", 1)[0].split(">", 1)[0]
            contigs.setdefault(name, len(contigs))
    return strings, contigs


def typed_floats(values: Sequence[float]) -> bytes:
    return _type_byte(len(values), FLOAT) + float_bytes(values)