"""Writers for small synthetic IDAT, BPM and EGT files (plan §16.4).

They produce structurally valid files for tests and synthetic data sets. Layouts mirror the readers in
:mod:`simurg.formats`; they are not meant to reproduce vendor files byte-for-byte.
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass, field

import numpy as np
import numpy.typing as npt

from simurg.formats import idat as idatfmt
from simurg.formats.binutil import BinaryWriter


@dataclass
class SynthIdat:
    address_ids: npt.NDArray[np.int32]
    mean: npt.NDArray[np.uint16]
    sd: npt.NDArray[np.uint16] | None = None
    n_beads: npt.NDArray[np.uint8] | None = None
    barcode: str = "200000000001"
    position: str = "R01C01"
    chip_type: str = "BeadChip 24x1x4"
    sample_id: str = ""
    plate: str = ""
    well: str = ""
    manifest: str = ""
    run_info: list[tuple[str, str, str, str, str]] = field(
        default_factory=lambda: [
            ("12/27/2023 5:52:21 PM", "Scan", "ScannerID=N0", "iScan Control Software", "4.2.1.729")
        ]
    )
    mid_block_itemsize: int = 4
    version: int = 3


def encode_idat(s: SynthIdat) -> bytes:
    n = len(s.address_ids)
    sd = s.sd if s.sd is not None else np.full(n, 100, np.uint16)
    nb = s.n_beads if s.n_beads is not None else np.full(n, 12, np.uint8)
    sections: list[tuple[int, bytes]] = []

    def sec(code: int, *parts: tuple[str, object]) -> None:
        w = BinaryWriter()
        for kind, val in parts:
            if kind == "i32":
                w.i32(int(val))  # type: ignore[call-overload]
            elif kind == "str":
                w.string(str(val))
            elif kind == "raw":
                w.write(val)  # type: ignore[arg-type]
            else:
                w.array(val, kind)  # type: ignore[arg-type]
        sections.append((code, w.getvalue()))

    mid_dtype = "<i4" if s.mid_block_itemsize == 4 else "u1"
    sec(idatfmt.N_SNPS, ("i32", n))
    sec(idatfmt.ILLUMINA_ID, ("<i4", s.address_ids))
    sec(idatfmt.SD, ("<u2", sd))
    sec(idatfmt.MEAN, ("<u2", s.mean))
    sec(idatfmt.NBEADS, ("u1", nb))
    sec(idatfmt.MID_BLOCK, ("i32", n), (mid_dtype, s.address_ids))
    run: list[tuple[str, object]] = [("i32", len(s.run_info))]
    run += [("str", part) for entry in s.run_info for part in entry]
    sec(idatfmt.RUN_INFO, *run)
    sec(idatfmt.RED_GREEN, ("raw", b"\x00\x00\x00\x00"))
    for code, val in (
        (idatfmt.MANIFEST, s.manifest),
        (idatfmt.BARCODE, s.barcode),
        (idatfmt.CHIP_TYPE, s.chip_type),
        (idatfmt.POSITION, s.position),
        (idatfmt.BEADSET, ""),
        (idatfmt.SAMPLE_ID, s.sample_id),
        (idatfmt.DESCRIPTION, ""),
        (idatfmt.PLATE, s.plate),
        (idatfmt.WELL, s.well),
    ):
        sec(code, ("str", val))
    sec(idatfmt.SAMPLE_COUNT, ("i32", 1))

    header = 4 + 8 + 4 + 10 * len(sections)
    out = BinaryWriter()
    out.write(b"IDAT")
    out.write(int(s.version).to_bytes(8, "little", signed=True))
    out.i32(len(sections))
    off = header
    for code, data in sections:
        out.u16(code)
        out.write(off.to_bytes(8, "little", signed=True))
        off += len(data)
    for _, data in sections:
        out.write(data)
    return out.getvalue()


@dataclass
class SynthLocus:
    name: str
    chrom: str = "1"
    pos: int = 1000
    snp: str = "[A/G]"
    address_a: int = 1
    address_b: int = 0
    assay_type: int = 0
    norm_id: int = 0
    ilmn_strand: str = "TOP"
    ref_strand: str = "+"
    source_strand: str = "TOP"
    intensity_only: int = 0
    exp_clusters: int = 3
    source_seq: str = ""
    probe_a: str = ""
    probe_b: str = ""
    ilmn_id: str | None = None


def encode_bpm(
    loci: Sequence[SynthLocus],
    *,
    manifest_name: str = "SYNTH-1v1_A1",
    entry_version: int = 8,
    bpm_version: int = 4,
    control_config: str = "",
    trailer: Sequence[str] = ("Illumina, Inc.",),
    entry_order: Sequence[int] | None = None,
) -> bytes:
    w = BinaryWriter()
    w.write(b"BPM")
    w.u8(1)
    w.i32(bpm_version)
    w.string(manifest_name)
    w.string(control_config)
    n = len(loci)
    w.i32(n)
    w.array(np.arange(1, n + 1), "<i4")
    for lc in loci:
        w.string(lc.name)
    w.array([lc.norm_id for lc in loci], "u1")
    order = list(entry_order) if entry_order is not None else list(range(n))
    for i in order:
        lc = loci[i]
        w.i32(entry_version)
        w.string(lc.ilmn_id or f"{lc.name}-0_T_F_1")
        w.string(lc.name)
        for _ in range(3):
            w.string("")
        w.i32(i + 1)
        w.string("")
        w.string(lc.ilmn_strand)
        w.string(lc.snp)
        w.string(lc.chrom)
        w.string("diploid")
        w.string("Homo sapiens")
        w.string(str(lc.pos))
        w.string("")  # TopGenomicSeq (v4 only)
        w.string("")  # customer strand
        w.i32(lc.address_a)
        w.i32(lc.address_b)
        w.string(lc.probe_a if entry_version == 4 else "")
        w.string(lc.probe_b if entry_version == 4 else "")
        w.string("38")
        w.string("synthetic")
        w.string("0")
        w.string(lc.source_strand)
        w.string(lc.source_seq if entry_version == 4 else "")
        if entry_version >= 6:
            w.u8(0)
            w.write(int(lc.exp_clusters).to_bytes(1, "little", signed=True))
            w.u8(lc.intensity_only)
            w.u8(lc.assay_type)
        if entry_version >= 7:
            for f in (0.25, 0.25, 0.25, 0.25):
                w.f32(f)
        if entry_version >= 8:
            w.string(lc.ref_strand)
    w.i32(len(trailer))
    for line in trailer:
        w.string(line)
    return w.getvalue()


def encode_egt(
    names: Sequence[str],
    *,
    theta_mean: npt.ArrayLike,
    r_mean: npt.ArrayLike,
    theta_dev: npt.ArrayLike | None = None,
    r_dev: npt.ArrayLike | None = None,
    n: npt.ArrayLike | None = None,
    gentrain: npt.ArrayLike | None = None,
    intensity_threshold: npt.ArrayLike | None = None,
    data_block_version: int = 9,
    manifest_name: str = "SYNTH-1v1_A1",
    corrupt_counts: bool = False,
) -> bytes:
    k = len(names)
    tm = np.asarray(theta_mean, np.float32).reshape(k, 3)
    rm = np.asarray(r_mean, np.float32).reshape(k, 3)
    td = np.asarray(theta_dev if theta_dev is not None else np.full((k, 3), 0.03), np.float32).reshape(k, 3)
    rd = np.asarray(r_dev if r_dev is not None else np.full((k, 3), 0.1), np.float32).reshape(k, 3)
    nn = np.asarray(n if n is not None else np.full((k, 3), 100), np.int32).reshape(k, 3)
    gt = np.asarray(gentrain if gentrain is not None else np.full(k, 0.8), np.float32)
    thr = np.asarray(intensity_threshold if intensity_threshold is not None else np.zeros(k), np.float32)
    w = BinaryWriter()
    w.i32(3)
    for s in ("< 6.1.3.0", "7.0.0", "7.0.0", "1.2.0", "1/1/2026 12:00:00 PM"):
        w.string(s)
    w.u8(1)
    w.string(manifest_name)
    w.i32(data_block_version)
    w.string("")
    w.i32(k)
    for i in range(k):
        w.array(nn[i], "<i4")
        w.array(rd[i], "<f4")
        w.array(rm[i], "<f4")
        w.array(td[i], "<f4")
        w.array(tm[i], "<f4")
        if data_block_version >= 7:
            w.f32(float(thr[i]))
            w.array(np.full(14, np.nan), "<f4")
        elif data_block_version >= 6:
            w.array(np.zeros(15), "<f4")
        elif data_block_version < 5:
            w.array(np.zeros(2), "<f4")
    for i in range(k):
        w.f32(0.5)
        w.f32(float(gt[i]))
        w.f32(float(gt[i]))
        w.u8(0)
    for _ in range(k):
        w.string("aa_ab_bb")
    for nm in names:
        w.string(nm)
    w.array(np.arange(1, k + 1), "<i4")
    counts = nn.copy()
    if corrupt_counts:
        counts[0, 0] += 1
    for i in range(k):
        w.array(counts[i], "<i4")
    if data_block_version == 9:
        w.array(np.zeros(k), "<f4")
    return w.getvalue()