"""Train the paper-budget DeepONet with hybrid exponential vRBA."""
import csv
import json
import math
import random
import sys
import time
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from model.vrba import VRBADeepONet, ExponentialVRBA
from common import arguments, config, path, accelerator, load_data, write_json, save_checkpoint, dataset_manifest


def run(cfg, resume=None, smoke=False):
    seed=cfg['training']['seed']
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    torch.set_num_threads(1)
    device=accelerator(cfg['training']['device'])
    torch.cuda.manual_seed_all(seed)
    torch.backends.cuda.matmul.allow_tf32=False
    torch.backends.cudnn.allow_tf32=False
    data=load_data(cfg)
    results=path(cfg['paths']['results']) / 'smoke' if smoke else path(cfg['paths']['results'])
    results.mkdir(parents=True,exist_ok=True)
    if not smoke and (results/'history.csv').exists() and resume is None:
        raise RuntimeError('Existing formal history; use --resume or a new results directory.')
    write_json(results/'data_manifest.json',dataset_manifest(data,cfg))
    write_json(results/'config_resolved.json',cfg)
    x=torch.as_tensor(data['x'],device=device)
    y=torch.as_tensor(data['y'],device=device)
    t=torch.as_tensor(data['time'],device=device)
    nt=cfg['data']['train']; bs=cfg['training']['batch_size']
    steps=nt//bs
    if nt%bs: raise ValueError('Training count must be divisible by batch size.')
    epochs=1 if smoke else cfg['training']['epochs']
    model=VRBADeepONet(**cfg['model']).to(device)
    nparam=sum(p.numel() for p in model.parameters())
    if nparam != 101100: raise ValueError(f'Paper parameter count mismatch: {nparam}')
    v=cfg['vrba']
    attention=ExponentialVRBA(nt,cfg['data']['output_points'],device,
        **{k:v[k] for k in ['eta','gamma','phi','c','initial_weight']})
    optimizer=torch.optim.Adam(model.parameters(),lr=cfg['training']['learning_rate'],
        weight_decay=cfg['training']['weight_decay'],betas=tuple(cfg['training']['adam_betas']),
        eps=cfg['training']['adam_epsilon'])
    scheduler=torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max=cfg['training']['epochs']*steps)
    best=float('inf'); best_epoch=0; start_epoch=0; step=0
    if resume:
        ck=torch.load(path(resume),map_location=device,weights_only=False)
        if ck['model_config']!=cfg['model']: raise ValueError('Resume architecture mismatch.')
        if ck['normalization']!=data['stats']: raise ValueError('Resume normalization mismatch.')
        model.load_state_dict(ck['model_state_dict']); optimizer.load_state_dict(ck['optimizer_state_dict'])
        scheduler.load_state_dict(ck['scheduler_state_dict']); attention.weights.copy_(ck['attention_weights'].to(device))
        start_epoch=ck['epoch']; step=ck['global_step']; best=ck['best_validation']; best_epoch=ck['best_epoch']
        torch.set_rng_state(ck['torch_rng'].cpu()); torch.cuda.set_rng_state_all([r.cpu() for r in ck['cuda_rng']])
        np.random.set_state(ck['numpy_rng']); random.setstate(ck['python_rng'])
    val=data['splits']['validation']
    fields=['epoch','global_step','train_weighted_loss','train_sampled_mse_normalized',
            'validation_mse_normalized','validation_relative_l2_normalized_batch_mean',
            'validation_relative_l2_physical_global','learning_rate','best_epoch','elapsed_seconds']
    mode='a' if resume else 'w'
    begin=time.perf_counter()
    environment=dict(python=sys.version,torch=torch.__version__,hip=torch.version.hip,
        device=torch.cuda.get_device_name(device),parameter_count=nparam,
        job_id=__import__('os').environ.get('SLURM_JOB_ID'),node=__import__('socket').gethostname())
    write_json(results/'environment.json',environment)
    print('RUN_START '+json.dumps(dict(environment=environment,epochs=epochs,batch_size=bs,steps_per_epoch=steps,smoke=smoke)),flush=True)
    checks={}
    with (results/'history.csv').open(mode,newline='',buffering=1) as f:
        writer=csv.DictWriter(f,fieldnames=fields)
        if not resume: writer.writeheader()
        for epoch in range(start_epoch,epochs):
            model.train(); losses=torch.zeros(2,device=device)
            for _ in range(steps):
                # Eq. 44 is the function marginal; do not apply a second exponential.
                with torch.no_grad():
                    indices=torch.multinomial(attention.probabilities(),bs,replacement=v['replacement'])
                optimizer.zero_grad(set_to_none=True)
                prediction=model(x[indices],t)
                if smoke and not checks:
                    assert prediction.shape==(bs,cfg['data']['output_points'])
                    checks['forward_pass']='passed'
                residual=prediction-y[indices]
                weights=attention.update(indices,residual,step)
                loss=(weights*residual).square().mean()
                loss.backward()
                if smoke and 'backward_pass' not in checks:
                    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())
                    checks['backward_pass']='passed'
                optimizer.step(); scheduler.step(); step+=1
                losses+=torch.stack((loss.detach(),residual.detach().square().mean()))
            model.eval()
            with torch.inference_mode():
                pv=model(x[val],t); rv=pv-y[val]
                val_mse=rv.square().mean().item()
                val_rel=torch.stack([torch.linalg.vector_norm(rv[i:i+bs])/torch.linalg.vector_norm(y[val][i:i+bs])
                    for i in range(0,cfg['data']['validation'],bs)]).mean().item()
                physical_true=y[val]*data['stats']['radius_std']+data['stats']['radius_mean']
                physical_rel=(torch.linalg.vector_norm(rv*data['stats']['radius_std'])/torch.linalg.vector_norm(physical_true)).item()
                train_loss,train_mse=(losses/steps).tolist()
            if not all(math.isfinite(vv) for vv in [train_loss,train_mse,val_mse,val_rel,physical_rel]):
                raise FloatingPointError(f'Nonfinite values at epoch {epoch+1}')
            improved=val_rel<best
            if improved: best=val_rel; best_epoch=epoch+1
            def checkpoint():
                return dict(format_version=1,model_state_dict=model.state_dict(),model_config=cfg['model'],
                    normalization=data['stats'],epoch=epoch+1,global_step=step,best_validation=best,best_epoch=best_epoch,
                    validation_selection='mean of normalized relative L2 over two batches of 50',
                    optimizer_state_dict=optimizer.state_dict(),scheduler_state_dict=scheduler.state_dict(),
                    attention_weights=attention.weights.detach().cpu(),torch_rng=torch.get_rng_state(),
                    cuda_rng=torch.cuda.get_rng_state_all(),numpy_rng=np.random.get_state(),python_rng=random.getstate(),
                    data_selection='original indices 0:1000; train 0:800, val 800:900, test 900:1000',
                    training_config=cfg['training'],vrba_config=cfg['vrba'])
            if not smoke and improved: save_checkpoint(path(cfg['paths']['best_checkpoint']),checkpoint())
            if not smoke and ((epoch+1)%cfg['training']['checkpoint_frequency']==0 or epoch+1==epochs):
                save_checkpoint(path(cfg['paths']['last_checkpoint']),checkpoint())
            row=dict(zip(fields,[epoch+1,step,train_loss,train_mse,val_mse,val_rel,physical_rel,
                scheduler.get_last_lr()[0],best_epoch,time.perf_counter()-begin]))
            writer.writerow(row)
            print('EPOCH '+json.dumps(row),flush=True)
            if smoke:
                domain = dict(relative_l2_normalized=val_rel,
                    residual_variance_absolute=rv.abs().var(unbiased=False).item(),
                    residual_linf=rv.abs().max().item())
                assert all(math.isfinite(value) for value in domain.values())
                write_json(results/'domain_metrics.json',domain)
                print('SMOKE_DOMAIN_METRICS '+json.dumps(domain),flush=True)
                checks.update(train_loop_dry_run='passed' ,validation_loop_dry_run='passed',
                    domain_metric_dry_run='passed: relative L2, residual variance and Linf; CL is not applicable to this experiment',
                    config_consistency='passed: Bx200 -> Bx100, 500x1 -> 500x100 -> Bx500; 101100 parameters')
                assert train_loss>0 and best<1e10
                write_json(results/'checks.json',dict(attempt=1,checks=checks))
    torch.cuda.synchronize()
    summary=dict(status='completed',completed_epochs=epochs,global_step=step,best_epoch=best_epoch,
        best_validation_relative_l2_normalized=best,early_stopping=False,smoke=smoke,
        duration_seconds=time.perf_counter()-begin)
    write_json(results/'training_summary.json',summary)
    print('TRAINING_COMPLETE '+json.dumps(summary),flush=True)


if __name__=='__main__':
    parser=arguments(__doc__)
    parser.add_argument('--resume',default=None,help='Full training-state checkpoint; restores optimizer, scheduler, vRBA and RNG.')
    parser.add_argument('--smoke',action='store_true',help='Separate one-epoch pipeline check; no formal checkpoints.')
    args=parser.parse_args()
    run(config(args.config),args.resume,args.smoke)