"""Complete joint epochs, exact-resume checkpoints, W&B and verified Hub backups."""
from pathlib import Path
import sys
sys.path.insert(0,str(Path(__file__).resolve().parent))
from multitask_model import *
import argparse,ast,copy,math,random
os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG',':4096:8')


def rng_state():
    return dict(torch_rng=torch.get_rng_state(),cuda_rng=torch.cuda.get_rng_state_all(),
        numpy_rng=np.random.get_state(),python_rng=random.getstate())


def restore_rng(state):
    torch.set_rng_state(state['torch_rng']);torch.cuda.set_rng_state_all(state['cuda_rng'])
    np.random.set_state(state['numpy_rng']);random.setstate(state['python_rng'])


def check_monroe_losses():
    """Compare losses and gradients with the pinned author's actual implementation."""
    import openadmet_evaluate as runtime
    runtime.restore_monroe()
    from monroe.mtl.UW import UW
    source=WORK/'external/Monroe/monroe/train/loss.py'
    tree=ast.parse(source.read_text());function=next(n for n in tree.body if isinstance(n,ast.FunctionDef) and n.name=='off_diag_cov_loss')
    namespace={'torch':torch};exec(compile(ast.Module(body=[function],type_ignores=[]),str(source),'exec'),namespace)
    torch.manual_seed(42026)
    z=torch.randn(19,37,requires_grad=True);z2=z.detach().clone().requires_grad_()
    ref_cov=namespace['off_diag_cov_loss'](z);our_cov=off_diag_cov_loss(z2)
    ref_cov.backward();our_cov.backward()
    assert torch.equal(ref_cov,our_cov) and torch.equal(z.grad,z2.grad)
    reference=UW(task_types={k:None for k in CONFIG['task_order']},encoder_class=None,decoders=nn.Identity(),device='cpu')
    reference.init_param()
    losses=torch.tensor([.2,.8,1.3],requires_grad=True)
    ours=losses.detach().clone().requires_grad_();scale=nn.Parameter(reference.loss_scale.detach().clone())
    ref=reference.backward(losses,extra_losses=torch.tensor([.125]))
    value=uncertainty_objective(ours,scale)+.125;value.backward()
    assert torch.equal(ref,value) and torch.equal(reference.loss_scale.grad,scale.grad) and torch.equal(losses.grad,ours.grad)
    result=dict(passed=True,monroe_revision=MONROE_REV,uw_value=float(value.detach()),
        uw_value_and_gradients_exact=True,covariance_value_and_gradients_exact=True,
        uw_source_sha256=sha(WORK/'external/Monroe/monroe/mtl/UW.py'),covariance_source_sha256=sha(source),
        no_normalization_of_objective_weights=True,decorrelation_outside_uncertainty=True,
        decorrelation_weight=CONFIG['decorrelation_weight'])
    write(OUT/'loss_implementation_validation.json',result)
    return result


class TrainingData:
    def __init__(self):
        self.cosmo=cosmo_frame();self.pairs=pd.read_parquet(DATA/'combisolv_pairs.parquet')
        self.c_cache=RawCache('cosmo');self.q_cache=RawCache('combisolv')
        self.tr=np.flatnonzero(self.cosmo.joint_split.to_numpy()=='train');self.va=np.flatnonzero(self.cosmo.joint_split.to_numpy()=='validation')
        self.qtr=np.flatnonzero(self.pairs.split.to_numpy()=='train');self.qva=np.flatnonzero(self.pairs.split.to_numpy()=='validation')
        transform=json.loads((DATA/'cosmo_label_transform.json').read_text())
        self.c_mean=np.asarray(transform['mean']);self.c_scale=np.asarray(transform['scale'])
        labels=self.cosmo[[f'p_{i:03d}' for i in range(153)]+['area_A2','volume_A3']].to_numpy(dtype=np.float64)
        loglabels=np.c_[np.log1p(labels[:,:153]),np.log(labels[:,153:])]
        self.c_target=torch.tensor((loglabels-self.c_mean)/self.c_scale,dtype=torch.float32,device='cuda')
        self.q_transform=json.loads((DATA/'combisolv_label_transform.json').read_text())
        self.q_target=torch.tensor((self.pairs.dGsolv_kcal_mol.to_numpy()-self.q_transform['mean'])/self.q_transform['scale'],dtype=torch.float32,device='cuda')
        assert len(self.c_cache.lengths)==len(self.cosmo) and len(self.q_target)==1000000
        assert torch.isfinite(self.c_target).all() and torch.isfinite(self.q_target).all()

    def plans(self,epoch):
        rng=np.random.default_rng(CONFIG['seed']+epoch)
        cp=list(batches(self.tr,self.c_cache.lengths,rng,max_size=CONFIG['cosmo_batch_size'],atom_cap=CONFIG['cosmo_atom_cap']))
        order=rng.permutation(self.qtr)
        qp=[order[i:i+CONFIG['qm_pair_batch_size']] for i in range(0,len(order),CONFIG['qm_pair_batch_size'])]
        assert sorted(i for batch in cp for i in batch)==self.tr.tolist()
        assert np.array_equal(np.sort(np.concatenate(qp)),self.qtr)
        return cp,qp,max(len(cp),len(qp))


def update(model,optimizer,data,ci,qi):
    optimizer.zero_grad(set_to_none=True)
    losses,cov,z,_,parts=model.cosmo_denoising(data.c_cache.sample(ci),data.c_target[ci])
    first=uncertainty_objective(losses,model.log_variance[:2])+CONFIG['decorrelation_weight']*.5*cov
    # Accumulate both objective graphs before a single optimizer step. This is algebraically
    # the full three-task UW sum and average of the two covariance penalties.
    first.backward()
    log=dict(denoising=float(losses[0].detach()),denoising_raw=float(losses[0].detach())*1088,
        cosmo=float(losses[1].detach()),cov_cosmo=float(cov.detach()),
        descriptor_std_cosmo=float(z.detach().float().std(dim=0).mean()),**{'cosmo_'+k:v for k,v in parts.items()})
    first_value=float(first.detach())
    del losses,cov,z,first
    qm,cov,pred,z=model.qm(data.q_cache,qi,data.pairs,data.q_target)
    second=uncertainty_objective(qm[None],model.log_variance[2:])+CONFIG['decorrelation_weight']*.5*cov
    second.backward()
    log.update(combisolv=float(qm.detach()),cov_combisolv=float(cov.detach()),
        descriptor_std_combisolv=float(z.detach().float().std(dim=0).mean()),
        total=first_value+float(second.detach()))
    norm=nn.utils.clip_grad_norm_(model.parameters(),CONFIG['gradient_clip'],error_if_nonfinite=True)
    assert all(math.isfinite(v) for v in log.values()),log
    optimizer.step()
    log['gradient_norm']=float(norm)
    log['combisolv_rmse_kcal_mol']=math.sqrt(log['combisolv'])*data.q_transform['scale']
    for i,task in enumerate(CONFIG['task_order']):
        log['log_variance_'+task]=float(model.log_variance[i].detach())
        log['precision_'+task]=float((1/(2*model.log_variance[i].exp())).detach())
    return log


@torch.no_grad()
def validate(model,data):
    old_rng=rng_state();torch.manual_seed(83026);torch.cuda.manual_seed_all(83026)
    model.eval();sums=np.zeros(3);count=0
    for ids in batches(data.va,data.c_cache.lengths,max_size=64,atom_cap=3600):
        losses,cov,_,_,_=model.cosmo_denoising(data.c_cache.sample(ids),data.c_target[ids])
        sums+=np.asarray([float(losses[0]),float(losses[1]),float(cov)])*len(ids);count+=len(ids)
    assert count==len(data.va)
    q_sse=0.;q_mae=0.;q_count=0
    for start in range(0,len(data.qva),CONFIG['qm_pair_batch_size']):
        ids=data.qva[start:start+CONFIG['qm_pair_batch_size']]
        loss,_,pred,_=model.qm(data.q_cache,ids,data.pairs,data.q_target)
        diff=(pred-data.q_target[ids])*data.q_transform['scale']
        q_sse+=float(diff.double().square().sum());q_mae+=float(diff.double().abs().sum());q_count+=len(ids)
    assert q_count==len(data.qva)
    restore_rng(old_rng)
    result=dict(denoising=float(sums[0]/count),denoising_raw=float(sums[0]/count*1088),
        cosmo=float(sums[1]/count),cov_cosmo=float(sums[2]/count),
        combisolv_rmse_kcal_mol=math.sqrt(q_sse/q_count),combisolv_mae_kcal_mol=q_mae/q_count,
        cosmo_rows=count,combisolv_rows=q_count)
    assert all(math.isfinite(v) for v in result.values()),result
    return result


def save_checkpoint(path,model,optimizer,scheduler,epoch,global_step,metrics):
    pending=path.with_name(path.name+'.partial');pending.mkdir(parents=True,exist_ok=True)
    for name,module in model.parts.items():torch.save(module.state_dict(),pending/name)
    torch.save(dict(log_variance=model.log_variance.detach().cpu(),task_order=CONFIG['task_order']),pending/'task_weighting.pth')
    torch.save(dict(epoch=epoch,global_step=global_step,optimizer=optimizer.state_dict(),scheduler=scheduler.state_dict(),
        config=CONFIG,dataset=json.loads((OUT/'prepared_data_hub.json').read_text()),**rng_state()),pending/'training_state.pt')
    write(pending/'metrics.json',metrics);write(pending/'config.json',CONFIG)
    for filename in ['cosmo_label_transform.json','combisolv_label_transform.json','manifest.json']:
        shutil.copyfile(DATA/filename,pending/('data_manifest.json' if filename=='manifest.json' else filename))
    shutil.copyfile(OUT/'initialization.json',pending/'initialization.json')
    (pending/'post_training_architecture_config.yaml').write_text((common.ORIGINAL/'post_training_architecture_config.yaml').read_text())
    write(pending/'manifest.json',dict(epoch=epoch,global_step=global_step,
        files={p.name:dict(bytes=p.stat().st_size,sha256=sha(p)) for p in pending.iterdir() if p.is_file() and p.name!='manifest.json'}))
    if path.exists():shutil.rmtree(path)
    pending.rename(path)


def upload_checkpoint(path):
    marker=path/'hub_verified.json'
    if marker.exists():return json.loads(marker.read_text())
    api=HfApi();api.create_repo(MODEL_REPO,repo_type='model',exist_ok=True)
    files=[p for p in path.iterdir() if p.is_file() and p.name!='hub_verified.json']
    for attempt in range(5):
        try:
            commit=api.create_commit(repo_id=MODEL_REPO,
                operations=[CommitOperationAdd(path_in_repo=path.name+'/'+p.name,path_or_fileobj=str(p)) for p in files],
                commit_message=f'Preserve complete joint checkpoint {path.name}: all heads, learned task weights and exact resume state')
            entries={s.rfilename:s for s in api.model_info(MODEL_REPO,revision=commit.oid,files_metadata=True).siblings}
            for p in files:
                if p.suffix not in ['.pth','.pt']:continue
                remote=entries[path.name+'/'+p.name]
                assert remote.lfs and remote.lfs.sha256==sha(p),p.name
            result=dict(repository=MODEL_REPO,revision=commit.oid,epoch=int(path.name.split('_')[-1]),verified=True)
            write(marker,result);return result
        except Exception:
            if attempt==4:raise
            print('HUB_UPLOAD_RETRY',traceback.format_exc(),flush=True);time.sleep(min(10*2**attempt,50))


def smoke(data):
    verified_path=OUT/'loss_implementation_validation.json'
    if verified_path.exists():
        verified=json.loads(verified_path.read_text())
        assert verified['passed']
        assert verified['uw_source_sha256']==sha(WORK/'external/Monroe/monroe/mtl/UW.py')
        assert verified['covariance_source_sha256']==sha(WORK/'external/Monroe/monroe/train/loss.py')
    else:check_monroe_losses()
    model=MultiTask();optimizer=optimizer_for(model)
    scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lambda e:1.)
    model.eval();baseline=RawCache('bigsoldb')
    with torch.no_grad():
        ref=np.load(WORK/'bigsoldb_evaluation/features/cosmo_epoch10.npz')['X']
        ours=np.empty_like(ref)
        for ids in common.batch_indices(np.arange(len(ref)),baseline.lengths,max_size=24,atom_cap=2400):
            _,_,z=model.encode(baseline.sample(ids));ours[ids]=z.cpu().numpy()
        difference=np.abs(ours-ref)
        assert np.allclose(ours,ref,atol=2e-4,rtol=2e-4),(float(difference.max()),float(difference.mean()))
    initial_difference=dict(molecules=len(ref),max_absolute=float(difference.max()),mean_absolute=float(difference.mean()))
    frozen={n:b.cpu().clone() for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers()}
    cp,qp,_=data.plans(1);ci=cp[10];qi=qp[10]
    initial=next(model.encoder.parameters()).detach().clone()
    model.train();torch.cuda.reset_peak_memory_stats();t=time.perf_counter()
    first=update(model,optimizer,data,ci,qi);duration=time.perf_counter()-t
    assert not torch.equal(initial,next(model.encoder.parameters()).detach())
    gradients={name:any(p.grad is not None and bool((p.grad!=0).any()) for p in module.parameters()) for name,module in model.parts.items()}
    assert all(gradients.values()),gradients
    assert model.log_variance.grad is not None and bool((model.log_variance.grad!=0).all())
    temp=OUT/'smoke_checkpoint';save_checkpoint(temp,model,optimizer,scheduler,0,1,first)
    expected=update(model,optimizer,data,ci,qi)
    expected_state={n:v.detach().cpu().clone() for n,v in model.state_dict().items()}
    model.load_parts(temp);state=torch.load(temp/'training_state.pt',map_location='cpu',weights_only=False)
    optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']);restore_rng(state)
    replay_started=time.perf_counter()
    actual=update(model,optimizer,data,ci,qi)
    steady_duration=time.perf_counter()-replay_started
    assert all(np.isclose(actual[k],expected[k],atol=1e-7,rtol=1e-7) for k in expected),(expected,actual)
    assert all(torch.equal(v.detach().cpu(),expected_state[n]) for n,v in model.state_dict().items()),'Resume changed next-step parameters'
    for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers():assert torch.equal(b.cpu(),frozen[n])
    peak=torch.cuda.max_memory_allocated()/1e9
    result=dict(passed=True,all_head_gradients=gradients,encoder_updated=True,task_weights_updated=True,
        optimizer_and_rng_resume_next_step_exact=True,normalization_buffers_unchanged=True,
        initial_bigsoldb_embedding_match=initial_difference,real_batch_loss=first,seconds_per_joint_step=duration,
        steady_seconds_per_joint_step=steady_duration,
        smoke_cosmo_molecules=len(ci),smoke_qm_pairs=len(qi),peak_gpu_GB=peak,config=CONFIG)
    write(OUT/'smoke_validation.json',result);shutil.rmtree(temp)
    status('smoke_passed',seconds_per_joint_step=duration,peak_gpu_GB=peak,initial_embedding_difference=initial_difference)


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--smoke',action='store_true');parser.add_argument('--check-losses',action='store_true');args=parser.parse_args()
    OUT.mkdir(exist_ok=True)
    lock=(OUT/'training.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    if args.check_losses:print(check_monroe_losses());return
    assert (OUT/'preparation_complete.json').exists(),'Data preparation has not finished'
    # Ensure repeated molecule-index gradients and CUDA GEMMs replay exactly.
    torch.use_deterministic_algorithms(True)
    torch.manual_seed(CONFIG['seed']);torch.cuda.manual_seed_all(CONFIG['seed']);np.random.seed(CONFIG['seed']);random.seed(CONFIG['seed'])
    data=TrainingData()
    if args.smoke:smoke(data);return
    smoke_result=json.loads((OUT/'smoke_validation.json').read_text())
    assert smoke_result['passed'] and smoke_result['config']==CONFIG
    checkpoints=sorted((OUT/'checkpoints').glob('epoch_[0-9][0-9]')) if (OUT/'checkpoints').exists() else []
    model=MultiTask(checkpoints[-1] if checkpoints else None);optimizer=optimizer_for(model)
    scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lambda e:.1+.9*.5*(1+math.cos(math.pi*e/CONFIG['epochs'])))
    frozen={n:b.cpu().clone() for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers()}
    start=1;global_step=0
    if checkpoints:
        state=torch.load(checkpoints[-1]/'training_state.pt',map_location='cpu',weights_only=False)
        assert state['config']==CONFIG
        assert state['dataset']==json.loads((OUT/'prepared_data_hub.json').read_text())
        optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']);restore_rng(state)
        start=state['epoch']+1;global_step=state['global_step'];upload_checkpoint(checkpoints[-1])
        write(OUT/'resume_state.json',dict(restored_epoch=start-1,optimizer_scheduler_rng_restored=True,global_step=global_step))
    import wandb
    wb_path=OUT/'wandb.json';wb=json.loads(wb_path.read_text()) if wb_path.exists() else dict(run_id=wandb.util.generate_id())
    write(wb_path,wb)
    run=wandb.init(entity='queimo',project='rem3di-reproduction',id=wb['run_id'],resume='allow',mode='online',
        name='pcqm2-cosmo10-combisolv-three-task-5epochs',group='rem3di-cosmo-combisolv',job_type='multitask-pretraining',
        config={**CONFIG,'dataset':json.loads((DATA/'manifest.json').read_text()),'initialization':json.loads((OUT/'initialization.json').read_text())},
        dir=str(OUT),settings=wandb.Settings(init_timeout=90))
    write(wb_path,dict(run_id=run.id,url=run.url));write(OUT/'config.json',CONFIG)
    run.define_metric('global_step');run.define_metric('train/*',step_metric='global_step')
    run.define_metric('epoch');run.define_metric('validation/*',step_metric='epoch')
    if checkpoints:restore_rng(state)
    try:
        for epoch in range(start,CONFIG['epochs']+1):
            model.train();cp,qp,total=data.plans(epoch);started=time.time();history=[];seen_c=seen_q=0
            status('training',epoch=epoch,epochs=CONFIG['epochs'],batches=total,cosmo_batches=len(cp),qm_batches=len(qp),wandb_url=run.url)
            for bi in range(total):
                ci=cp[bi%len(cp)];qi=qp[bi%len(qp)]
                log=update(model,optimizer,data,ci,qi);history.append(log);global_step+=1;seen_c+=len(ci);seen_q+=len(qi)
                if bi==0 or (bi+1)%20==0 or bi+1==total:
                    run.log(dict(global_step=global_step,epoch=epoch,**{'train/'+k:v for k,v in log.items()}))
                    status('training',epoch=epoch,epochs=CONFIG['epochs'],batch=bi+1,batches=total,global_step=global_step,
                        cosmo_molecules_seen=seen_c,combisolv_pairs_seen=seen_q,seconds=time.time()-started,
                        losses={k:log[k] for k in ['denoising','cosmo','combisolv','total']},wandb_url=run.url)
            assert seen_c>=len(data.tr) and seen_q>=len(data.qtr)
            status('validating',epoch=epoch,wandb_url=run.url)
            validation=validate(model,data)
            for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers():assert torch.equal(b.cpu(),frozen[n])
            scheduler.step()
            result=dict(epoch=epoch,global_step=global_step,seconds=time.time()-started,
                train={k:float(np.mean([h[k] for h in history])) for k in history[0]},validation=validation,
                coverage=dict(cosmo_unique_rows=len(data.tr),combisolv_unique_rows=len(data.qtr),
                    cosmo_rows_with_repeats=seen_c,combisolv_rows_with_repeats=seen_q,every_training_row_covered=True),
                task_log_variances=model.log_variance.detach().cpu().tolist())
            target=OUT/'checkpoints'/f'epoch_{epoch:02d}'
            save_checkpoint(target,model,optimizer,scheduler,epoch,global_step,result)
            status('uploading_checkpoint',epoch=epoch,wandb_url=run.url)
            hub=upload_checkpoint(target)
            write(OUT/'latest_epoch.json',dict(**result,hub=hub,wandb_url=run.url))
            run.log(dict(epoch=epoch,**{'validation/'+k:v for k,v in validation.items()},
                **{'epoch/'+k:v for k,v in result['train'].items()},**{'backup/verified_epoch':epoch,'backup/revision':hub['revision']}))
            # Run frozen Kumo in a separate process after releasing all training GPU tensors.
            # The driver performs all requested checkpoint probes after training finishes.
        write(OUT/'training_complete.json',dict(epochs=CONFIG['epochs'],wandb_url=run.url,completed_unix=time.time()))
        run.summary.update(dict(training_complete=True,epochs=CONFIG['epochs']));run.finish()
        status('training_complete',epochs=CONFIG['epochs'])
    except Exception:
        status('failed_training',error=traceback.format_exc());run.finish(exit_code=1);raise


if __name__=='__main__':
    try:main()
    except Exception:status('failed_training',error=traceback.format_exc());raise