"""Master: masked IMU reconstruction and supervised local/global odometry."""
import math
import torch
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from .lab_bigru.model import WindowEncoder


class TCNBlock(nn.Module):
    def __init__(self, hidden, dilation):
        super().__init__()
        self.norm=nn.LayerNorm(hidden)
        self.conv=nn.Conv1d(hidden,hidden,3,padding=dilation,dilation=dilation)
        self.proj=nn.Linear(hidden,hidden)
        nn.init.zeros_(self.proj.weight);nn.init.zeros_(self.proj.bias)
    def forward(self,x,valid):
        z=self.norm(x)*valid[:,:,None]
        z=self.conv(z.transpose(1,2)).transpose(1,2)
        return (x+self.proj(torch.nn.functional.gelu(z)))*valid[:,:,None]


class Master(nn.Module):
    def __init__(self,mean,std,hidden=256):
        super().__init__()
        self.hidden=hidden
        self.encoder=WindowEncoder(mean,std)
        self.project=nn.Sequential(nn.Linear(384,hidden),nn.LayerNorm(hidden),nn.GELU(),nn.Dropout(.1))
        self.tcn=nn.ModuleList([TCNBlock(hidden,d) for d in (1,2,4)])
        self.temporal=nn.GRU(hidden,hidden,2,batch_first=True,bidirectional=True,dropout=.1)
        self.norm=nn.LayerNorm(hidden*2)
        layer=nn.TransformerEncoderLayer(hidden,4,hidden*4,.1,batch_first=True,norm_first=True,activation='gelu')
        self.global_net=nn.TransformerEncoder(layer,2,enable_nested_tensor=False)
        # Independent initial layer draws rather than cloned identical weights.
        for layer in self.global_net.layers:
            for module in layer.modules():
                if isinstance(module,nn.Linear):module.reset_parameters()
            layer.self_attn._reset_parameters()
        self.fuse=nn.Linear(hidden,hidden*2)
        nn.init.zeros_(self.fuse.weight);nn.init.zeros_(self.fuse.bias)
        self.direction=nn.Sequential(nn.Linear(hidden*2,hidden),nn.GELU(),nn.Dropout(.1),nn.Linear(hidden,3))
        self.speed=nn.Sequential(nn.Linear(hidden*2,hidden),nn.GELU(),nn.Dropout(.1),nn.Linear(hidden,1))
        self.residual=nn.Sequential(nn.Linear(hidden*2,hidden//2),nn.GELU(),nn.Linear(hidden//2,3))
        nn.init.zeros_(self.residual[-1].weight);nn.init.zeros_(self.residual[-1].bias)
        self.mask_embedding=nn.Parameter(torch.zeros(hidden))
        self.reconstruct=nn.Sequential(nn.Linear(hidden*2,hidden),nn.GELU(),nn.Linear(hidden,240))

    def forward(self,raw,lengths,starts,mask=None):
        b,t=raw.shape[:2];device=raw.device
        valid=torch.arange(t,device=device)[None]<lengths.to(device)[:,None]
        if mask is not None:
            raw=torch.where(mask[:,:,None,None],self.encoder.input_mean[None,None,:,None],raw)
        raw=torch.where(valid[:,:,None,None],raw, self.encoder.input_mean[None,None,:,None])
        flat=raw.reshape(-1,12,20)
        # Encoding chunks controls convolution workspace, without changing sample statistics.
        enc=torch.cat([self.encoder(z) for z in flat.split(256)]).reshape(b,t,384)
        x=self.project(enc)
        if mask is not None:x=x+mask[:,:,None]*self.mask_embedding
        x=x*valid[:,:,None]
        for block in self.tcn:x=block(x,valid)
        packed=pack_padded_sequence(x,lengths.cpu(),batch_first=True,enforce_sorted=False)
        out,_=self.temporal(packed)
        h,_=pad_packed_sequence(out,batch_first=True,total_length=t)
        h=self.norm(h)
        group=(torch.arange(t,device=device)[None]+starts.to(device)[:,None]%10)//10
        groups=(t+18)//10
        pooled=x.new_zeros(b,groups,self.hidden)
        pooled.scatter_add_(1,group[:,:,None].expand(-1,-1,self.hidden),x*valid[:,:,None])
        counts=x.new_zeros(b,groups)
        counts.scatter_add_(1,group,valid.to(x.dtype))
        pooled=pooled/counts.clamp_min(1)[:,:,None]
        pos=torch.arange(groups,device=device,dtype=torch.float32)[:,None]
        freq=torch.exp(torch.arange(0,self.hidden,2,device=device)*(-math.log(10000)/self.hidden))
        pe=torch.zeros(groups,self.hidden,device=device)
        pe[:,0::2]=torch.sin(pos*freq);pe[:,1::2]=torch.cos(pos*freq)
        g=self.global_net(pooled+pe.to(pooled.dtype),src_key_padding_mask=counts==0)
        g=g.gather(1,group[:,:,None].expand(-1,-1,self.hidden))
        h=h+self.fuse(g)
        if mask is not None:return self.reconstruct(h).reshape(b,t,12,20)
        d=self.direction(h).float()
        v=torch.nn.functional.normalize(d,dim=-1,eps=1e-6)*torch.nn.functional.softplus(self.speed(h).float())
        return (v+self.residual(h).float())*valid[:,:,None]


def block_mask(lengths,time,generator):
    result=torch.zeros(len(lengths),time,dtype=torch.bool)
    for b,n in enumerate(lengths.tolist()):
        if n<2:raise ValueError('need masked and visible tokens')
        count=max(1,min(n-1,round(n*.2)))
        while int(result[b].sum())<count:
            start=int(torch.randint(n,(1,),generator=generator))
            width=int(torch.randint(5,21,(1,),generator=generator))
            for i in range(start,min(n,start+width)):
                if int(result[b].sum())==count:break
                result[b,i]=True
    return result


def pretrain_loss(pred,raw,mask,model):
    target=(raw.float()-model.encoder.input_mean[None,None,:,None])/model.encoder.input_std[None,None,:,None]
    error=(pred.float()-target).square().mean((2,3))
    counts=mask.sum(1)
    if (counts==0).any():raise ValueError('empty reconstruction mask')
    return ((error*mask).sum(1)/counts).mean()


def supervised_loss(pred,batch,ramp):
    target=batch['target'];valid=torch.arange(pred.shape[1],device=pred.device)[None]<batch['lengths'].to(pred.device)[:,None]
    speed=torch.linalg.vector_norm(pred,dim=-1);truth=torch.linalg.vector_norm(target,dim=-1)
    err=torch.linalg.vector_norm(pred-target,dim=-1)
    mean=lambda x:((x*valid).sum(1)/valid.sum(1)).mean()
    dense=mean(err*torch.where(speed<truth,2.,1.))
    mag=mean((speed-truth).abs())
    ones=[];disp={k:[] for k in (2,5,10)}
    for b,(n,start) in enumerate(zip(batch['lengths'].tolist(),batch['starts'].tolist())):
        offset=(-start)%10;count=(n-offset)//10
        if count<1:continue
        delta=(pred[b,offset:offset+count*10]-target[b,offset:offset+count*10]).reshape(count,10,3).mean(1)
        ones.append(torch.linalg.vector_norm(delta,dim=-1).mean())
        rotation=batch['rotation'][b,offset+9:offset+count*10:10]
        world=torch.einsum('nij,nj->ni',rotation.float(),delta.float())
        cumulative=torch.cat([world.new_zeros(1,3),world.cumsum(0)])
        for k in disp:
            if count>=k:disp[k].append(torch.linalg.vector_norm(cumulative[k:]-cumulative[:-k],dim=-1).mean()/k)
    zero=pred.sum()*0
    one=torch.stack(ones).mean() if ones else zero
    ds=[torch.stack(v).mean() for v in disp.values() if v]
    displacement=torch.stack(ds).mean() if ds else zero
    parts=dict(dense=dense,speed=mag,one_second=one,displacement=displacement)
    return dense+ramp*(.25*mag+.25*one+.1*displacement),parts


def transfer(model,state):
    excluded=('direction.','speed.','residual.','reconstruct.','mask_embedding')
    own=model.state_dict(); keys={k for k in own if not k.startswith(excluded)}
    source={k for k in state if not k.startswith(excluded)}
    if keys!=source:raise ValueError('backbone key mismatch')
    for k in keys:
        if own[k].shape!=state[k].shape or own[k].dtype!=state[k].dtype:raise ValueError('backbone tensor mismatch '+k)
        own[k]=state[k]
    model.load_state_dict(own,strict=True)


def window_starts(n,window=640,stride=320):
    if n<1:raise ValueError('empty trajectory')
    if n<=window:return [0]
    return sorted(set(list(range(0,n-window+1,stride))+[n-window]))


def validate_parent(payload,seed,regime,inputs,steps,smoke):
    c=payload['contract']
    if (c['seed']!=seed or c['regime']!=regime or c['phase']!='pretrain'
            or c['historical_weights'] is not False or c['inputs']!=inputs
            or (not smoke and (payload['step']!=steps or c['smoke']))):
        raise ValueError('pretraining parent identity/data/step mismatch')


def load_f3_weights(model,state):
    """Strict mapping of every frozen F3 tensor; fresh residual paths stay zero."""
    mapping={'encoder.':'encoder.','input_projection.':'project.','temporal.':'temporal.',
             'output_norm.':'norm.','output.':'direction.','mag_head.':'speed.'}
    own=model.state_dict();translated={}
    for name,value in state.items():
        prefix=next((p for p in mapping if name.startswith(p)),None)
        if prefix is None:raise ValueError('unexpected F3 parameter '+name)
        target=mapping[prefix]+name[len(prefix):]
        if target not in own or value.shape!=own[target].shape or value.dtype!=own[target].dtype:
            raise ValueError('F3 tensor mismatch '+name)
        translated[target]=value
    required={k for k in own if k.startswith(tuple(mapping.values()))}
    if set(translated)!=required:raise ValueError('incomplete F3 checkpoint')
    own.update(translated);model.load_state_dict(own,strict=True)


def validate_continuation(payload,seed,regime,inputs):
    c=payload['contract']
    if (c['seed']!=seed or c['regime']!=regime or c['phase']!='finetune'
            or c['smoke'] or c['inputs']!=inputs or c['config']['name']!='f3_sd_warm'
            or payload['step']!=18000 or c['config']['finetune_steps']!=18000):
        raise ValueError('continuation requires matching completed Warm Master')