"""ArrayTrain cluster training (plan §12)."""

from __future__ import annotations

from pathlib import Path

import numpy as np

from simurg.formats.egt import read_egt, write_egt
from simurg.train.clusters import fit_block, impute_missing, learn_relations


def _block(rng: np.random.Generator, n_loci: int, s: int, freqs: np.ndarray) -> tuple[np.ndarray, ...]:
    centres = np.array([0.05, 0.5, 0.95])
    theta = np.empty((n_loci, s), np.float32)
    r = np.empty((n_loci, s), np.float32)
    truth = np.empty((n_loci, s), np.int64)
    for i in range(n_loci):
        p = freqs[i]
        g = rng.binomial(2, p, s)
        truth[i] = g
        theta[i] = np.clip(centres[g] + rng.normal(0, 0.025, s), 0, 1)
        r[i] = np.where(g == 1, 1.0, 0.85) + rng.normal(0, 0.08, s)
    return theta, r, truth


def test_fit_three_clusters() -> None:
    rng = np.random.default_rng(1)
    theta, r, _ = _block(rng, 40, 200, np.full(40, 0.4))
    fit = fit_block(theta, r, np.ones_like(theta, bool), np.zeros(40, bool))
    assert (fit.n > 0).all()
    assert np.allclose(fit.theta_mean.mean(axis=0), [0.05, 0.5, 0.95], atol=0.02)
    assert np.all(fit.n.sum(axis=1) >= 190)  # nearly every sample assigned
    assert fit.train_score.min() > 0.7 and fit.cluster_sep.min() > 0.9


def test_monomorphic_and_haploid() -> None:
    rng = np.random.default_rng(2)
    theta = np.clip(0.3 + rng.normal(0, 0.02, (2, 150)), 0, 1).astype(np.float32)  # one cluster at θ 0.3
    r = np.ones_like(theta)
    fit = fit_block(theta, r, np.ones_like(theta, bool), np.zeros(2, bool))
    assert (fit.n[:, 1] == 0).all() and (fit.n[:, 0] > 0).all()  # a lone cluster is homozygous (AA by θ < 0.5)
    theta, r, _ = _block(rng, 3, 100, np.full(3, 0.5))
    theta = np.where(np.abs(theta - 0.5) < 0.2, 0.95, theta).astype(np.float32)  # haploid: no hets
    fit = fit_block(theta, r, np.ones_like(theta, bool), np.ones(3, bool))
    assert (fit.n[:, 1] == 0).all()


def test_impute_missing_clusters() -> None:
    rng = np.random.default_rng(3)
    theta, r, _ = _block(rng, 60, 200, np.full(60, 0.45))
    rare, r2, _ = _block(rng, 5, 200, np.full(5, 0.0))  # only AA observed
    fit = fit_block(np.vstack([theta, rare]), np.vstack([r, r2]), np.ones((65, 200), bool), np.zeros(65, bool))
    rel = learn_relations(fit)
    impute_missing(fit, rel)
    assert (fit.n[60:, 1:] == 0).all()
    assert np.all(np.abs(fit.theta_mean[60:, 1] - 0.5) < 0.08) and np.all(fit.theta_mean[60:, 2] > 0.85)
    assert np.isfinite(fit.r_mean).all() and (fit.theta_dev > 0).all()


def test_egt_round_trip(tmp_path: Path) -> None:  # plan §12 step 5
    k = 4
    n = np.array([[10, 5, 0]] * k, np.int32)
    p = tmp_path / "t.egt"
    write_egt(p, [f"rs{i}" for i in range(k)], np.arange(k, dtype=np.int32), n,
              np.tile(np.float32([0.1, 0.5, 0.9]), (k, 1)), np.full((k, 3), 0.03, np.float32),
              np.ones((k, 3), np.float32), np.full((k, 3), 0.1, np.float32), np.full(k, 0.5, np.float32),
              np.full(k, 0.8, np.float32), "M.bpm")  # fmt: skip
    e = read_egt(p)
    assert e.data_block_version == 9 and e.extra_per_record == 15 and e.names == [f"rs{i}" for i in range(k)]
    assert (e.n == n).all() and np.allclose(e.gentrain_score, 0.8) and np.allclose(e.theta_mean[0], [0.1, 0.5, 0.9])


def test_compressed_and_two_cluster_labelling() -> None:
    """Compressed loci (BB at θ ≈ 0.45) and AB + BB two-cluster loci are labelled by HWE + position, not by fixed
    θ windows (the GSE165845 failure modes, docs/validation/V2_arraytrain.md)."""
    rng = np.random.default_rng(4)
    s = 200
    g = rng.binomial(2, 0.6, s)  # compressed: AA 0.05, AB 0.25, BB 0.45
    theta_c = np.clip(np.array([0.05, 0.25, 0.45])[g] + rng.normal(0, 0.02, s), 0, 1)
    g2 = rng.binomial(2, 0.9, s)
    g2[g2 == 0] = 1  # rare allele never homozygous: two clusters, AB (0.5) + BB (0.95)
    theta_2 = np.clip(np.array([0.05, 0.5, 0.95])[g2] + rng.normal(0, 0.025, s), 0, 1)
    theta = np.vstack([theta_c, theta_2]).astype(np.float32)
    fit = fit_block(theta, np.ones_like(theta), np.ones_like(theta, bool), np.zeros(2, bool))
    assert (fit.n[0] > 0).all() and np.allclose(fit.theta_mean[0], [0.05, 0.25, 0.45], atol=0.02)
    assert fit.n[1, 0] == 0 and fit.n[1, 1] > 0 and fit.n[1, 2] > 0
    assert abs(fit.n[1, 1] - (g2 == 1).sum()) <= 3