"""Autoregressive and self-speculative decoding on a LayeredDecoder.

One cycle of self-speculative decoding:

1. Drafting. The last committed token (not yet cached) and then each draft
   token climb the layers; at every head allowed by the policy we ask the
   policy whether to exit. An exit yields a draft token from that head; the
   K/V of the deeper drafting layers are filled from the exit hidden state
   (CALM-style copy) so later draft tokens can attend to it. A token with no
   exit, the draft budget, or a drafted EOS ends the draft.
2. Verification. Every row of the cycle has passed through layer c_min (the
   shallowest exit the policy can use) with exact inputs, so its hidden state
   there is exact. One pass runs layers [c_min, L) over all rows, overwriting
   the approximate K/V, and yields the full-model distribution at each row.
3. Acceptance. Standard speculative sampling (Leviathan et al. 2023): accept
   draft x with prob min(1, p(x)/q(x)), else resample from norm(max(p-q, 0));
   if all drafts are accepted, sample a bonus token from the last row. At T=0
   this is exact argmax matching. The output law (greedy: the output) equals
   the full model's.

Cache invariant between cycles: positions [0, len(seq)-1) are exact in every
layer; the last committed token is processed at the next cycle.
"""

from __future__ import annotations

import bisect
import time
from dataclasses import dataclass, field

import torch
import torch.nn.functional as F

from .heads import ExitHead, entropy, head_draft_distribution
from .layered import LayeredDecoder


# ---------------------------------------------------------------------- policies
class ExitPolicy:
    """Which heads may be tried, in which order, and when to exit."""

    head_order: list[int]          # indices into `heads`, shallow to deep
    max_draft: int

    def exit(self, k: int, logits: torch.Tensor, n_drafted: int) -> bool:
        raise NotImplementedError

    def decide(self, k: int, conf: float, n_drafted: int) -> bool:
        """Exit decision from the head confidence (graph engine)."""
        raise NotImplementedError


class StaticPolicy(ExitPolicy):
    """Always exit at head k; draft exactly `length` tokens (LayerSkip-like)."""

    def __init__(self, k: int, length: int):
        self.head_order = [k]
        self.max_draft = length

    def exit(self, k, logits, n_drafted):
        return True

    def decide(self, k, conf, n_drafted):
        return True


class ThresholdPolicy(ExitPolicy):
    """Exit at the first head (in `head_order`) whose score passes its threshold.

    score: "maxprob" (max over V_k ∪ ⊥, 0 if ⊥ wins), "neg_entropy", or
    "one_minus_reject". With a single head this is the Kangaroo rule; with all
    heads it is DSSD.
    """

    def __init__(self, heads: list[ExitHead], thresholds: dict[int, float],
                 max_draft: int = 32, score: str = "maxprob", cum_gamma: float = 0.0,
                 acc_maps: dict | None = None):
        self.heads = heads
        # Optional per-head calibrated acceptance a_hat_k(conf) (isotonic, piecewise linear
        # knots/values): thresholds and the cumulative product then apply to a_hat.
        self.acc_maps = {int(k): (list(v[0]), list(v[1])) for k, v in (acc_maps or {}).items()}
        self.thresholds = thresholds
        self.head_order = sorted(thresholds)
        self.max_draft = max_draft
        self.score = score
        # Cumulative stop (g): a draft token is only emitted while the product of the
        # exit confidences of the block stays >= cum_gamma (proxy for the probability
        # that the whole block survives verification, cf. SpecDec++).
        self.cum_gamma = cum_gamma
        self._cum = 1.0

    def start_block(self) -> None:
        self._cum = 1.0

    def _cum_ok(self, conf: float, k: int | None = None) -> bool:
        # cum_gamma: a scalar, or a per-head dict {k: gamma_k} (F7 option B: the exit test
        # at head k is cum * a_hat_k >= gamma_k)
        g = self.cum_gamma.get(k, 0.0) if isinstance(self.cum_gamma, dict) else self.cum_gamma
        if g <= 0.0:
            self._cum *= conf
            return True
        if self._cum * conf < g:
            return False
        self._cum *= conf
        return True

    def exit(self, k, logits, n_drafted):
        head = self.heads[k]
        if self.score == "maxprob":
            _, _, s = head_draft_distribution(logits, head.reject, 1.0)
        elif self.score == "neg_entropy":
            s = -entropy(logits)
            if head.reject and int(logits.argmax(-1)) == logits.shape[-1] - 1:
                return False
        elif self.score == "one_minus_reject":
            _, p_rej, s = head_draft_distribution(logits, head.reject, 1.0)
            s = 1.0 - p_rej
        else:
            raise ValueError(self.score)
        return float(s) >= self.thresholds[k] and self._cum_ok(float(s), k)

    def _ahat(self, k, conf):
        if k not in self.acc_maps:
            return conf
        xs, ys = self.acc_maps[k]
        if conf <= xs[0]:
            return ys[0]
        if conf >= xs[-1]:
            return ys[-1]
        j = bisect.bisect_right(xs, conf)
        x0, x1, y0, y1 = xs[j - 1], xs[j], ys[j - 1], ys[j]
        return y0 + (y1 - y0) * (conf - x0) / (x1 - x0 if x1 > x0 else 1.0)

    def decide(self, k, conf, n_drafted):
        if conf <= 0.0 and self.heads[k].reject:  # ⊥ won: never draft (T4: gamma_k > a_hat(⊥))
            return False
        a = self._ahat(k, conf)
        return a >= self.thresholds[k] and self._cum_ok(a, k)


class ScriptedPolicy(ExitPolicy):
    """Testing: draft token i exits at head schedule[i % len] (None = no exit)."""

    def __init__(self, n_heads: int, schedule: list, max_draft: int = 15):
        self.head_order = list(range(n_heads))
        self.schedule = schedule
        self.max_draft = max_draft

    def _want(self, n_drafted):
        return self.schedule[n_drafted % len(self.schedule)]

    def exit(self, k, logits, n_drafted):
        return self._want(n_drafted) == k

    def decide(self, k, conf, n_drafted):
        return self._want(n_drafted) == k


class AdaptiveThresholdPolicy(ThresholdPolicy):
    """Single head, threshold adapted online toward a target acceptance rate
    (Draft & Verify / AdaEDL style): after each cycle, raise the threshold if the
    cycle's acceptance fell below `target`, lower it otherwise."""

    def __init__(self, heads, k: int, gamma0: float, target: float, step: float = 0.01,
                 max_draft: int = 15):
        super().__init__(heads, {k: gamma0}, max_draft=max_draft, score="maxprob")
        self.k, self.target, self.step = k, target, step

    def update(self, n_drafted: int, n_accepted: int) -> None:
        if n_drafted == 0:
            return
        g = self.thresholds[self.k] + (self.step if n_accepted / n_drafted < self.target else -self.step)
        self.thresholds[self.k] = min(0.999, max(0.0, g))


# ---------------------------------------------------------------------- stats
@dataclass
class GenStats:
    new_tokens: int = 0
    cycles: int = 0                 # full-model passes after prefill
    drafted: int = 0
    accepted: int = 0
    empty_drafts: int = 0           # cycles whose draft had length 0 (fallback)
    rejected_cycles: int = 0        # cycles ended by a rejected draft
    ar_steps: int = 0               # online control: plain AR steps (lambda > 1, no probe)
    lam_start: float = float("nan") # online control: lambda at the start / end of the request
    lam_end: float = float("nan")
    gamma_start: float = float("nan")  # cost safety (NoSlowdown): gamma at the start / end
    gamma_end: float = float("nan")
    ar_step_median: float = float("nan")  # AR runs: median per-token step time
    lab_n: int = 0                  # online control: labels sent to the controller (and their sum)
    lab_s: float = 0.0
    n_jumps: int = -1               # online control: re-anchoring jumps of the ACI (cumulative)
    pld_drafted: int = 0            # prompt lookup: copied drafts / accepted copied drafts
    pld_accepted: int = 0
    cost_units: float = 0.0         # online control: sum of cycle costs (1 = one AR step)
    h_max: float = float("-inf")    # max over drafting cycles of cost - (1 + xi) * emitted
    excess_end: float = 0.0         # sum over the request of cost - (1 + xi) * emitted
    excess_max: float = float("-inf")  # running maximum of that sum within the request
    exits: dict = field(default_factory=dict)       # head index -> count
    sources: list = field(default_factory=list)     # per output token: copy/head/star/draft/model
    head_evals: int = 0
    layer_rows: int = 0
    kv_fill_rows: int = 0
    vocab_proj_rows: int = 0
    seconds: float = 0.0
    prefill_seconds: float = 0.0

    @property
    def tau(self) -> float:
        """Tokens committed per full-model pass (incl. correction/bonus)."""
        return self.new_tokens / max(self.cycles, 1)

    def as_dict(self) -> dict:
        d = dict(self.__dict__)
        d.pop("sources")  # per-token list: kept out of the bench jsonl
        d["tau"] = self.tau
        d["acceptance"] = self.accepted / max(self.drafted, 1)
        # drafts emitted on the output prefix = accepted + one rejected per rejecting cycle
        d["acceptance_prefix"] = self.accepted / max(self.accepted + self.rejected_cycles, 1)
        d["tokens_per_s"] = self.new_tokens / max(self.seconds, 1e-9)
        return d


def _sync(device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


def _sample(probs: torch.Tensor, gen: torch.Generator | None) -> int:
    return int(torch.multinomial(probs, 1, generator=gen))


def _target_token(logits: torch.Tensor, temperature: float, gen) -> int:
    if temperature == 0.0:
        return int(logits.argmax())
    return _sample(F.softmax(logits / temperature, dim=-1), gen)


# ---------------------------------------------------------------------- decoders
@torch.no_grad()
def generate_ar(dec: LayeredDecoder, prompt: list[int], max_new: int,
                temperature: float = 0.0, eos_ids: set[int] | None = None,
                gen: torch.Generator | None = None) -> tuple[list[int], GenStats]:
    eos_ids = eos_ids or set()
    st = GenStats()
    dec.reset_counters()
    ids = torch.tensor(prompt, device=dec.device)
    _sync(dec.device)
    t0 = time.perf_counter()
    logits = dec.prefill(ids)
    _sync(dec.device)
    st.prefill_seconds = time.perf_counter() - t0
    out = [_target_token(logits, temperature, gen)]
    t1 = time.perf_counter()
    while len(out) < max_new and out[-1] not in eos_ids:
        pos = len(prompt) + len(out) - 1
        logits = dec.full_step(torch.tensor([out[-1]], device=dec.device), pos)[0]
        out.append(_target_token(logits, temperature, gen))
        st.cycles += 1
    _sync(dec.device)
    st.seconds = time.perf_counter() - t1 + st.prefill_seconds
    st.new_tokens = len(out)
    st.layer_rows = dec.counters["layer_rows"]
    st.vocab_proj_rows = dec.counters["vocab_proj_rows"]
    return out, st


@torch.no_grad()
def generate_spec(dec: LayeredDecoder, heads: list[ExitHead], head_layers: list[int],
                  policy: ExitPolicy, prompt: list[int], max_new: int,
                  temperature: float = 0.0, eos_ids: set[int] | None = None,
                  gen: torch.Generator | None = None,
                  kv_fill: str = "lazy", trace: list | None = None) -> tuple[list[int], GenStats]:
    """Self-speculative decoding; head k reads the hidden state after head_layers[k] layers.

    kv_fill:
      "lazy"      regime A (reference): before a draft row climbs past the depth of
                  earlier rows of the block, those rows are caught up exactly (batched),
                  so every head score depends only on the exact prefix. Verification
                  then runs only layers [c_max, L).
      "copy"      regime B: CALM state copy of the K/V of skipped drafting layers;
                  verification runs [c_min, L).
      "recompute" regime B variant: run the skipped drafting layers for exited tokens.
    trace: if a list, receives (position, head, logits) for every head evaluation.
    """
    eos_ids = eos_ids or set()
    st = GenStats()
    dec.reset_counters()
    L = dec.spec.num_layers
    order = policy.head_order
    c_min = head_layers[order[0]]
    c_max = head_layers[order[-1]]
    ids = torch.tensor(prompt, device=dec.device)

    _sync(dec.device)
    t0 = time.perf_counter()
    logits = dec.prefill(ids)
    _sync(dec.device)
    st.prefill_seconds = time.perf_counter() - t0
    out = [_target_token(logits, temperature, gen)]

    t1 = time.perf_counter()
    while len(out) < max_new and out[-1] not in eos_ids:
        pos0 = len(prompt) + len(out) - 1
        inputs = [out[-1]]
        h_min: list[torch.Tensor] = []    # hidden after c_min layers, per row
        drafts: list[int] = []
        q_rows: list[torch.Tensor | None] = []
        budget = min(policy.max_draft, max_new - len(out))
        if hasattr(policy, "start_block"):
            policy.start_block()

        if kv_fill == "lazy":
            drafts, q_rows, H, verify_from = _draft_lazy(
                dec, heads, head_layers, policy, out[-1], pos0, budget, temperature, gen,
                eos_ids, st, trace, context=(prompt + out[:-1]) if trace is not None else None)
            inputs = [out[-1]] + drafts
            H = dec.run_layers(H, pos0, verify_from, L)
            target = dec.final_logits(H)
            st.cycles += 1
            st.drafted += len(drafts)
            if not drafts:
                st.empty_drafts += 1
            n_acc, committed = _accept(drafts, q_rows, target, temperature, gen)
            st.accepted += n_acc
            for t in committed:
                out.append(t)
                if t in eos_ids or len(out) >= max_new:
                    break
            continue

        # ---------------- drafting (regime B)
        while True:
            j = len(inputs) - 1
            if len(drafts) >= budget or (drafts and drafts[-1] in eos_ids):
                break  # last input is processed by the verification pass
            pos = pos0 + j
            h = dec.run_layers(dec.embed(torch.tensor([inputs[j]], device=dec.device)),
                               pos, 0, c_min)
            h_min.append(h)
            depth = c_min
            exited = None
            for k in order:
                if head_layers[k] > depth:
                    h = dec.run_layers(h, pos, depth, head_layers[k])
                    depth = head_layers[k]
                logits_k = heads[k](h)[0]
                st.head_evals += 1
                if trace is not None:
                    trace.append((pos, k, logits_k.detach().clone()))
                if policy.exit(k, logits_k, len(drafts)):
                    exited = k
                    break
            if exited is None:
                break
            head = heads[exited]
            q_k, _, _ = head_draft_distribution(logits_k, head.reject, temperature)
            if temperature == 0.0:
                tok_local = int(q_k.argmax())
                q_full = None
            else:
                tok_local = _sample(q_k, gen)
                q_full = head.scatter_full(q_k)
            tok = int(head.to_token_ids(torch.tensor(tok_local, device=dec.device)))
            if depth < c_max:
                if kv_fill == "copy":
                    dec.fill_kv_from_hidden(h, pos, depth, c_max)
                elif kv_fill == "recompute":
                    dec.run_layers(h, pos, depth, c_max)
                else:
                    raise ValueError(kv_fill)
            st.exits[exited] = st.exits.get(exited, 0) + 1
            drafts.append(tok)
            q_rows.append(q_full)
            inputs.append(tok)

        # ---------------- verification (one full-model pass)
        if len(h_min) < len(inputs):  # last draft token never entered the stack
            j = len(inputs) - 1
            h_min.append(dec.run_layers(
                dec.embed(torch.tensor([inputs[j]], device=dec.device)), pos0 + j, 0, c_min))
        H = torch.cat(h_min, dim=0)
        H = dec.run_layers(H, pos0, c_min, L)
        target = dec.final_logits(H)               # (len(inputs), V)
        st.cycles += 1
        st.drafted += len(drafts)
        if not drafts:
            st.empty_drafts += 1

        # ---------------- acceptance
        n_acc, committed = _accept(drafts, q_rows, target, temperature, gen)
        st.accepted += n_acc

        for t in committed:
            out.append(t)
            if t in eos_ids or len(out) >= max_new:
                break

    _sync(dec.device)
    st.seconds = time.perf_counter() - t1 + st.prefill_seconds
    st.new_tokens = len(out)
    st.layer_rows = dec.counters["layer_rows"]
    st.kv_fill_rows = dec.counters["kv_fill_rows"]
    st.vocab_proj_rows = dec.counters["vocab_proj_rows"]
    return out, st


def _accept(drafts, q_rows, target, temperature, gen):
    """Speculative acceptance (Leviathan et al.): returns (#accepted, committed tokens)."""
    committed: list[int] = []
    for i, d in enumerate(drafts):
        if temperature == 0.0:
            if int(target[i].argmax()) != d:
                committed.append(int(target[i].argmax()))
                return len(committed) - 1, committed
        else:
            p = F.softmax(target[i] / temperature, dim=-1)
            q = q_rows[i]
            ratio = (p[d] / q[d].clamp_min(1e-20)).clamp(max=1.0)
            u = torch.rand((), device=p.device, generator=gen)
            if not bool(u < ratio):
                resid = (p - q).clamp_min(0.0)
                z = resid.sum()
                committed.append(_sample(resid / z if z > 0 else p, gen))
                return len(committed) - 1, committed
        committed.append(d)
    n_acc = len(committed)
    committed.append(_target_token(target[len(drafts)], temperature, gen))
    return n_acc, committed


def _draft_lazy(dec, heads, head_layers, policy, first_token, pos0, budget, temperature, gen,
                eos_ids, st, trace, context=None):
    """Regime-A drafting. Returns (drafts, q_rows, hidden rows at c_max, c_max).

    With `trace`, records (position, head, logits, exact input context) per head evaluation.
    """
    order = policy.head_order
    c_min, c_max = head_layers[order[0]], head_layers[order[-1]]
    rows: list[dict] = []  # {"pos", "h", "depth"} for every input row of the block

    def advance(target: int) -> None:
        while True:
            lag = [r for r in rows if r["depth"] < target]
            if not lag:
                return
            u = min(r["depth"] for r in lag)
            v = min([r["depth"] for r in lag if r["depth"] > u] + [target])
            batch = [r for r in lag if r["depth"] == u]
            out = dec.run_layers(torch.cat([r["h"] for r in batch]), 0, u, v,
                                 positions=[r["pos"] for r in batch])
            for i, r in enumerate(batch):
                r["h"], r["depth"] = out[i:i + 1], v

    def new_row(tok: int, pos: int) -> dict:
        r = {"pos": pos, "h": dec.embed(torch.tensor([tok], device=dec.device)), "depth": 0}
        rows.append(r)
        return r

    drafts: list[int] = []
    q_rows: list = []
    tok = first_token
    while True:
        row = new_row(tok, pos0 + len(drafts))
        if len(drafts) >= budget or (drafts and drafts[-1] in eos_ids):
            break  # trailing row: only verified
        exited = None
        for k in order:
            advance(head_layers[k])
            logits_k = heads[k](row["h"])[0]
            st.head_evals += 1
            if trace is not None:
                ctx = context + [first_token] + drafts
                trace.append((row["pos"], k, logits_k.detach().clone(), ctx))
            if policy.exit(k, logits_k, len(drafts)):
                exited = k
                break
        if exited is None:
            break
        head = heads[exited]
        q_k, _, _ = head_draft_distribution(logits_k, head.reject, temperature)
        if temperature == 0.0:
            tok_local, q_full = int(q_k.argmax()), None
        else:
            tok_local = _sample(q_k, gen)
            q_full = head.scatter_full(q_k)
        tok = int(head.to_token_ids(torch.tensor(tok_local, device=dec.device)))
        st.exits[exited] = st.exits.get(exited, 0) + 1
        drafts.append(tok)
        q_rows.append(q_full)
        if hasattr(policy, "update_draft"):
            policy.update_draft()
    advance(c_max)
    return drafts, q_rows, torch.cat([r["h"] for r in rows]), c_max