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

from __future__ import annotations

from pathlib import Path

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


def test_cn_only_genes() -> None:  # PGX-006
    from simurg.pipeline.pgx_star import cn_only_gene

    assert [cn_only_gene("s", "s", "GSTM1", c)[0]["genotype"] for c in (0, 1, 2, 3, 5, None)] == [
        "*0/*0", "*1/*0", "*1/*1", "*1x2/*1", "*1xN/*1", ""]  # fmt: skip


def test_variant_based_gene() -> None:  # PGX-021: per-variant notation, diplotype by lowest activity
    from simurg.pipeline.pgx_star import StarOptions, call_variant_gene

    variants = [("1", 100, "C", "T"), ("1", 200, "G", "A"), ("1", 300, "A", "G")]
    carries = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.int8)
    m = GeneModel("DPYD", variants, [("h0", "rs1"), ("h1", "rs2"), ("h2", "rs3")],
                  ["Reference", "c.A", "c.B", "c.C"], carries, np.array([0.9, 0.05, 0.03, 0.02]), None, "Reference",
                  "variant_based", {"Reference": 1.0, "c.A": 1.0, "c.B": 0.0, "c.C": 0.5},
                  {"Reference": "Normal function", "c.A": "Normal function", "c.B": "No function",
                   "c.C": "Decreased function"})  # fmt: skip
    recs = {("1", 100): [{"id": "p1", "ref": "C", "alts": ["T"], "gt": "0/1", "gs": "0.5", "baf": "0.5"}],
            ("1", 200): [{"id": "p2", "ref": "G", "alts": ["A"], "gt": "0/1", "gs": "0.5", "baf": "0.5"}],
            ("1", 300): [{"id": "p3", "ref": "A", "alts": ["G"], "gt": "0/1", "gs": "0.5", "baf": "0.5"}]}  # fmt: skip
    gj, rows, sup, _, _ = call_variant_gene("s", "s", m, recs, set(), StarOptions([], Path("x"), Path("y")))
    assert gj is not None and gj["genotype"] == "rs1 C/T; rs2 G/A; rs3 A/G"
    assert gj["feasible_solutions"] == ["c.B/c.C"]  # three carried: the two of lowest activity
    assert rows[0]["call_type"] == "variant" and len(sup) == 3
    recs = {("1", 100): [{"id": "p1", "ref": "C", "alts": ["T"], "gt": "0/0", "gs": "0.5", "baf": "0"}]}
    gj, *_ = call_variant_gene("s", "s", m, recs, set(), StarOptions([], Path("x"), Path("y")))
    assert gj is not None and gj["genotype"] == "Reference" and gj["feasible_solutions"] == ["Reference/Reference"]


def test_cyp2d6_structures() -> None:  # PGX-013 CN rules
    from simurg.pgx.star import Cyp2d6Structure, cyp2d6_structure, structure_candidates

    assert cyp2d6_structure(3, 2, 2) == Cyp2d6Structure(2, 0, 1, 0)  # *68 unit (NA12878 *3/*68+*4)
    assert cyp2d6_structure(3, 3, 2) == Cyp2d6Structure(2, 1, 0, 0)  # *36 unit
    assert cyp2d6_structure(1, 2, 2) == Cyp2d6Structure(1, 0, 0, 1)  # *13 hybrid
    assert cyp2d6_structure(2, 2, 2) == Cyp2d6Structure(2, 0, 0, 0)
    m = _model()
    m.haplotypes = ["*1", "*4", "*10", "*2", "*5"]
    names = {n for _, n in structure_candidates(m, Cyp2d6Structure(2, 0, 1, 0))}
    assert ("*1", "*68+*4") in names and ("*68+*1", "*2") in names
    names = {n for _, n in structure_candidates(m, Cyp2d6Structure(2, 1, 0, 0))}
    assert ("*36+*10", "*4") in names and all("*36+*10" in n for n in names)