"""Matched one-epoch fold5 continuation pilots; never overwrite supplied weights."""
import argparse,time
from pathlib import Path
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import config as C
from dataset import PASTISPatchDataset,pretrain_patch_ids,pad_collate,crop_ij
from model import UTAE_ASPP,UTAEPrediction
from temporal_cnn import valid_from_index_positions
from notebook_helpers import write_json

class StructuredPrediction(UTAEPrediction):
    def forward(self,x,pos):
        b,t,c,h,w=x.shape
        # Four consecutive acquisitions x 4X4 pixels: all ten bands masked together.
        coarse=torch.rand(b,(t+3)//4,(h+3)//4,(w+3)//4,device=x.device)<self.mask_ratio
        hidden=coarse.repeat_interleave(4,1).repeat_interleave(4,2).repeat_interleave(4,3)[:,:t,:h,:w]
        visible=(~hidden)[:,:,None].expand_as(x)
        masked=x*visible+self.mask_token*(~visible)
        out,spatial=self.utae(masked,pos)
        recon=self.linear(out).permute(0,3,4,1,2)
        recon2=self.midlinear(spatial.permute(0,2,3,1)).permute(0,3,1,2).reshape(b,t,c,h,w)
        return recon,recon2,x,visible.float()

def main():
    ap=argparse.ArgumentParser();ap.add_argument('--pretrain_pth',required=True);ap.add_argument('--epochs',type=int,default=3);ap.add_argument('--policy',choices=['inherited','structured'],required=True);ap.add_argument('--limit_batches',type=int,default=0);ap.add_argument('--out_dir',required=True);args=ap.parse_args()
    torch.set_num_threads(4);torch.manual_seed(3407);torch.backends.cuda.matmul.allow_tf32=True;torch.backends.cudnn.allow_tf32=True
    out=Path(args.out_dir);out.mkdir(parents=True,exist_ok=True)
    if (out/'completed.json').exists():return
    enc=UTAE_ASPP(n_channels=C.N_CHANNELS,d_model=C.D_MODEL,n_heads=C.N_HEADS)
    enc.load_state_dict(torch.load(args.pretrain_pth,map_location='cpu',weights_only=False)['model_state_dict'],strict=True)
    model=(StructuredPrediction if args.policy=='structured' else UTAEPrediction)(enc,mask_ratio=.4).cuda()
    # Identical freshly initialized reconstruction heads in both arms; source encoder identical.
    dl=DataLoader(PASTISPatchDataset(pretrain_patch_ids(),load_target=False),batch_size=4,shuffle=True,num_workers=2,pin_memory=True,persistent_workers=True,collate_fn=pad_collate)
    opt=torch.optim.AdamW(model.parameters(),lr=1e-5,betas=(.9,.95),weight_decay=C.PRE_WD);scaler=torch.amp.GradScaler('cuda')
    start_epoch=0;start_batch=0;steps=0;total=0.;hidden_count=valid_count=0.
    # Materialize one fixed shuffled permutation for batch-boundary recovery.
    generator=torch.Generator().manual_seed(3407);order=torch.randperm(496,generator=generator).tolist()
    dl=DataLoader(torch.utils.data.Subset(dl.dataset,order),batch_size=4,shuffle=False,num_workers=2,pin_memory=True,persistent_workers=True,collate_fn=pad_collate)
    if (out/'resume.tar').exists():
        ck=torch.load(out/'resume.tar',weights_only=False);model.load_state_dict(ck['model']);opt.load_state_dict(ck['optimizer']);scaler.load_state_dict(ck['scaler']);start_epoch=ck.get('epoch',0);start_batch=ck['next_batch'];steps=ck['steps'];total=ck['total'];hidden_count=ck['hidden_count'];valid_count=ck['valid_count'];torch.set_rng_state(ck['rng']);torch.cuda.set_rng_state_all(ck['cuda_rng'])
    started=time.monotonic();model.train()
    for epoch in range(start_epoch,args.epochs):
        if epoch != start_epoch:
            start_batch=0;steps=0;total=0.;hidden_count=valid_count=0.
        for bi,((x,pos,_),_) in enumerate(dl):
            if bi<start_batch:continue
            if args.limit_batches and bi>=args.limit_batches:break
            x,pos=x.cuda(),pos.cuda();valid=valid_from_index_positions(pos)[:,:,None,None,None]
            for i in range(4):
                for j in range(4):
                    xc,_=crop_ij(x,None,i,j,4);opt.zero_grad(set_to_none=True)
                    with torch.autocast('cuda',dtype=torch.float16):r,r2,target,mask=model(xc,pos)
                    weight=(1-mask)*valid;denom=weight.sum().clamp_min(1.)
                    loss=(((r.float()-target)**2*weight).sum()+((r2.float()-target)**2*weight).sum())/denom
                    if not torch.isfinite(loss):raise RuntimeError('Nonfinite reconstruction loss')
                    scaler.scale(loss).backward();scaler.unscale_(opt);torch.nn.utils.clip_grad_norm_(model.parameters(),1.);scaler.step(opt);scaler.update()
                    steps+=1;total+=loss.item();hidden_count+=weight.sum().item();valid_count+=valid.expand_as(xc).sum().item()
            ck={'model':model.state_dict(),'optimizer':opt.state_dict(),'scaler':scaler.state_dict(),'epoch':epoch,'next_batch':bi+1,'steps':steps,'total':total,'hidden_count':hidden_count,'valid_count':valid_count,'rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all()}
            tmp=out/'resume.tmp';torch.save(ck,tmp);tmp.replace(out/'resume.tar')
            write_json(out/'progress.json',{'epoch':epoch+1,'epochs':args.epochs,'batches_completed':bi+1,'batches':len(dl),'steps':steps,'loss':total/steps,'mask_fraction':hidden_count/max(valid_count,1),'seconds':time.monotonic()-started})
            if bi%5==0:print(args.policy,'epoch',epoch+1,'batch',bi+1,'/',len(dl),'steps',steps,'loss',total/steps,'mask_fraction',hidden_count/max(valid_count,1),'seconds',time.monotonic()-started,flush=True)
        if not args.limit_batches:assert steps==1984
    torch.save({'epoch':args.epochs-1,'model_state_dict':model.utae.state_dict(),'policy':args.policy,'source':'latest.utae.tar','pilot_epochs':args.epochs},out/'encoder_32X32.tar')
    write_json(out/'completed.json',{'policy':args.policy,'epochs':args.epochs,'total_steps':args.epochs*1984,'steps_last_epoch':steps,'mean_reconstruction_loss':total/steps,'mask_fraction':hidden_count/max(valid_count,1),'seconds':time.monotonic()-started,'pilot':True,'labels_used':False})
if __name__=='__main__':main()