"""Star-allele calling (spec PGX-012 … PGX-017)."""

from __future__ import annotations

import numpy as np

from simurg.pgx.star import CALLED, MASKED, NOT_COVERED, GeneModel, Observations, call_gene, gs_error
from simurg.pipeline.pgx_star import observations


def _model() -> GeneModel:
    # variants: v0 defines *2, v1 defines *3, v2 is carried by *4 and *3 wildcard, v3 not on the array
    variants = [("1", 100, "A", "G"), ("1", 200, "C", "T"), ("1", 300, "G", "A"), ("1", 400, "T", "C")]
    carries = np.array([[0, 0, 0, 0],   # *1
                        [1, 0, 0, 0],   # *2
                        [0, 1, -1, 0],  # *3 (v2 wildcard)
                        [0, 0, 1, 0],   # *4
                        [0, 0, 0, 1]],  # *5 (only v3, not covered)
                       dtype=np.int8)  # fmt: skip
    return GeneModel("G", variants, [(f"v{i}", None) for i in range(4)], ["*1", "*2", "*3", "*4", "*5"], carries,
                     np.array([0.7, 0.1, 0.1, 0.05, 0.05]), deletion_allele="*0", reference="*1")  # fmt: skip


def _obs(dosage: list[int], state: list[int] | None = None) -> Observations:
    st = np.array(state or [CALLED, CALLED, CALLED, NOT_COVERED], dtype=np.int8)
    return Observations(st, np.array(dosage, dtype=np.int8), np.full(4, 0.2), np.full(4, np.nan))


def test_error_model() -> None:  # PGX-012 (calibrated, K-029)
    assert gs_error(np.array([0.16, 0.9, 0.1, np.nan])).tolist() == [1e-3, 1e-3, 0.05, 0.05]


def test_simple_diplotypes() -> None:  # PGX-013/014
    m = _model()
    sols, feas = call_gene(m, _obs([1, 0, 0, -1]))
    assert sols[0].genotype == "*1/*2" and sols[0].score > 0.9 and feas == [sols[0]]
    sols, _ = call_gene(m, _obs([0, 1, 1, -1]))
    assert sols[0].genotype == "*1/*3"  # *3 with its wildcard at v2 explains the v2 het too (vs *3/*4)
    sols, _ = call_gene(m, _obs([2, 0, 0, -1]))
    assert sols[0].genotype == "*2/*2"


def test_array_equivalence() -> None:  # PGX-017: *5 differs from *1 only at an uncovered variant
    m = _model()
    sols, _ = call_gene(m, _obs([0, 0, 0, -1]))
    top = sols[0]
    assert top.genotype == "*1/*1" and ("*1", "*5") in top.members and ("*5", "*5") in top.members
    assert abs(sum(s.raw_score for s in sols) - 1) < 1e-9


def test_copy_number_forms() -> None:  # PGX-013 structural candidates
    m = _model()
    obs = Observations(np.array([CALLED, CALLED, CALLED, NOT_COVERED], np.int8), np.array([1, 0, 0, -1], np.int8),
                       np.full(4, 0.5), np.array([2 / 3, 0.0, 0.0, np.nan]))  # fmt: skip
    sols, _ = call_gene(m, obs, cn=3)
    assert sols[0].genotype == "*2x2/*1" and sols[0].copy_number == 3
    obs.baf[0] = 1.0
    sols, _ = call_gene(m, obs, cn=1)
    assert sols[0].genotype == "*2/*0"


def test_masked_record() -> None:  # PGX-010
    m = _model()
    recs = {("1", 100): [{"id": "p1,p1b", "ref": "A", "alts": ["G"], "gt": "1/1", "gs": "0.3", "baf": "1"}],
            ("1", 200): [{"id": "p2", "ref": "C", "alts": ["T"], "gt": "0/1", "gs": "0.3", "baf": "0.5"}]}  # fmt: skip
    obs, ploidy = observations(m, recs, masked={"p1b"})
    assert ploidy == 2 and obs.state[0] == MASKED and obs.state[1] == CALLED and obs.dosage[1] == 1
    assert obs.state[2] == NOT_COVERED