from __future__ import annotations
import torch
import torch.nn.functional as F

class ResidualTimestepCache:
    """Training-free residual/velocity cache for multi-step flow inference.

    Reuse is guarded by state cosine similarity, shape equality and a maximum timestep gap.
    It is intended only for >=4-step students; 1-2 step students should bypass it entirely.
    """
    def __init__(self,threshold=.985,max_reuse=2,max_timestep_gap=300.0):
        self.threshold=threshold;self.max_reuse=max_reuse;self.max_timestep_gap=max_timestep_gap;self.reset()
    def reset(self):self.prev_state=None;self.prev_output=None;self.prev_t=None;self.reuse_count=0
    @staticmethod
    def _feat(x):return x.feats if hasattr(x,'feats') else x
    @staticmethod
    def _t_scalar(t):
        if t is None:return None
        if torch.is_tensor(t):return float(t.detach().float().mean().cpu())
        return float(t)
    def similarity(self,x):
        if self.prev_state is None:return -1.
        a=self._feat(x).float();b=self._feat(self.prev_state).float()
        if a.shape!=b.shape:return -1.
        a=a.flatten(1);b=b.flatten(1);return float(F.cosine_similarity(a,b,dim=-1).mean())
    def get(self,x,t=None):
        sim=self.similarity(x);ts=self._t_scalar(t);time_ok=True
        if self.prev_t is not None and ts is not None:time_ok=abs(ts-self.prev_t)<=self.max_timestep_gap
        if time_ok and sim>=self.threshold and self.reuse_count<self.max_reuse:
            self.reuse_count+=1;return self.prev_output,{'hit':True,'similarity':sim,'t_gap':None if ts is None or self.prev_t is None else abs(ts-self.prev_t)}
        return None,{'hit':False,'similarity':sim,'t_gap':None if ts is None or self.prev_t is None else abs(ts-self.prev_t)}
    def put(self,x,out,t=None):
        self.prev_state=x;self.prev_output=out;self.prev_t=self._t_scalar(t);self.reuse_count=0

class CachedFlowModel:
    def __init__(self,model,threshold=.985,max_reuse=2,max_timestep_gap=300.0):
        self.model=model;self.cache=ResidualTimestepCache(threshold,max_reuse,max_timestep_gap);self.hits=0;self.calls=0
    def __call__(self,x,xs,t,cond,*args,**kwargs):
        self.calls+=1;cached,meta=self.cache.get(x,t)
        if cached is not None:self.hits+=1;return cached
        out=self.model(x,xs,t,cond,*args,**kwargs);self.cache.put(x,out,t);return out
    @property
    def hit_rate(self):return self.hits/max(1,self.calls)
    def reset(self):self.cache.reset();self.hits=0;self.calls=0