"""Phase 7: cached incremental decoding == full-prefix recomputation (tolerance-based), bypass bitwise."""
import torch

from prisme.cache import RequestBatchCache
from prisme.generation import step_logits, generate
from prisme.flops import action_cost_proxy


def test_incremental_equals_full_prefix_forced(model_nonzero, cfg, tokens):
    m = model_nonzero
    g = torch.Generator().manual_seed(3)
    acts = torch.randint(0, cfg.n_actions + 1, (3, 12, 4), generator=g)
    full = m(tokens, acts)["logits"]
    cache = RequestBatchCache(layers=m.backbone.new_cache())
    outs = []
    # prefill 5 tokens then decode one at a time
    lg, _, _ = step_logits(m, tokens[:, :5], cache, 4, forced_actions=acts[:, :5])
    outs.append(lg)
    for t in range(5, 12):
        lg, _, _ = step_logits(m, tokens[:, t:t + 1], cache, 4, forced_actions=acts[:, t:t + 1])
        outs.append(lg)
    inc = torch.cat(outs, 1)
    assert torch.allclose(inc, full, atol=1e-4, rtol=1e-4)
    assert cache.length == 12 and cache.bank.length == 12


def test_incremental_equals_full_prefix_policy(model_nonzero, cfg, tokens):
    m = model_nonzero
    with torch.no_grad():
        m.controller.heads.bias[0] = 0.5   # make the controller edit sometimes
        m.controller.heads.weight.normal_(std=0.5)
    cost = torch.tensor(action_cost_proxy(cfg, 12))
    full = m.forward_policy(tokens, 3, cost)
    cache = RequestBatchCache(layers=m.backbone.new_cache())
    outs, traces = [], []
    for t in range(12):
        lg, tr, _ = step_logits(m, tokens[:, t:t + 1], cache, 3, cost=cost)
        outs.append(lg); traces.append(tr)
    inc = torch.cat(outs, 1)
    assert torch.allclose(inc, full["logits"], atol=1e-4, rtol=1e-4)
    assert torch.equal(torch.cat(traces, 1), full["actions"])
    assert int(full["n_edits"].sum()) > 0


def test_budget_zero_bypass_bitwise_and_no_workspace(model_nonzero, cfg, tokens):
    m = model_nonzero
    ref = m.backbone.head(m.backbone(tokens))
    cache = RequestBatchCache(layers=m.backbone.new_cache())
    outs = []
    for t in range(12):
        lg, tr, ne = step_logits(m, tokens[:, t:t + 1], cache, 0)
        assert tr is None and int(ne.sum()) == 0
        outs.append(lg)
    inc = torch.cat(outs, 1)
    assert torch.allclose(inc, ref, atol=1e-4)
    # the explicit bypass branch of the full forward is bitwise identical to the backbone
    assert torch.equal(m(tokens, None)["logits"], ref)


def test_cache_isolation_between_requests(model_nonzero, cfg, tokens):
    m = model_nonzero
    c1 = RequestBatchCache(layers=m.backbone.new_cache())
    c2 = RequestBatchCache(layers=m.backbone.new_cache())
    step_logits(m, tokens[:1, :6], c1, 2, forced_actions=torch.zeros(1, 6, 2, dtype=torch.long))
    lg_a, _, _ = step_logits(m, tokens[1:2, :6], c2, 2, forced_actions=torch.zeros(1, 6, 2, dtype=torch.long))
    c3 = RequestBatchCache(layers=m.backbone.new_cache())
    lg_b, _, _ = step_logits(m, tokens[1:2, :6], c3, 2, forced_actions=torch.zeros(1, 6, 2, dtype=torch.long))
    assert torch.equal(lg_a, lg_b)
    assert c1.length == 6 and c2.length == 6
    c1.reset(); assert c1.length == 0 and c1.bank.length == 0 and c2.length == 6


def test_generate_runs_and_caps_budget(model_nonzero, cfg, tokens):
    m = model_nonzero
    cost = torch.tensor(action_cost_proxy(cfg, 12))
    out = generate(m, tokens[:, :4], 6, max_budget=8, cost=cost)
    assert out["tokens"].shape == (3, 10)
    assert out["n_edits"] is None or int(out["n_edits"].max()) <= cfg.r_max
    out0 = generate(m, tokens[:, :4], 6, max_budget=0)
    assert out0["tokens"].shape == (3, 10)