"""GT motion-proxy ROI and GT-referenced dynamics. No actor/search sees future GT."""
import torch
import torch.nn.functional as F
from .rewards import CodecWorker
from .batch import TensorBatch

@torch.no_grad()
def motion_regions(truth,previous):
    """Shared across candidates: GT-derived, never shrinks after predicted disappearance."""
    prev=torch.cat([previous[None],truth[:-1]],0)
    motion=(truth-prev).abs().mean(1,keepdim=True)
    thresholds=motion.flatten(1).quantile(.90,dim=1).clamp_min(.02)
    masks=(motion>thresholds[:,None,None,None]).float()
    masks=F.max_pool2d(masks,11,stride=1,padding=5)
    fallback=masks.mean((1,2,3))<.01
    masks[fallback]=1
    crops=[];height,width=truth.shape[-2:]
    for mask in masks:
        y,x=torch.where(mask[0]>0)
        y0,y1=int(y.min()),int(y.max())+1;x0,x1=int(x.min()),int(x.max())+1
        cy=(y0+y1)//2;cx=(x0+x1)//2
        h=max(64,y1-y0);w=max(64,x1-x0)
        y0=max(0,min(cy-h//2,height-h));x0=max(0,min(cx-w//2,width-w))
        crops.append((y0,min(height,y0+h),x0,min(width,x0+w)))
    return masks,crops,fallback

class MotionReward:
    def __init__(self,codec,perceptual,scales=None,alpha=.5,beta=.25):
        self.worker=CodecWorker(codec,perceptual);self.perceptual=perceptual
        self.scales=scales or dict(global_scale=1.,roi_scale=1.,delta_scale=1.)
        self.alpha=alpha;self.beta=beta

    @torch.no_grad()
    def components(self,pred,pixels):
        truth=pixels[0,2:];previous=pixels[0,1]
        n,h=pred.shape[:2];real=truth[None].expand(n,-1,-1,-1,-1)
        mae=(pred-real).abs().mean((2,3,4))
        percept=self.worker.perceptual_loss(TensorBatch(dict(real=real,pred=pred))).batch['perceptual_loss']
        masks,crops,fallback=motion_regions(truth,previous)
        roi_mae=((pred-real).abs()*masks).sum((2,3,4))/(3*masks.sum((1,2,3)))[None]
        roi_percept=torch.empty((n,h),device=pred.device)
        for t,(y0,y1,x0,x1) in enumerate(crops):
            r=truth[t:t+1,:,y0:y1,x0:x1]
            for start in range(0,n,2):
                pp=pred[start:start+2,t,:,y0:y1,x0:x1]
                roi_percept[start:start+len(pp),t]=self.perceptual(pp*2-1,r.expand(len(pp),-1,-1,-1)*2-1).flatten()
        pred_prev=torch.cat([previous[None,None].expand(n,1,-1,-1,-1),pred[:,:-1]],1)
        true_prev=torch.cat([previous[None],truth[:-1]],0)
        delta=((pred-pred_prev)-(truth-true_prev)[None]).abs().mean((2,3,4))
        return dict(mae=mae,lpips=percept,global_loss=mae+percept,
                    roi_loss=roi_mae+roi_percept,roi_mae=roi_mae,roi_lpips=roi_percept,
                    delta_loss=delta,roi_fallback=fallback,roi_area=masks.mean((1,2,3)))

    def combine(self,c):
        s=self.scales
        return -(c['global_loss']+self.alpha*s['global_scale']/s['roi_scale']*c['roi_loss']+
                 self.beta*s['global_scale']/s['delta_scale']*c['delta_loss']).mean(-1)

    @torch.no_grad()
    def decode(self,frames,context):
        codes=torch.stack([f['response'] for f in frames],1)
        assert codes.min()>=0 and codes.max()<4375
        return self.worker.detokenize(TensorBatch(dict(tokens=codes,ctx_tokens=context.expand(len(codes),-1,-1))),None).batch['pixels'][:,1:]

    @torch.no_grad()
    def __call__(self,frames,context,actions,pixels):
        c=self.components(self.decode(frames,context),pixels)
        reward=self.combine(c)
        assert torch.isfinite(reward).all()
        metrics={'critic/recon_loss/mean':float(c['mae'].mean()),'critic/perceptual_loss/mean':float(c['lpips'].mean())}
        metrics.update({k:float(v.float().mean()) for k,v in c.items()})
        return reward,metrics