"""Scientific correctness tests for the cache adapters (Section 7.2 of the brief). Runs a real small model."""
import os, json, pytest, numpy as np
import mlx.core as mx
from wp30.backends.mlx_backend import MLXModel, deep_copy_cache
from wp30.cache.sparse import select_positions, quantize_cache, to_sparse_all, cache_tensor_bytes, dense_bytes_formula
from wp30.cache.serialize import build_manifest, serialize_cache, deserialize_cache, IncompatibleCache
from wp30.cache.selection import select_cache, sink_recent_indices, random_indices

REPO = os.environ.get("WP30_TEST_MODEL", "mlx-community/Qwen2.5-1.5B-Instruct-bf16")
DOC = ("Project Varkolith-7 was proposed by Teodora Qvistrand on 14 March 2031 with a budget of 4,381,220 CAD. On 2 April 2031 the budget was "
       "revised to 3,907,415 CAD and the project was approved. The site is Lot 83-K in Chibougamau. The supplier identifier is ZR-77Q-0931. ") * 4
Q = "What is the supplier identifier? Answer with the identifier only."

def kl_div(a, b):
    pa = mx.softmax(a); pb = mx.softmax(b)
    return float((pa * (mx.log(pa + 1e-12) - mx.log(pb + 1e-12))).sum())

def max_dprob(a, b):
    return float(mx.abs(mx.softmax(a) - mx.softmax(b)).max())

@pytest.fixture(scope="module")
def m():
    return MLXModel(REPO)

@pytest.fixture(scope="module")
def m32(m):
    """float32 copy of the model for tolerance-free algorithmic checks."""
    mm = MLXModel(REPO)
    mm.model.set_dtype(mx.float32)
    mx.eval(mm.model.parameters())
    return mm

def split_ids(model):
    sys_msg = "Answer from the document."
    full = model.chat_prompt_tokens([{"role": "system", "content": sys_msg}, {"role": "user", "content": "Document:\n" + DOC + "\n\nQuestion: " + Q}])
    pre_txt = model.tokenizer.apply_chat_template([{"role": "system", "content": sys_msg}, {"role": "user", "content": "Document:\n" + DOC + "\n\n"}], add_generation_prompt=False, tokenize=False)
    pre = model.encode(pre_txt)
    n = 0
    while n < min(len(pre), len(full)) and pre[n] == full[n]:
        n += 1
    return full, full[:n], full[n:]

def test_prefix_equivalence_fp32(m32):
    full, pre, rest = split_ids(m32)
    c_full = m32.new_cache(); m32.prefill(full[:-1], c_full); ref = m32.next_logits(c_full, full[-1:])
    c_pre, _ = m32.prefill(pre); lg = m32.next_logits(c_pre, rest)
    # Registered tolerance (DECISIONS D-006): MLX kernels are not bitwise deterministic across prefill chunkings even in
    # fp32 (measured max|dlogit| ~0.1 between one-chunk and token-by-token prefill); the criterion is distributional:
    # KL(ref || cached) < 1e-5, max |dprob| < 1e-4, identical top-1.
    assert kl_div(ref, lg) < 1e-5 and max_dprob(ref, lg) < 1e-4
    assert int(mx.argmax(lg)) == int(mx.argmax(ref))

def test_prefix_equivalence_bf16_top1(m):
    full, pre, rest = split_ids(m)
    c_full = m.new_cache(); m.prefill(full[:-1], c_full); ref = m.next_logits(c_full, full[-1:])
    c_pre, _ = m.prefill(pre); lg = m.next_logits(c_pre, rest)
    assert int(mx.argmax(lg)) == int(mx.argmax(ref))
    assert kl_div(ref, lg) < 1e-3 and max_dprob(ref, lg) < 1e-2  # bf16 registered tolerance; top-1 must agree
    # greedy continuation agrees
    g1 = m.generate_text(full, max_tokens=8)["text"]
    c_pre, _ = m.prefill(pre); g2 = m.generate_from_cache(c_pre, rest, max_tokens=8)["text"]
    assert g1 == g2

def test_no_duplicated_boundary_token(m):
    full, pre, rest = split_ids(m)
    assert pre + rest == full and len(rest) >= 1
    bos = m.info.bos_token_id
    if bos is not None:
        assert full.count(bos) <= 1

def test_identity_roundtrip_and_bytes(m):
    full, pre, rest = split_ids(m)
    c_pre, _ = m.prefill(pre)
    lg = m.next_logits(deep_copy_cache(c_pre), rest)
    man = build_manifest(m.info, pre, c_pre, "doc", "B1", {})
    payload = serialize_cache(c_pre, man)
    c2, man2 = deserialize_cache(payload, m.info, man["prefix_token_hash"])
    lg2 = m.next_logits(c2, rest)
    assert float(mx.abs(lg2 - lg).max()) == 0.0
    tb = cache_tensor_bytes(c_pre)
    assert tb == dense_bytes_formula(m.info.num_layers, len(pre), m.info.num_kv_heads, m.info.head_dim, 2)
    assert len(payload) >= tb and len(payload) < tb * 1.01 + 65536

def test_loader_rejects_wrong_prefix_and_model(m):
    full, pre, rest = split_ids(m)
    c_pre, _ = m.prefill(pre)
    payload = serialize_cache(c_pre, build_manifest(m.info, pre, c_pre, "doc", "B1", {}))
    with pytest.raises(IncompatibleCache):
        deserialize_cache(payload, m.info, "deadbeef")
    import copy
    other = copy.copy(m.info); other.config_hash = "x"
    with pytest.raises(IncompatibleCache):
        deserialize_cache(payload, other)

def test_full_retention_control_exact(m):
    full, pre, rest = split_ids(m)
    c_pre, _ = m.prefill(pre); lg = m.next_logits(deep_copy_cache(c_pre), rest)
    T = len(pre)
    c_all = select_positions(c_pre, [np.arange(T)] * m.info.num_layers)
    assert float(mx.abs(m.next_logits(c_all, rest) - lg).max()) == 0.0
    # per-head 2-D indices identical across heads must also be exact
    idx2 = np.tile(np.arange(T), (m.info.num_kv_heads, 1))
    c_all2 = select_positions(c_pre, [idx2] * m.info.num_layers)
    assert float(mx.abs(m.next_logits(c_all2, rest) - lg).max()) == 0.0

def test_sparse_mask_path_matches_token_by_token_fp32(m32):
    full, pre, rest = split_ids(m32)
    c_pre, _ = m32.prefill(pre)
    T = len(pre)
    idx = sink_recent_indices(T, T // 2)
    a = select_positions(c_pre, [idx] * m32.info.num_layers); lg_a = m32.next_logits(a, rest)
    b = select_positions(c_pre, [idx] * m32.info.num_layers)
    for t in rest[:-1]:
        m32.next_logits(b, [t])
    lg_b = m32.next_logits(b, rest[-1:])
    assert kl_div(lg_a, lg_b) < 1e-5 and int(mx.argmax(lg_a)) == int(mx.argmax(lg_b))
    # logical offset preserved: new tokens positioned at T, T+1, ...
    assert a[0].offset == T + len(rest) and a[0].n_kept == len(idx) + len(rest)

def test_positions_preserved_vs_relabel(m32):
    """Retaining positions with the original offset must differ from relabelling them as 0..k-1 (which would be StreamingLLM-style)."""
    full, pre, rest = split_ids(m32)
    c_pre, _ = m32.prefill(pre); T = len(pre)
    idx = sink_recent_indices(T, T // 2)
    a = select_positions(c_pre, [idx] * m32.info.num_layers)
    b = select_positions(c_pre, [idx] * m32.info.num_layers, logical_offset=len(idx))
    assert float(mx.abs(m32.next_logits(a, rest) - m32.next_logits(b, rest)).max()) > 0.0

def test_packed_quantization_bytes_and_execution(m):
    full, pre, rest = split_ids(m)
    c_pre, _ = m.prefill(pre)
    dense = cache_tensor_bytes(c_pre)
    for bits, lo, hi in [(8, 1.6, 2.0), (4, 3.0, 4.0), (2, 5.0, 8.0)]:
        cq = quantize_cache(deep_copy_cache(c_pre), bits)
        b = cache_tensor_bytes(cq)
        assert lo <= dense / b <= hi, (bits, dense / b)
        k0 = cq[0].keys_and_values()[0]
        assert k0[0].dtype == mx.uint32  # packed words, not a wide dtype
        g = m.generate_from_cache(cq, rest, max_tokens=4)
        assert g["output_tokens"] >= 1

def test_no_shared_mutable_state(m):
    full, pre, rest = split_ids(m)
    c_pre, _ = m.prefill(pre)
    b1 = deep_copy_cache(c_pre); b2 = deep_copy_cache(c_pre)
    m.next_logits(b1, rest)
    assert b2[0].offset == c_pre[0].offset == len(pre)
    assert float(mx.abs(m.next_logits(b2, rest) - m.next_logits(deep_copy_cache(c_pre), rest)).max()) == 0.0

def test_attention_selection_runs_and_keeps_window(m):
    full, pre, rest = split_ids(m)
    c_pre, _ = m.prefill(pre); T = len(pre)
    out, info = select_cache(m, c_pre, "attention", 0.25, window_token_ids=rest, window_keep=8)
    assert out[0].n_kept == info["keep"] and out[0].offset == T
    rp = np.array(out[0].retained_positions)
    assert rp.ndim == 2 and all((rp[h] >= T - 8).sum() == 8 for h in range(rp.shape[0]))
    g = m.generate_from_cache(out, rest, max_tokens=6)
    assert g["output_tokens"] >= 1

def test_no_message_control_does_not_know_answer(m):
    ids = m.chat_prompt_tokens([{"role": "system", "content": "Answer."}, {"role": "user", "content": Q}])
    g = m.generate_text(ids, max_tokens=12)
    assert "ZR-77Q-0931" not in g["text"]