import json
import os

import pytest
import torch

from mouh.jev.types import JevMessageType as T
from mouh.model import (
    ByteTokenizer, MouhConfig, MouhForCausalLM, estimate_params, preset_config, type_ids_for, PRESETS,
)
from mouh.model.layers import GatedSSM, _segsum, ssd_chunked, ssm_step
from mouh.model.moe import DeterministicMoE

torch.manual_seed(0)


@pytest.fixture(scope="module")
def tiny():
    torch.manual_seed(0)
    return MouhForCausalLM(preset_config("mouh-tiny")).eval()


def test_param_estimate_matches_real_count(tiny):
    real = sum(p.numel() for p in tiny.parameters())
    assert real == estimate_params(tiny.config).total


@pytest.mark.parametrize("name", sorted(PRESETS))
def test_param_estimate_matches_meta_instantiation(name):
    cfg = preset_config(name)
    if name != "mouh-tiny":
        with torch.device("meta"):
            m = MouhForCausalLM(cfg)
    else:
        m = MouhForCausalLM(cfg)
    assert sum(p.numel() for p in m.parameters()) == estimate_params(cfg).total


def test_layer_pattern_is_five_ssm_per_attention():
    cfg = preset_config("mouh-1.5b")
    pattern = ["A" if cfg.is_attention_layer(i) else "S" for i in range(cfg.num_hidden_layers)]
    assert "".join(pattern) == ("SSSSSA" * 4)
    with pytest.raises(ValueError):
        MouhConfig(num_hidden_layers=10)


def test_expert_count_and_topk():
    for name in PRESETS:
        cfg = preset_config(name)
        assert cfg.num_experts == 8
    assert preset_config("mouh-1.5b").num_experts_per_tok == 4


def test_spec_budget_arithmetic_documented():
    # top-4/8 : actifs >= 50 % des paramètres experts => total 1,5 Md => actifs >= ~750 M
    r = estimate_params(preset_config("mouh-1.5b"))
    assert 1.45e9 < r.total < 1.55e9 and r.active > 0.75e9
    r = estimate_params(preset_config("mouh-a400m"))
    assert abs(r.active - 400e6) / 400e6 < 0.01
    r = estimate_params(preset_config("mouh-1.5b-top1"))
    assert 1.45e9 < r.total < 1.55e9 and abs(r.active - 400e6) / 400e6 < 0.02


def test_ssd_chunked_matches_sequential_recurrence():
    torch.manual_seed(1)
    b, T, H, P, N = 2, 37, 3, 4, 5
    x = torch.randn(b, T, H, P)
    a = -torch.rand(b, T, H)
    B, C = torch.randn(b, T, N), torch.randn(b, T, N)
    y, fin = ssd_chunked(x, a, B, C, chunk=8)
    state = torch.zeros(b, H, P, N)
    ys = []
    for t in range(T):
        yt, state = ssm_step(state, x[:, t], a[:, t], B[:, t], C[:, t])
        ys.append(yt)
    ref = torch.stack(ys, 1)
    assert torch.allclose(y, ref, atol=1e-4, rtol=1e-4)
    assert torch.allclose(fin, state, atol=1e-4, rtol=1e-4)
    # état initial non nul
    init = torch.randn(b, H, P, N)
    y2, fin2 = ssd_chunked(x, a, B, C, chunk=8, init_state=init)
    state = init.clone(); ys = []
    for t in range(T):
        yt, state = ssm_step(state, x[:, t], a[:, t], B[:, t], C[:, t]); ys.append(yt)
    assert torch.allclose(y2, torch.stack(ys, 1), atol=1e-4, rtol=1e-4)
    assert torch.allclose(fin2, state, atol=1e-4, rtol=1e-4)


def test_ssm_block_streaming_equals_full(tiny):
    ssm = [b.mixer for b in tiny.blocks if isinstance(b.mixer, GatedSSM)][0].eval()
    x = torch.randn(2, 30, tiny.config.hidden_size)
    full, _ = ssm(x)
    cache, outs = None, []
    for t in range(30):
        o, cache = ssm(x[:, t:t + 1], cache)
        outs.append(o)
    assert torch.allclose(full, torch.cat(outs, 1), atol=1e-4, rtol=1e-4)
    # préremplissage puis flux
    o1, cache = ssm(x[:, :17])
    o2, _ = ssm(x[:, 17:], cache)
    assert torch.allclose(full, torch.cat([o1, o2], 1), atol=1e-4, rtol=1e-4)


def test_cached_decoding_matches_full_forward(tiny):
    ids = torch.randint(0, 256, (2, 25))
    with torch.no_grad():
        full = tiny(ids).logits
        o = tiny(ids[:, :10], use_cache=True)
        pieces = [o.logits]
        cache = o.cache
        for t in range(10, 25):
            o = tiny(ids[:, t:t + 1], cache=cache)
            cache = o.cache
            pieces.append(o.logits)
        stepped = torch.cat(pieces, 1)
    assert torch.allclose(full, stepped, atol=2e-3, rtol=2e-3), (full - stepped).abs().max()


def test_causality(tiny):
    a = torch.randint(0, 256, (1, 20))
    b = a.clone(); b[0, 15:] = torch.randint(0, 256, (5,))
    with torch.no_grad():
        la, lb = tiny(a).logits, tiny(b).logits
    assert torch.allclose(la[:, :15], lb[:, :15], atol=1e-4)


def test_forward_is_deterministic(tiny):
    ids = torch.randint(0, 256, (2, 16))
    with torch.no_grad():
        l1, l2 = tiny(ids, message_type_ids=type_ids_for(T.CODE_PROPOSAL, 2, ids.device)).logits, \
                 tiny(ids, message_type_ids=type_ids_for(T.CODE_PROPOSAL, 2, ids.device)).logits
    assert torch.equal(l1, l2)


def test_moe_router_topk_deterministic_and_type_prior():
    cfg = preset_config("mouh-tiny")
    moe = DeterministicMoE(cfg).eval()
    x = torch.randn(1, 12, cfg.hidden_size)
    moe(x, torch.tensor([0])); t1 = moe.last_trace
    moe(x, torch.tensor([0])); t2 = moe.last_trace
    assert torch.equal(t1.experts, t2.experts) and torch.equal(t1.weights, t2.weights)
    assert t1.experts.shape == (12, 4)
    assert torch.allclose(t1.weights.sum(-1), torch.ones(12), atol=1e-5)
    # un prior de type très fort impose les experts de la table JEV (CORRECTION_REQUEST -> n°5,1,7,3)
    strong = preset_config("mouh-tiny", router_type_prior_strength=1000.0)
    moe2 = DeterministicMoE(strong).eval()
    moe2(x, type_ids_for(T.CORRECTION_REQUEST, 1, x.device))
    chosen = {tuple(r) for r in moe2.last_trace.experts.tolist()}
    assert chosen == {(4, 0, 6, 2)}


def test_moe_ties_break_to_lowest_index():
    cfg = preset_config("mouh-tiny")
    moe = DeterministicMoE(cfg).eval()
    with torch.no_grad():
        moe.router.weight.zero_()
    moe(torch.randn(1, 3, cfg.hidden_size))
    assert moe.last_trace.experts.tolist() == [[0, 1, 2, 3]] * 3


def test_moe_capacity_limits_load_in_training():
    cfg = preset_config("mouh-tiny", capacity_factor=1.0)
    moe = DeterministicMoE(cfg).train()
    with torch.no_grad():
        moe.router.weight.zero_()  # tous les tokens visent les mêmes experts -> débordement garanti
    N = 32
    moe(torch.randn(1, N, cfg.hidden_size))
    cap = -(-int(1.0 * N * 4) // 8)
    assert moe.last_trace.load.max().item() <= cap
    assert (~moe.last_trace.kept).any()


def test_moe_balancing_bias_moves_against_overload():
    cfg = preset_config("mouh-tiny")
    moe = DeterministicMoE(cfg).train()
    with torch.no_grad():
        moe.router.weight.zero_()
    moe(torch.randn(1, 64, cfg.hidden_size))
    moe.step_routing_bias()
    assert moe.routing_bias[0] < 0 and moe.routing_bias[7] >= 0


def test_moe_expert_offload_roundtrip():
    cfg = preset_config("mouh-tiny")
    moe = DeterministicMoE(cfg).eval()
    x = torch.randn(1, 8, cfg.hidden_size)
    ref = moe(x)
    moe.offload_experts([0, 1], "cpu"); moe.load_experts([0, 1], "cpu")
    assert torch.equal(moe(x), ref)


def test_jev_decisions_are_typed(tiny):
    ids = torch.randint(0, 256, (1, 12))
    msg_choice, tool_choice = tiny.decide(ids, T.TASK_SUBMISSION)
    assert msg_choice.chosen in {t.value for t in T} and abs(sum(msg_choice.score.values) - 1) < 1e-9
    assert tool_choice.chosen in ("none", "PythonExecutionSkill", "PytestRunnerSkill", "FileSystemSkill", "Z3SolverSkill")


def test_generate_greedy_deterministic_and_cache_consistent(tiny):
    ids = torch.tensor([ByteTokenizer().encode("def f():")])
    a = tiny.generate_greedy(ids, max_new_tokens=12)
    b = tiny.generate_greedy(ids, max_new_tokens=12)
    assert torch.equal(a, b) and a.shape[1] >= ids.shape[1] + 1
    # le décodage avec cache équivaut à un décodage glouton sans cache
    seq = ids
    with torch.no_grad():
        for _ in range(12):
            nxt = tiny(seq).logits[:, -1].argmax(-1, keepdim=True)
            seq = torch.cat([seq, nxt], 1)
            if nxt.item() == tiny.config.eos_token_id:
                break
    assert torch.equal(a[:, : seq.shape[1]], seq)


def test_gradients_flow_to_all_parameter_groups():
    m = MouhForCausalLM(preset_config("mouh-tiny")).train()
    ids = torch.randint(0, 256, (2, 64))
    out = m(ids, labels=ids)
    out.loss.backward()
    missing = [n for n, p in m.named_parameters() if p.grad is None]
    assert not missing, missing[:5]
    assert all(torch.isfinite(p.grad).all() for p in m.parameters() if p.grad is not None)


def test_tiny_model_overfits_a_batch():
    torch.manual_seed(0)
    m = MouhForCausalLM(preset_config("mouh-tiny", router_aux_loss_coef=0.0)).train()
    opt = torch.optim.AdamW(m.parameters(), lr=3e-3)
    ids = torch.tensor([ByteTokenizer().encode("def add(a, b):\n    return a + b\n") * 1])
    first = None
    for step in range(60):
        out = m(ids, labels=ids)
        out.loss.backward(); opt.step(); opt.zero_grad()
        first = first or out.loss.item()
    assert out.loss.item() < first * 0.35


def test_safetensors_roundtrip_and_config_json(tmp_path, tiny):
    tiny.save_pretrained(tmp_path)
    assert (tmp_path / "model.safetensors").exists()
    cfg = json.loads((tmp_path / "config.json").read_text())
    assert cfg["model_type"] == "mouh" and cfg["architectures"] == ["MouhForCausalLM"]
    assert cfg["num_experts"] == 8 and cfg["num_experts_per_tok"] == 4 and cfg["attn_every"] == 6
    loaded = MouhForCausalLM.from_pretrained(tmp_path).eval()
    ids = torch.randint(0, 256, (1, 20))
    with torch.no_grad():
        assert torch.allclose(tiny(ids).logits, loaded(ids).logits, atol=1e-5)
    from transformers import AutoModelForCausalLM
    auto = AutoModelForCausalLM.from_pretrained(tmp_path)
    assert isinstance(auto, MouhForCausalLM)


def test_root_config_json_matches_default_preset():
    root = os.path.join(os.path.dirname(__file__), "..", "config.json")
    saved = json.load(open(root))
    cfg = preset_config()
    for k in ("hidden_size", "num_hidden_layers", "num_experts", "num_experts_per_tok",
              "expert_intermediate_size", "attn_every", "vocab_size", "model_type"):
        assert saved[k] == (getattr(cfg, k) if k != "model_type" else "mouh"), k


def test_routing_trace_is_exported_as_validated_bus_messages(tiny):
    from mouh.jev import JevBus
    from mouh.jev.types import new_id
    from mouh.model import routing_messages
    ids = torch.randint(0, 256, (1, 24))
    with torch.no_grad():
        tiny(ids, message_type_ids=type_ids_for(T.CODE_PROPOSAL, 1, ids.device))
    bus = JevBus()
    msgs = routing_messages(tiny, new_id(), T.CODE_PROPOSAL)
    assert len(msgs) == tiny.config.num_hidden_layers
    for m in msgs:
        out = bus.publish(m)  # validation stricte du schéma par le bus
        assert out[0].payload.experts and len(out[0].payload.experts) == 4