"""Replay final temporal CNN training, preserving the scored loss/augmentation/RNG logic."""
import argparse,json,random,time
from pathlib import Path
import numpy as np
import torch
from torch.utils.data import DataLoader
import config as C
from dataset import PASTISPatchDataset,train_patch_ids,val_patch_ids,crop_ij,pad_collate
from temporal_cnn import make_model
from training_feature_clusters import metrics_cm
from augmentation import aug_dihedral,aug_temporal_subsample,aug_frame_dropout,aug_spectral,lr_factor
OUT=Path(__file__).resolve().parents[1]/"runs"
CEREALS=[2,3,4,6,10,11,17,18]
def write_json(path,data):
    path.parent.mkdir(parents=True,exist_ok=True);path.write_text(json.dumps(data,indent=2))

def loader(ids,batch,shuffle=False):
    return DataLoader(PASTISPatchDataset(ids),batch_size=batch,shuffle=shuffle,num_workers=2,pin_memory=True,persistent_workers=True,collate_fn=pad_collate)


def target_for(y,classes):
    target=torch.full_like(y,-100)
    for i,c in enumerate(classes):target[y==c]=i
    return target


@torch.inference_mode()
def validate(model,dl,classes):
    model.eval();cm=np.zeros((20,20),np.int64)
    for (x,pos,_),y in dl:
        x,pos=x.cuda(),pos.cuda()
        for i,j in C.FT_CROP_IJ:
            xc,yc=crop_ij(x,y,i,j,4)
            with torch.autocast('cuda',dtype=torch.float16):pred=model(xc,pos).argmax(1)
            pred=np.array(classes)[pred.cpu().numpy().ravel()];gt=yc.numpy().ravel()
            valid=np.isin(gt,classes);cm+=np.bincount(20*gt[valid]+pred[valid],minlength=400).reshape(20,20)
    result=metrics_cm(cm)
    if len(classes)<18:result['selection_score']=float(np.array(result['per_class_iou'])[np.array(classes)-1].mean())
    else:result['selection_score']=result['miou']
    return result


def train(args):
    classes=CEREALS if args.kind=='cereal' else list(range(1,19))
    folder=OUT/f'32X32_{args.kind}_seed{args.seed}';folder.mkdir(parents=True,exist_ok=True)
    model=make_model(args.kind,len(classes)).cuda();optimizer=torch.optim.AdamW(model.parameters(),lr=1e-3,weight_decay=.01)
    torch.manual_seed(args.seed);random.seed(args.seed)
    scaler=torch.amp.GradScaler('cuda');td=loader(train_patch_ids(),2,True);vd=loader(val_patch_ids(),2)
    start=0;best=-1.
    if (folder/'latest.tar').exists():
        ck=torch.load(folder/'latest.tar',weights_only=False);model.load_state_dict(ck['state']);optimizer.load_state_dict(ck['optimizer']);scaler.load_state_dict(ck['scaler']);start=ck['epoch']+1;best=ck['best'];torch.set_rng_state(ck['rng']);torch.cuda.set_rng_state_all(ck['cuda_rng']);random.setstate(ck['python_rng'])
    for epoch in range(start,args.epochs):
        model.train();started=time.monotonic();total=0.;steps=0
        for (x,pos,_),y in td:
            x,pos,y=x.cuda(),pos.cuda(),y.cuda()
            for i,j in C.FT_CROP_IJ:
                xc,yc=crop_ij(x,y,i,j,4);xc,yc=aug_dihedral(xc,yc);xc,pc=aug_temporal_subsample(xc,pos);xc=aug_frame_dropout(xc);xc=aug_spectral(xc)
                target=target_for(yc,classes);optimizer.zero_grad(set_to_none=True)
                factor=lr_factor(epoch*76+steps,args.epochs*76,3*76,.01)
                for g in optimizer.param_groups:g['lr']=1e-3*factor
                with torch.autocast('cuda',dtype=torch.float16):logits=model(xc,pc)
                loss=torch.nn.functional.cross_entropy(logits.float(),target,ignore_index=-100,label_smoothing=.05) if (target!=-100).any() else logits.float().sum()*0
                if not torch.isfinite(loss):raise RuntimeError('Nonfinite specialist loss')
                scaler.scale(loss).backward();scaler.unscale_(optimizer);torch.nn.utils.clip_grad_norm_(model.parameters(),1.);scaler.step(optimizer);scaler.update();total+=loss.item();steps+=1
        assert steps==76
        metric=validate(model,vd,classes)
        if metric['selection_score']>best:
            best=metric['selection_score'];torch.save({'state':model.state_dict(),'classes':classes,'kind':args.kind,'seed':args.seed,'epoch':epoch,'validation':metric},folder/'best.tar')
        torch.save({'state':model.state_dict(),'optimizer':optimizer.state_dict(),'scaler':scaler.state_dict(),'epoch':epoch,'best':best,'rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all(),'python_rng':random.getstate()},folder/'latest.tar')
        write_json(folder/'progress.json',{'epoch_completed':epoch+1,'epochs':args.epochs,'best_selection_score':best,'seconds_per_epoch':time.monotonic()-started})
        print('epoch',epoch+1,'loss',total/steps,'selection',metric['selection_score'],'best',best,'seconds',time.monotonic()-started,flush=True)


if __name__=='__main__':
    ap=argparse.ArgumentParser();ap.add_argument('--kind',choices=['cnn_control','attention'],default='cnn_control');ap.add_argument('--seed',type=int,default=3407);ap.add_argument('--epochs',type=int,default=100);args=ap.parse_args()
    torch.set_num_threads(4);torch.manual_seed(args.seed);random.seed(args.seed);np.random.seed(args.seed);torch.backends.cuda.matmul.allow_tf32=True;torch.backends.cudnn.allow_tf32=True;train(args)