"""Phase P3: normalization numerics, fit and apply (plan §7.2–§7.4, Appendix A)."""

from __future__ import annotations

import numpy as np
import pytest
from typer.testing import CliRunner

from simurg.signal.normalization import numerics as nm
from simurg.signal.normalization.apply import normalize, polar
from simurg.signal.normalization.fit import (
    ROBUST_THRESHOLD,
    _closest_binned,
    _closest_slow,
    fit_transforms,
    normalize_single_bin_single_channel,
)

F32 = np.float32


def test_percentile_vendor_rule() -> None:
    v = np.arange(1, 11, dtype=np.float32)  # 1..10
    # n=10, p=50: i1=5, f=0 -> i1=4; x1=45, x2=55; y = 5 + (6-5)/10*(50-45) = 5.5
    assert nm.percentile(v, 50) == F32(5.5)
    assert nm.percentile(v, 0) == F32(1) and nm.percentile(v, 100) == F32(10)
    assert np.isnan(nm.percentile(np.zeros(0, np.float32), 50))


def test_trimmean_and_iqr() -> None:
    # n=20: vendor bounds P10 = 2.5 and P90 = 18.5 -> keeps 3..18 (the 1000 outlier is trimmed)
    v = np.array([*range(1, 20), 1000], dtype=np.float32)
    assert nm.trimmean(v, 20) == pytest.approx(10.5)
    # with only 5 values the vendor rule puts P90 at the maximum, so nothing is trimmed from the top
    assert nm.trimmean(np.array([1, 2, 3, 4, 100], np.float32), 20) == pytest.approx(22.0)
    assert nm.iqr(np.arange(0, 101, dtype=np.float32)) == pytest.approx(50.0, abs=1.0)


def test_seqsum_is_sequential() -> None:
    a = np.array([0.1, 0.2, 0.3, 1e16, -1e16, 0.4])
    assert nm.seqsum(a) == 0.4  # left-to-right; a compensated sum would give 1.0


@pytest.mark.parametrize("madsigma", [nm.madsigma_new, nm.madsigma_old])
def test_robustfit_recovers_lines_despite_outliers(madsigma) -> None:  # type: ignore[no-untyped-def]
    rng = np.random.default_rng(3)
    x = rng.uniform(0, 100, 300).astype(np.float32)
    y = (2.5 * x + 7 + rng.normal(0, 0.5, 300)).astype(np.float32)
    y[:20] += 500  # gross outliers
    b, m = nm.robustfit1(x, y, madsigma)
    assert float(m) == pytest.approx(2.5, abs=0.02) and float(b) == pytest.approx(7, abs=1.0)
    m0 = nm.robustfit0(x, (3.0 * x).astype(np.float32), madsigma)
    assert float(m0) == pytest.approx(3.0, rel=1e-6)


def test_linspace_and_unique() -> None:
    s = nm.linspace_f(5, F32(0), F32(1))
    assert s.tolist() == [0.0, 0.25, 0.5, 0.75, 1.0] and s.dtype == np.float32
    assert nm.unique_indices(np.array([3, -1, 3, 0, -1])).tolist() == [0, 3]


def test_nearest_neighbour_searches() -> None:
    rng = np.random.default_rng(0)
    xr = rng.uniform(0, 1000, 5000).astype(np.float32)
    yr = rng.uniform(0, 1000, 5000).astype(np.float32)
    sweep = nm.linspace_f(400, F32(0), F32(xr.max()))
    zeros = np.zeros(400, np.float32)
    exact = _closest_slow(xr, yr, sweep, zeros)
    binned = _closest_binned(xr, yr, sweep, zeros)
    assert (binned >= -1).all() and (binned < 5000).all()
    # the binned search is approximate but finds the exact neighbour for most sweep points
    found = binned >= 0
    assert (binned[found] == exact[found]).mean() > 0.9


def synth_bin(
    n: int, rng: np.random.Generator, *, offset=(300.0, 500.0), scale=(8000.0, 12000.0), theta=0.05, shear=0.03
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:  # type: ignore[type-arg]
    """Canonical AA/AB/BB clusters distorted by the inverse of the vendor transform.

    Like real data, each genotype cluster is a streak along its ray (intensity varies 0.3–1.6×), which is what
    the vendor's axis-sweep line fits rely on.
    """
    g = rng.choice(3, n, p=[0.4, 0.3, 0.3])
    inten = rng.uniform(0.3, 1.6, n)
    cx = inten * np.select([g == 0, g == 1, g == 2], [1.0, 0.5, 0.02]) + rng.normal(0, 0.01, n)
    cy = inten * np.select([g == 0, g == 1, g == 2], [0.02, 0.5, 1.0]) + rng.normal(0, 0.01, n)
    tx3, ty2 = cx * scale[0], cy * scale[1]
    tx2 = tx3 + shear * ty2
    c, s = np.cos(theta), np.sin(theta)
    tx = c * tx2 - s * ty2
    ty = s * tx2 + c * ty2
    raw_x = np.clip(np.round(tx + offset[0]), 0, 65535).astype(np.uint16)
    raw_y = np.clip(np.round(ty + offset[1]), 0, 65535).astype(np.uint16)
    return raw_x, raw_y, g


@pytest.mark.parametrize("version", ["1.2.0", "1.1.2"])
def test_fit_undoes_known_distortion(version) -> None:  # type: ignore[no-untyped-def]
    rng = np.random.default_rng(11)
    rx, ry, g = synth_bin(3000, rng)
    bins = np.zeros(3000, np.uint8)
    xf = fit_transforms(rx, ry, bins, version)
    assert len(xf) == 1
    # the AA streak itself is tilted by atan(0.02·12000/8000) ≈ 0.03 rad, which the fit also removes
    assert 0.05 < float(xf[0].theta) < 0.1
    nx, ny = normalize(rx, ry, np.zeros(3000, np.int32), xf)
    th, r = polar(nx, ny)
    assert np.nanmedian(th[g == 0]) < 0.08 and np.nanmedian(th[g == 2]) > 0.92
    assert 0.35 < np.nanmedian(th[g == 1]) < 0.65
    assert 0.8 < np.nanmedian(r[g != 1]) < 1.3


def test_bins_ordered_and_small_bins() -> None:
    rng = np.random.default_rng(5)
    rx1, ry1, _ = synth_bin(1000, rng)
    rx2, ry2, _ = synth_bin(150, rng)  # < 192 loci: offset/rotation/shear skipped
    rx = np.concatenate([rx2, rx1])
    ry = np.concatenate([ry2, ry1])
    bins = np.concatenate([np.full(150, 7, np.uint8), np.full(1000, 3, np.uint8)])
    xf = fit_transforms(rx, ry, bins, "1.2.0")
    assert len(xf) == 2  # ascending bin id: 3 then 7
    assert xf[0].theta != 0 and xf[1].theta == 0 and xf[1].offset_x == 0 and xf[1].shear == 0
    with pytest.raises(ValueError):
        fit_transforms(rx[:5], ry[:5], np.zeros(5, np.uint8))
    with pytest.raises(ValueError):
        fit_transforms(rx, ry, bins, "9.9")


def test_single_channel_model_for_small_infinium_i_bins() -> None:
    rng = np.random.default_rng(8)
    n = 120
    a = rng.uniform(2000, 9000, n)
    b = rng.uniform(200, 600, n)
    x = np.where(rng.random(n) < 0.5, a, b).astype(np.float32)
    y = np.where(x > 1000, b, a).astype(np.float32)
    xf = normalize_single_bin_single_channel(x, y)
    assert xf.offset_x == xf.offset_y and xf.scale_x == xf.scale_y and xf.theta == 0 and xf.shear == 0
    assert xf.rr12 == 1 and np.isfinite(xf.scale_x)
    bins = np.full(n, 105, np.uint8)  # Infinium I bin (id >= 100), < 192 loci
    rx, ry = x.astype(np.uint16), y.astype(np.uint16)
    via_fit = fit_transforms(rx, ry, bins, "1.2.0")[0]
    assert via_fit.scale_x == via_fit.scale_y  # 1.2.0 → 2-parameter model
    assert fit_transforms(rx, ry, bins, "1.1.2")[0].scale_x != via_fit.scale_x or n >= ROBUST_THRESHOLD


def test_polar_edge_cases() -> None:
    nx = np.array([0, 1, 0, 1], np.float32)
    ny = np.array([0, 0, 1, 1], np.float32)
    th, r = polar(nx, ny)
    assert np.isnan(th[0]) and np.isnan(r[0])
    assert th[1:].tolist() == [0.0, 1.0, 0.5] and r[1:].tolist() == [1.0, 1.0, 2.0]


def test_normalize_cli_on_ingest_output(tmp_path) -> None:  # type: ignore[no-untyped-def]
    import pyarrow.parquet as pq

    from simurg.cli.main import app
    from simurg.synth.files import SynthIdat, SynthLocus, encode_bpm, encode_idat

    rng = np.random.default_rng(2)
    n = 600
    rx, ry, _ = synth_bin(n, rng)
    addr = np.arange(1, n + 1, dtype=np.int32)
    for ch, mean in (("Red", rx), ("Grn", ry)):
        s = SynthIdat(
            address_ids=addr, mean=mean, n_beads=np.full(n, 12, np.uint8), barcode="200000000001", position="R01C01"
        )
        (tmp_path / f"200000000001_R01C01_{ch}.idat").write_bytes(encode_idat(s))
    (tmp_path / "m.bpm").write_bytes(encode_bpm([SynthLocus(f"rs{i}", address_a=int(a)) for i, a in enumerate(addr)]))
    out = tmp_path / "out"
    runner = CliRunner()
    assert (
        runner.invoke(
            app,
            [
                "ingest",
                "--quiet",
                "--bpm",
                str(tmp_path / "m.bpm"),
                "--idats",
                str(tmp_path),
                "--out",
                str(out),
                "--threads",
                "1",
            ],
        ).exit_code
        == 0
    )
    r = runner.invoke(app, ["normalize", "--quiet", "--in", str(out), "--threads", "1"])
    assert r.exit_code == 0, r.output
    t = pq.read_table(next(out.rglob("*.normalized.parquet")))
    assert t.num_rows == n and set(t.column_names) == {"norm_x", "norm_y", "theta", "r"}
    tr = pq.read_table(next(out.rglob("*.transforms.parquet")))
    assert tr.num_rows == 1 and tr.column("n_loci").to_pylist() == [n]
    r = runner.invoke(app, ["normalize", "--quiet", "--in", str(out)])
    assert r.exit_code == 2 and "E_OUTPUT_EXISTS" in r.output