"""Successor training protocol: matched32-byte prefix +32-byte continuation.
Preserves original200 structures; corrects conditioning coverage, not ancestry.
Historical models200.py/results remain unchanged for audit.
"""
import math
import torch
from torch.nn import functional as F
from models200 import Architecture as BaseArchitecture,SMIN,SMAX,cdf,sigma
class Architecture(BaseArchitecture):
    PREFIX=32;TARGETS=32
    def training_mask(self,t):
        assert t==32
        n=self.PREFIX+t-1;k=self.spec['generation_block_size'];i=torch.arange(t);group=(i//k)*k
        mask=torch.zeros((n+t,n+t),dtype=torch.bool);mask[:n,:n]=torch.ones((n,n),dtype=torch.bool).tril()
        mask[n:,:n]=torch.arange(n)[None,:]<(self.PREFIX+group[:,None])
        mask[n:,n:]=(i[None,:]//k)==(i[:,None]//k)
        return mask
    def denoise(self,clean,z,sig,block,training=True):
        if not training:return super().denoise(clean,z,sig,block,False)
        assert clean.shape[1]==63 and z.shape[1]==32
        b,t,d=z.shape;past=self.emb(clean);n=past.shape[1];sd=.5
        cin=torch.rsqrt(sig.square()+sd**2);cout=sig*sd*cin;cskip=sd**2/(sig.square()+sd**2)
        x=torch.cat([past,z*cin[:,None,None]],1);cond=x.new_zeros((b,n+t,1));cond[:,n:]=(.25*sig.log())[:,None,None]
        positions=torch.cat([torch.arange(n),torch.arange(self.PREFIX,self.PREFIX+t)])
        hidden=self.blocks[block](x,self.training_mask(t),positions,cond)[:,n:]
        return self.logits(hidden*cout[:,None,None]+z*cskip[:,None,None])
    def loss(self,batch,g):
        assert batch.shape[1]==64,'Exactly32prefix+32target bytes required'
        if not self.diff:
            logits=self(batch[:,:-1])[:,self.PREFIX-1:];target=batch[:,self.PREFIX:]
            return F.cross_entropy(logits.flatten(0,1),target.flatten()),{'active_blocks':len(self.blocks),'conditioning':'32clean-prefix,32continuation; causalteacherforcing'}
        interval=int(torch.randint(len(self.blocks),(1,),generator=g));lo,hi=self.bands[interval:interval+2]
        delta=math.log(hi)-math.log(lo);lo=max(SMIN,math.exp(math.log(lo)-.1*delta));hi=min(SMAX,math.exp(math.log(hi)+.1*delta))
        u=torch.rand(batch.shape[0],generator=g).tolist();sig=torch.tensor([sigma(cdf(lo)+(cdf(hi)-cdf(lo))*v) for v in u])
        target=F.normalize(self.emb(batch[:,self.PREFIX:]),dim=-1);z=target+sig[:,None,None]*torch.randn(target.shape,generator=g)
        block=len(self.blocks)-1-interval;logits=self.denoise(batch[:,:-1],z,sig,block,True)
        ce=F.cross_entropy(logits.flatten(0,1),batch[:,self.PREFIX:].flatten(),reduction='none').reshape(batch.shape[0],self.TARGETS)
        weight=(sig.square()+.25)/(sig*.5).square()
        return (ce*weight[:,None]).mean(),{'active_blocks':1,'block':block,'objective':'weighteddenoisingCE_not_ARNLL','conditioning':'matched32prefix plus permittedpreviousgroups'}
    def generate(self,prefix,new_tokens,seed=1,passes=4):
        if passes<2:raise ValueError('Diffusion passes must be >=2')
        if prefix.ndim!=2 or prefix.shape[1]<1 or not 1<=new_tokens<=128:raise ValueError('Invalid generation bounds')
        return super().generate(prefix,new_tokens,seed,passes)