"""The four cache methods on the HY-WorldPlay DiT.  Same policies as
``cachelib/methods.py`` (see there for the rationale of every deviation from the
upstream implementations); only the block-level plumbing differs:

* a step is the 54-block vision loop of ``forward_vision``; the preamble
  (``img_in``, time/action embedding, RoPE tables) and ``final_layer`` always run;
* the indicator is block 0's modulated input, ``modulate(img_norm1(img), shift, scale)``;
* TaylorSeer forecasts each block's pre-gate attention and MLP features and
  re-applies the current step's gates;
* the selective forward keeps the other tokens' k/v in a per-layer bank filled at
  the last full step (HY-WorldPlay's denoising steps do not write the KV cache, so
  there are no slots to leave stale -- the bank plays that role).
"""

import math
from dataclasses import dataclass

import numpy as np
import torch

from .blocks import block_forward, block_gates, taylor_add_o0, taylor_add_o1


@dataclass
class StepCtx:
    model: object
    img: torch.Tensor          # [B, S, C] after img_in
    vec: torch.Tensor          # [B*S, C]
    freqs_cis: tuple           # (cos, sin) [S, D]
    viewmats: torch.Tensor     # [B, S, 4, 4]
    Ks: torch.Tensor           # [B, S, 3, 3]
    kv_cache: list
    grid: tuple                # (tt, th, tw)
    block_idx: int
    step_idx: int
    forced_full: bool
    last_cacheable_step: int


def run_blocks_full(ctx, bank=None, features=None):
    """The stock block loop when nothing has to be recorded, else the re-implementation."""
    model, img = ctx.model, ctx.img
    if bank is None and features is None:
        for index, block in enumerate(model.double_blocks):
            model.attn_param["layer-name"] = f"double_block_{index + 1}"
            img, _ = block(bi_inference=False, ar_txt_inference=False,
                           ar_vision_inference=True, img=img, vec=ctx.vec,
                           freqs_cis=ctx.freqs_cis, attn_param=model.attn_param,
                           block_idx=index, viewmats=ctx.viewmats, Ks=ctx.Ks,
                           kv_cache=ctx.kv_cache, cache_vision=False)
        return img
    if bank is not None:
        bank["viewmats"], bank["Ks"], bank["freqs_cis"] = ctx.viewmats, ctx.Ks, ctx.freqs_cis
    for index, block in enumerate(model.double_blocks):
        model.attn_param["layer-name"] = f"double_block_{index + 1}"
        img = block_forward(block, index, img, ctx.vec, ctx.freqs_cis, ctx.viewmats,
                            ctx.Ks, ctx.kv_cache, bank=bank, features=features)
    return img


def run_blocks_selective(ctx, sel, bank):
    """Recompute only the tokens ``sel``; returns their new hidden rows [B, len(sel), C]."""
    model = ctx.model
    assert ctx.img.shape[0] == 1, "selective forward assumes batch size 1"
    img = ctx.img[:, sel]
    vec = ctx.vec[sel]
    cos, sin = ctx.freqs_cis
    freqs = (cos[sel], sin[sel])
    viewmats, Ks = ctx.viewmats[:, sel], ctx.Ks[:, sel]
    for index, block in enumerate(model.double_blocks):
        model.attn_param["layer-name"] = f"double_block_{index + 1}"
        img = block_forward(block, index, img, vec, freqs, viewmats, Ks, ctx.kv_cache,
                            sel=sel, bank=bank)
    return img


def modulated_input(ctx):
    from hyvideo.models.transformers.modules.modulate_layers import modulate
    b0 = ctx.model.double_blocks[0]
    shift, scale = b0.img_mod(ctx.vec).chunk(6, dim=-1)[:2]
    return modulate(b0.img_norm1(ctx.img), shift=shift, scale=scale)


def rel_l1(cur, prev):
    diff = (cur - prev).abs().float().mean()
    base = prev.abs().float().mean() + 1e-8
    return (diff / base).item()


def rel_l1_per_token(cur, prev):
    diff = (cur - prev).abs().float().mean(dim=(0, -1))
    base = prev.abs().float().mean(dim=(0, -1)) + 1e-8
    return diff / base


class CacheMethod:
    name = "base"

    def __init__(self, coefficients=None, **kw):
        self.coefficients = list(coefficients) if coefficients else None
        self.extra = kw

    def rescale(self, value):
        if self.coefficients is None:
            return value
        return float(np.poly1d(self.coefficients)(value))

    def indicator(self, ctx):
        return modulated_input(ctx)

    def reset_video(self):
        pass

    def begin_chunk(self, block_idx):
        pass

    def forward(self, ctx):
        raise NotImplementedError

    def config(self):
        return {"name": self.name, "indicator": "modulated_input"}


class NoCache(CacheMethod):
    name = "none"

    def forward(self, ctx):
        return run_blocks_full(ctx), 1.0


class TeaCache(CacheMethod):
    name = "teacache"

    def __init__(self, thresh=0.0, **kw):
        super().__init__(**kw)
        self.thresh = float(thresh)
        self.reset_video()

    def reset_video(self):
        self.acc = 0.0
        self.prev_ind = None
        self.prev_residual = None

    def forward(self, ctx):
        ind = self.indicator(ctx)
        if ctx.forced_full or self.prev_ind is None or self.prev_residual is None:
            should_calc, self.acc = True, 0.0
        else:
            self.acc += self.rescale(rel_l1(ind, self.prev_ind))
            should_calc = self.acc >= self.thresh
            if should_calc:
                self.acc = 0.0
        self.prev_ind = ind
        if should_calc:
            ori = ctx.img
            img = run_blocks_full(ctx)
            self.prev_residual = img - ori
            return img, 1.0
        return ctx.img + self.prev_residual, 0.0

    def config(self):
        return dict(super().config(), thresh=self.thresh)


class FlowCache(CacheMethod):
    """Independent accumulator per group of latent frames inside the chunk
    (``group_size`` frames per group; 4 latent frames per chunk -> 4 groups)."""

    name = "flowcache"

    def __init__(self, thresh=0.0, group_size=1, **kw):
        super().__init__(**kw)
        self.thresh = float(thresh)
        self.group_size = int(group_size)
        self.reset_video()

    def reset_video(self):
        self.begin_chunk(-1)

    def begin_chunk(self, block_idx):
        self.acc, self.prev_ind, self.residual, self.bank = {}, {}, None, None

    def _groups(self, ctx):
        tt, th, tw = ctx.grid
        per = th * tw
        gs = min(self.group_size, tt)
        return [(s * per, min(s + gs, tt) * per) for s in range(0, tt, gs)]

    def forward(self, ctx):
        ind = self.indicator(ctx)
        groups = self._groups(ctx)
        if ctx.forced_full or self.residual is None:
            ori = ctx.img
            self.bank = {}
            img = run_blocks_full(ctx, bank=self.bank)
            self.residual = img - ori
            for g, (lo, hi) in enumerate(groups):
                self.acc[g] = 0.0
                self.prev_ind[g] = ind[:, lo:hi]
            return img, 1.0
        recompute = []
        for g, (lo, hi) in enumerate(groups):
            cur = ind[:, lo:hi]
            self.acc[g] = self.acc.get(g, 0.0) + self.rescale(rel_l1(cur, self.prev_ind[g]))
            if self.acc[g] >= self.thresh:
                self.acc[g] = 0.0
                recompute.append(g)
            self.prev_ind[g] = cur
        img = ctx.img + self.residual
        if not recompute:
            return img, 0.0
        sel = torch.cat([torch.arange(groups[g][0], groups[g][1], device=img.device)
                         for g in recompute])
        new = run_blocks_selective(ctx, sel, self.bank)
        img = img.index_copy(1, sel, new)
        self.residual = self.residual.index_copy(1, sel, new - ctx.img[:, sel])
        return img, sel.numel() / ctx.img.shape[1]

    def config(self):
        return dict(super().config(), thresh=self.thresh, group_size=self.group_size)


class TaylorSeer(CacheMethod):
    name = "taylorseer"

    def __init__(self, interval=1.0, max_order=1, **kw):
        super().__init__(**kw)
        self.interval = float(interval)
        self.max_order = int(max_order)
        self.reset_video()

    def reset_video(self):
        self._pattern_acc = 0.0
        self.begin_chunk(-1)

    def begin_chunk(self, block_idx):
        self.cache = {}
        self.activated = []
        lo, hi = int(math.floor(self.interval)), int(math.ceil(self.interval))
        if lo == hi:
            self.chunk_interval = max(1, lo)
        else:
            self._pattern_acc += self.interval - lo
            if self._pattern_acc >= 1.0 - 1e-9:
                self._pattern_acc -= 1.0
                self.chunk_interval = max(1, hi)
            else:
                self.chunk_interval = max(1, lo)
        self._since_full = 0

    def _should_calc(self, ctx):
        if ctx.forced_full or not self.activated:
            return True
        return (self._since_full + 1) >= self.chunk_interval

    def forward(self, ctx):
        if self._should_calc(ctx):
            img = self._record(ctx)
            self.activated.append(ctx.step_idx)
            self._since_full = 0
            return img, 1.0
        self._since_full += 1
        return self._forecast(ctx, ctx.step_idx - self.activated[-1]), 0.0

    def _record(self, ctx):
        dt = ctx.step_idx - self.activated[-1] if self.activated else 1
        # A derivative is only worth its memory if a later step of this chunk can
        # still be forecast from it.
        want_deriv = self.max_order >= 1 and ctx.step_idx < ctx.last_cacheable_step
        feats = {}
        img = run_blocks_full(ctx, features=feats)
        for i, f in feats.items():
            prev = self.cache.get(i)
            new = {}
            for key in ("attn", "mlp"):
                entry = {0: f[key]}
                if want_deriv and prev is not None and dt > 0 and 0 in prev[key]:
                    entry[1] = (f[key] - prev[key][0]) / dt
                new[key] = entry
            self.cache[i] = new
        return img

    def _forecast(self, ctx, distance):
        img = ctx.img
        d = float(distance)
        for i, block in enumerate(ctx.model.double_blocks):
            g1, g2 = block_gates(block, ctx.vec)
            a, m = self.cache[i]["attn"], self.cache[i]["mlp"]
            if 1 in a and 1 in m:
                img = taylor_add_o1(img, a[0], a[1], m[0], m[1], g1, g2, d)
            else:
                img = taylor_add_o0(img, a[0], m[0], g1, g2)
        return img

    def config(self):
        return dict(super().config(), interval=self.interval, max_order=self.max_order)


class MotionCache(CacheMethod):
    name = "motioncache"

    def __init__(self, thresh=0.0, weight_norm="mean", weight_floor=0.3,
                 min_update_ratio=0.0, **kw):
        super().__init__(**kw)
        self.thresh = float(thresh)
        self.weight_norm = weight_norm
        self.weight_floor = float(weight_floor)
        self.min_update_ratio = float(min_update_ratio)
        self.reset_video()

    def reset_video(self):
        self.prev_chunk_last_frame = None
        self.begin_chunk(-1)

    def begin_chunk(self, block_idx):
        self.acc = self.prev_ind = self.residual = self.weights = self.bank = None

    def _motion_weights(self, out, grid):
        tt, th, tw = grid
        spatial = th * tw
        o = out.view(out.shape[0], tt, spatial, out.shape[-1])
        diffs = []
        for fi in range(tt):
            cur = o[:, fi]
            if fi == 0:
                if self.prev_chunk_last_frame is None:
                    diffs.append(None)
                    continue
                prev = self.prev_chunk_last_frame
            else:
                prev = o[:, fi - 1]
            d = (cur - prev).abs().mean(dim=(0, 2))
            diffs.append(d / (prev.abs().mean(dim=(0, 2)) + 1e-8))
        if diffs[0] is None:
            diffs[0] = diffs[1].clone() if tt > 1 else torch.ones(spatial, device=out.device)
        fd = torch.stack(diffs).float()
        self.prev_chunk_last_frame = o[:, -1].detach().clone()
        if self.weight_norm == "max":
            w = fd / (fd.max(dim=1, keepdim=True)[0] + 1e-8)
        elif self.weight_norm == "max_rescale":
            lo, hi = fd.min(dim=1, keepdim=True)[0], fd.max(dim=1, keepdim=True)[0]
            w = self.weight_floor + (1 - self.weight_floor) * (fd - lo) / (hi - lo + 1e-8)
        else:
            w = fd / (fd.mean(dim=1, keepdim=True) + 1e-8)
        return w.reshape(-1)

    def forward(self, ctx):
        ind = self.indicator(ctx)
        L = ctx.img.shape[1]
        if ctx.forced_full or self.residual is None:
            ori = ctx.img
            self.bank = {}
            img = run_blocks_full(ctx, bank=self.bank)
            self.residual = img - ori
            self.prev_ind = ind
            self.acc = torch.zeros(L, device=img.device, dtype=torch.float32)
            self.weights = self._motion_weights(img, ctx.grid)
            return img, 1.0
        dist = rel_l1_per_token(ind, self.prev_ind)
        if self.coefficients is not None:
            c = self.coefficients
            dist = sum(c[i] * dist ** (len(c) - 1 - i) for i in range(len(c)))
        self.acc = self.acc + dist * self.weights
        self.prev_ind = ind
        need = self.acc >= self.thresh
        sel = torch.nonzero(need, as_tuple=False).flatten()
        if 0 < sel.numel() < self.min_update_ratio * L:
            sel = sel[:0]
        else:
            self.acc = torch.where(need, torch.zeros_like(self.acc), self.acc)
        img = ctx.img + self.residual
        if sel.numel() == 0:
            return img, 0.0
        new = run_blocks_selective(ctx, sel, self.bank)
        img = img.index_copy(1, sel, new)
        self.residual = self.residual.index_copy(1, sel, new - ctx.img[:, sel])
        return img, sel.numel() / L

    def config(self):
        return dict(super().config(), thresh=self.thresh, weight_norm=self.weight_norm,
                    min_update_ratio=self.min_update_ratio)



class DirectReuse(CacheMethod):
    """Naive cache baseline: every cacheable step reuses the last computed step's
    residual, no indicator and no threshold (`FRFF` / `FRRF` / `FRRR`)."""

    name = "reuse"

    def reset_video(self):
        self.prev_residual = None

    def forward(self, ctx):
        if ctx.forced_full or self.prev_residual is None:
            ori = ctx.img
            img = run_blocks_full(ctx)
            self.prev_residual = img - ori
            return img, 1.0
        return ctx.img + self.prev_residual, 0.0


class CalibrationProbe(CacheMethod):
    name = "calibrate"

    def __init__(self, **kw):
        super().__init__(**kw)
        self.samples = []
        self.prev_ind = self.prev_res = None

    def begin_chunk(self, block_idx):
        self.prev_ind = self.prev_res = None

    def forward(self, ctx):
        ind = self.indicator(ctx)
        ori = ctx.img
        img = run_blocks_full(ctx)
        res = img - ori
        if self.prev_ind is not None:
            self.samples.append({"x": rel_l1(ind, self.prev_ind), "y": rel_l1(res, self.prev_res),
                                 "block": ctx.block_idx, "step": ctx.step_idx})
        self.prev_ind, self.prev_res = ind, res
        return img, 1.0


class ExactCheck(CacheMethod):
    """Test-only: cacheable steps go through the re-implemented full loop (with bank
    recording) or through the selective path with *every* token selected.  Both
    must reproduce the stock forward bit for bit."""

    name = "exactcheck"

    def __init__(self, mode="reimpl", **kw):
        super().__init__(**kw)
        self.mode = mode
        self.bank = None

    def begin_chunk(self, block_idx):
        self.bank = None

    def forward(self, ctx):
        if self.mode == "reimpl" or ctx.forced_full or self.bank is None:
            self.bank = {}
            return run_blocks_full(ctx, bank=self.bank), 1.0
        sel = torch.arange(ctx.img.shape[1], device=ctx.img.device)
        return run_blocks_selective(ctx, sel, self.bank), 1.0


REGISTRY = {"none": NoCache, "reuse": DirectReuse, "calibrate": CalibrationProbe, "teacache": TeaCache,
            "flowcache": FlowCache, "taylorseer": TaylorSeer, "motioncache": MotionCache,
            "exactcheck": ExactCheck}


def build_method(name, **kw):
    if name not in REGISTRY:
        raise KeyError(f"unknown cache method {name!r}; have {sorted(REGISTRY)}")
    return REGISTRY[name](**{k: v for k, v in kw.items() if v is not None})