"""Ten-epoch supervised COSMO continuation of the selected PCQM epoch 2."""
import argparse
import fcntl
import gc
import hashlib
import json
import math
import os
from pathlib import Path
import random
import shutil
import sys
import time
import traceback

WORK = Path(__file__).resolve().parent.parent
AUDIT = Path('/marimo/rem3di_training/corrected_20261027/evaluation/paper_audit')
ORIGINAL = AUDIT/'author_checkpoint/exp138/polar'
sys.path[:0] = [str(AUDIT),str(AUDIT/'author_legacy_source'),'/marimo/Rem3Di']
os.environ.update(CUDA_VISIBLE_DEVICES='0',OMP_NUM_THREADS='8',MKL_NUM_THREADS='8',WANDB_MODE='online')
from dotenv import load_dotenv
load_dotenv('/marimo/.env',override=False)
import numpy as np
import pandas as pd
import torch
from torch import nn
import yaml
from huggingface_hub import HfApi, hf_hub_download, snapshot_download, CommitOperationAdd
from threedscriptors.configuration.architecture_config import EncoderOnlyArchitectureConfig
from threedscriptors.data_handling.sample import Sample as LegacySample
from remedi.data_handling.sample import Sample, yield_molecules_collate_fn
from build_author_mace_runtime import build_author_mace_runtime

torch.serialization.add_safe_globals([slice])
torch.set_num_threads(8)


def write(path,value):
    path.parent.mkdir(parents=True,exist_ok=True)
    tmp=path.with_suffix(path.suffix+'.partial')
    tmp.write_text(json.dumps(value,indent=2,allow_nan=False));tmp.replace(path)


def sha(path):
    h=hashlib.sha256()
    with path.open('rb') as f:
        for b in iter(lambda:f.read(8<<20),b''):h.update(b)
    return h.hexdigest()


def status(phase,**kw):
    row=dict(phase=phase,pid=os.getpid(),updated_unix=time.time(),**kw)
    write(WORK/'training_status.json',row)
    print('STATUS',json.dumps(row),flush=True)


def backbone(path):
    d=yaml.safe_load((ORIGINAL/'post_training_architecture_config.yaml').read_text())
    d.pop('decoder_config',None);d['kind']='encoder_only';d['mace_config']=None
    cfg=EncoderOnlyArchitectureConfig.model_validate(d)
    pre,enc=cfg._build_preprocessor(),cfg._build_encoder()
    # Preserve the learned chiral weights in the checkpoint's original precision.
    pre.atomic_preprocessor.double();pre.geometric_preprocessor.float();enc.float()
    for name,module in [('encoder.pth',enc),('atomic_preprocessor.pth',pre.atomic_preprocessor),('geometric_preprocessor.pth',pre.geometric_preprocessor)]:
        module.load_state_dict(torch.load(path/name,map_location='cpu',weights_only=True),strict=True)
    return pre.cuda(),enc.cuda()


class COSMORegressor(nn.Module):
    def __init__(self,pre,enc):
        super().__init__();self.pre=pre;self.encoder=enc
        self.head=nn.Sequential(nn.Linear(640,256),nn.SiLU(),nn.Linear(256,155)).cuda()

    def forward(self,sample):
        z=self.encoder(self.pre(sample))
        if not torch.is_tensor(z):z=z.flat
        return self.head(z)


def batches(indices,lengths,rng=None,max_size=96,atom_cap=4000):
    order=np.asarray(indices).copy()
    if rng is not None:rng.shuffle(order)
    for start in range(0,len(order),1024):
        chunk=sorted(order[start:start+1024],key=lambda i:lengths[i]);ids=[];atoms=0
        for i in chunk:
            if ids and (atoms+lengths[i]>atom_cap or len(ids)>=max_size):
                yield ids;ids=[];atoms=0
            ids.append(int(i));atoms+=lengths[i]
        if ids:yield ids


def inverse(z,mean,scale):
    t=z*scale+mean
    # Clamp only the metrics' physical reconstruction to the valid range.
    return torch.cat([torch.expm1(t[:,:153]).clamp_min(0),torch.exp(t[:,153:])],dim=1)


def loss_terms(pred,target):
    e=(pred-target).square()
    bins=e[:,:153].mean();area=e[:,153].mean();volume=e[:,154].mean()
    return (bins+area+volume)/3,dict(sigma=float(bins.detach()),area=float(area.detach()),volume=float(volume.detach()))


def upload_epoch(api,cfg,path):
    marker=path/'hub_verified.json'
    if marker.exists():return json.loads(marker.read_text())
    names=['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth','cosmo_head.pth',
           'training_state.pt','metrics.json','manifest.json','label_transform.json','post_training_architecture_config.yaml','cosmo_decoder_config.json']
    ops=[CommitOperationAdd(path_in_repo=f'epoch_{path.name.split("_")[-1]}/{n}',path_or_fileobj=str(path/n)) for n in names]
    commit=api.create_commit(repo_id=cfg['model_repository'],operations=ops,
                             commit_message=f'Save complete COSMO epoch {path.name}: weights, optimizer, scheduler, RNG and transform')
    entries={s.rfilename:s for s in api.model_info(cfg['model_repository'],revision=commit.oid,files_metadata=True).siblings}
    for n in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth','cosmo_head.pth','training_state.pt']:
        entry=entries[f'{path.name}/{n}']
        remote=entry.lfs.sha256 if entry.lfs else None
        assert remote==sha(path/n),(n,remote)
    result=dict(epoch=int(path.name.split('_')[-1]),repository=cfg['model_repository'],revision=commit.oid,verified=True,time=time.time())
    write(marker,result)
    return result


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--smoke',action='store_true')
    parser.add_argument('--verify-resume',action='store_true');args=parser.parse_args()
    lock=(WORK/'training.lock').open('a')
    try:fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    except BlockingIOError:
        print('TRAINING_ALREADY_RUNNING',flush=True);return
    cfg=json.loads((WORK/'config.json').read_text())
    api=HfApi()
    status('loading_published_dataset')
    reference_path=WORK/'dataset_reference.json'
    reference=json.loads(reference_path.read_text()) if reference_path.exists() else {}
    revision=reference.get('revision') or api.dataset_info(cfg['dataset_repository']).sha
    data_root=Path(snapshot_download(cfg['dataset_repository'],repo_type='dataset',revision=revision,
                                      allow_patterns=['data/*','manifest.json','label_transform.json','pm6_source.json']))
    manifest=json.loads((data_root/'manifest.json').read_text())
    assert manifest['rows']==361583
    # Check every downloaded training shard against the independently assembled manifest.
    for name,digest in manifest['files'].items():assert sha(data_root/name)==digest,name
    transform=json.loads((data_root/'label_transform.json').read_text())
    frames=[]
    for split in ['train','validation']:
        frame=pd.concat([pd.read_parquet(p) for p in sorted((data_root/'data').glob(split+'-*.parquet'))],ignore_index=True)
        assert len(frame)==manifest['splits'][split]
        frame['is_validation']=split=='validation';frames.append(frame)
    df=pd.concat(frames,ignore_index=True)
    tr=np.flatnonzero(~df.is_validation.values);va=np.flatnonzero(df.is_validation.values)
    assert not set(df.molecule_group.iloc[tr])&set(df.molecule_group.iloc[va])
    target_cols=[f'p_{i:03d}' for i in range(153)]+['area_A2','volume_A3']
    labels=df[target_cols].to_numpy(np.float64)
    transformed=np.concatenate([np.log1p(labels[:,:153]),np.log(labels[:,153:])],axis=1)
    mean=np.asarray(transform['mean']);scale=np.asarray(transform['scale'])
    assert np.allclose(transformed[tr].mean(0),mean,atol=1e-10)
    assert np.allclose(np.maximum(transformed[tr].std(0),1e-3),scale,atol=1e-10)
    targets=torch.tensor((transformed-mean)/scale,dtype=torch.float32,device='cuda')
    mean=torch.tensor(mean,dtype=torch.float32,device='cuda');scale=torch.tensor(scale,dtype=torch.float32,device='cuda')
    positions=[torch.tensor(np.stack(x),dtype=torch.float32) for x in df.positions]
    numbers=[torch.tensor(np.asarray(x),dtype=torch.long) for x in df.atomic_numbers]
    lengths=np.array([len(z) for z in numbers])
    assert np.isfinite(labels).all()
    write(WORK/'dataset_reference.json',dict(repository=cfg['dataset_repository'],revision=revision,manifest=manifest))
    completed=sorted((WORK/'checkpoints').glob('epoch_[0-9][0-9]')) if (WORK/'checkpoints').exists() else []
    if completed and not args.smoke:
        init=completed[-1]
    else:
        init=WORK/'initial_epoch_02';init.mkdir(exist_ok=True)
        for name in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth','manifest.json']:
            src=Path(hf_hub_download(cfg['initial_checkpoint_repository'],
                filename=cfg['initial_checkpoint_path']+'/'+name,revision=cfg['initial_checkpoint_revision']))
            if not (init/name).exists():shutil.copyfile(src,init/name)
    initial_manifest=json.loads((init/'manifest.json').read_text())
    for name in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth']:
        assert sha(init/name)==initial_manifest['files'][name]['sha256']
    torch.manual_seed(cfg['seed']);torch.cuda.manual_seed_all(cfg['seed']);np.random.seed(cfg['seed']);random.seed(cfg['seed'])
    pre,enc=backbone(init)
    model=COSMORegressor(pre,enc)
    optimizer=torch.optim.AdamW([
        {'params':list(pre.parameters())+list(enc.parameters()),'lr':cfg['learning_rate_backbone']},
        {'params':model.head.parameters(),'lr':cfg['learning_rate_head']}],weight_decay=cfg['weight_decay'])
    # Epoch schedule has a stable horizon even when atom-budget packing varies.
    scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lambda e:.1+.9*.5*(1+math.cos(math.pi*e/cfg['epochs'])))
    mace=build_author_mace_runtime().cuda().eval().requires_grad_(False)
    captures=[]
    hook=mace.preprocessor.atomic_preprocessor.register_forward_pre_hook(lambda m,a:captures.append(a[0].detach()))
    frozen_buffers={n:b.detach().cpu().clone() for n,b in pre.atomic_preprocessor.invariant_normalization.named_buffers()}

    def sample(ids):
        batch=yield_molecules_collate_fn([Sample(atomic_positions=positions[i],atomic_numbers=numbers[i],total_charge=0.,multiplicity=1.) for i in ids]).to_('cuda')
        captures.clear()
        with torch.no_grad():mace.preprocessor(batch)
        assert len(captures)==1
        raw=captures.pop().clone()
        return LegacySample(embeddings=raw,atomic_positions=batch.atomic_positions,padding_mask=batch.padding_mask)

    def update(ids):
        optimizer.zero_grad(set_to_none=True)
        pred=model(sample(ids));loss,parts=loss_terms(pred,targets[ids])
        assert torch.isfinite(loss),'Nonfinite COSMO loss'
        loss.backward()
        norm=nn.utils.clip_grad_norm_(model.parameters(),1.,error_if_nonfinite=True)
        optimizer.step()
        return float(loss.detach()),parts,float(norm)

    if args.smoke:
        ids=tr[:8].tolist();before=next(enc.parameters()).detach().clone()
        model.train();result=update(ids)
        assert not torch.equal(before,next(enc.parameters()).detach()),'Encoder did not update'
        assert any(p.grad is not None and bool((p.grad!=0).any()) for p in pre.parameters())
        assert all(p.grad is None for p in mace.parameters())
        model.eval()
        with torch.no_grad():
            s=sample(ids);ref=model(s);physical=inverse(targets[ids],mean,scale)
            assert np.allclose(physical.cpu().numpy(),labels[ids],rtol=1e-5,atol=1e-4)
        state={k:v.cpu() for k,v in model.state_dict().items()}
        path=WORK/'smoke_model.pt';torch.save(state,path)
        model.load_state_dict(torch.load(path,map_location='cpu',weights_only=True))
        with torch.no_grad():assert torch.allclose(ref,model(s),atol=1e-6,rtol=1e-6)
        for n,b in pre.atomic_preprocessor.invariant_normalization.named_buffers():assert torch.equal(b.cpu(),frozen_buffers[n])
        write(WORK/'smoke_validation.json',dict(passed=True,loss=result[0],gradient_norm=result[2],
             decoder_outputs=155,encoder_updated=True,preprocessor_gradients=True,mace_frozen=True,
             transform_roundtrip=True,checkpoint_roundtrip=True,dataset_revision=revision))
        path.unlink();status('smoke_passed');return

    smoke=json.loads((WORK/'smoke_validation.json').read_text())
    assert smoke['passed'] and smoke['dataset_revision']==revision,'Smoke validation must match the pinned dataset'
    import wandb
    api.create_repo(cfg['model_repository'],repo_type='model',exist_ok=True)
    wb=json.loads((WORK/'wandb.json').read_text()) if (WORK/'wandb.json').exists() else {'run_id':wandb.util.generate_id()}
    write(WORK/'wandb.json',wb)
    run=None
    if not args.verify_resume:
        run=wandb.init(entity=cfg['wandb_entity'],project=cfg['wandb_project'],id=wb['run_id'],resume='allow',mode='online',
        name='rem3di-pcqm-epoch2-cosmo-10epochs',group='rem3di-cosmo',job_type='cosmo-posttraining',
        config={**cfg,'dataset_revision':revision,'loss':'(mean sigma z-score MSE + A z-score MSE + V z-score MSE)/3',
                'train_rows':len(tr),'validation_rows':len(va),'mace':'frozen; original PM6 coordinates; no geometry augmentation'},
            dir=str(WORK),settings=wandb.Settings(init_timeout=90))
        write(WORK/'wandb.json',dict(run_id=run.id,url=run.url))
        run.define_metric('global_step');run.define_metric('train/*',step_metric='global_step')
        run.define_metric('epoch');run.define_metric('validation/*',step_metric='epoch')
    start_epoch=1;global_step=0
    if completed:
        last=completed[-1]
        saved_manifest=json.loads((last/'manifest.json').read_text())
        for name in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth','cosmo_head.pth','training_state.pt']:
            assert sha(last/name)==saved_manifest['files'][name]['sha256'],name
        for n,module in [('encoder.pth',enc),('atomic_preprocessor.pth',pre.atomic_preprocessor),('geometric_preprocessor.pth',pre.geometric_preprocessor),('cosmo_head.pth',model.head)]:
            module.load_state_dict(torch.load(last/n,map_location='cpu',weights_only=True))
        state=torch.load(last/'training_state.pt',map_location='cpu',weights_only=False)
        assert state['dataset_revision']==revision,'Resume must preserve the original dataset revision'
        assert state['config']==cfg,'Resume configuration differs from the saved run'
        assert state['epoch']==int(last.name.split('_')[-1])
        optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler'])
        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'])
        start_epoch=state['epoch']+1;global_step=state['global_step']
        upload_epoch(api,cfg,last)
        write(WORK/'resume_state.json',dict(restored_epoch=state['epoch'],next_epoch=start_epoch,
             global_step=global_step,dataset_revision=revision,optimizer_restored=True,scheduler_restored=True,rng_restored=True))
        print('RESUMED',json.dumps(dict(epoch=state['epoch'],next_epoch=start_epoch,global_step=global_step)),flush=True)
    if args.verify_resume:
        assert completed,'No completed checkpoint to verify'
        model.eval()
        with torch.no_grad():
            output=model(sample(tr[:8].tolist()))
            loss,_=loss_terms(output,targets[tr[:8]])
            physical=inverse(output,mean,scale)
        assert output.shape==(8,155) and torch.isfinite(output).all() and torch.isfinite(physical).all()
        assert torch.isfinite(loss) and optimizer.state and scheduler.last_epoch==start_epoch-1
        write(WORK/'resume_validation.json',dict(passed=True,restored_epoch=start_epoch-1,next_epoch=start_epoch,
             global_step=global_step,loss=float(loss),optimizer_state_tensors=len(optimizer.state),
             scheduler_epoch=scheduler.last_epoch,dataset_revision=revision,finite_predictions=True))
        hook.remove();status('resume_verified',next_epoch=start_epoch,global_step=global_step);return
    try:
        for epoch in range(start_epoch,cfg['epochs']+1):
            model.train();started=time.time();loss_sum=0.;count=0
            rng=np.random.default_rng(cfg['seed']+epoch)
            plan=list(batches(tr,lengths,rng,max_size=cfg['max_molecules_per_batch'],atom_cap=cfg['max_atoms_per_batch']))
            status('training',epoch=epoch,epochs=cfg['epochs'],batches=len(plan),wandb_url=run.url)
            for bi,ids in enumerate(plan,1):
                loss,parts,norm=update(ids)
                global_step+=1;loss_sum+=loss*len(ids);count+=len(ids)
                if bi==1 or bi%25==0:
                    row=dict(global_step=global_step,epoch=epoch,batch=bi,batches=len(plan),
                        **{'train/'+k:v for k,v in dict(loss=loss,gradient_norm=norm,backbone_lr=optimizer.param_groups[0]['lr'],head_lr=optimizer.param_groups[1]['lr'],**parts).items()})
                    run.log(row)
                    status('training',**row,wandb_url=run.url,molecules_seen=count,seconds=time.time()-started)
            model.eval();val_sum=0.;val_count=0;sse=np.zeros(3);mae=np.zeros(3);sumarea=0.
            with torch.no_grad():
                for ids in batches(va,lengths,max_size=cfg['max_molecules_per_batch'],atom_cap=cfg['max_atoms_per_batch']):
                    pred=model(sample(ids));loss,_=loss_terms(pred,targets[ids])
                    val_sum+=float(loss)*len(ids);val_count+=len(ids)
                    physical=inverse(pred,mean,scale).cpu().numpy()
                    assert np.isfinite(physical).all()
                    err=physical-labels[ids]
                    sse += [float((err[:,:153]**2).sum()/153),float((err[:,153]**2).sum()),float((err[:,154]**2).sum())]
                    mae += [float(abs(err[:,:153]).sum()/153),float(abs(err[:,153]).sum()),float(abs(err[:,154]).sum())]
                    sumarea+=float(abs(physical[:,:153].sum(1)-physical[:,153]).sum())
            assert count==len(tr) and val_count==len(va)
            for n,b in pre.atomic_preprocessor.invariant_normalization.named_buffers():assert torch.equal(b.cpu(),frozen_buffers[n])
            scheduler.step()
            metrics=dict(epoch=epoch,global_step=global_step,train_loss=loss_sum/count,validation_loss=val_sum/val_count,
                validation_sigma_rmse=float(np.sqrt(sse[0]/val_count)),validation_area_rmse=float(np.sqrt(sse[1]/val_count)),
                validation_volume_rmse=float(np.sqrt(sse[2]/val_count)),validation_sigma_mae=float(mae[0]/val_count),
                validation_area_mae=float(mae[1]/val_count),validation_volume_mae=float(mae[2]/val_count),
                validation_profile_area_inconsistency=float(sumarea/val_count),seconds=time.time()-started)
            assert all(math.isfinite(v) for v in metrics.values())
            target=WORK/'checkpoints'/f'epoch_{epoch:02d}'
            pending=target.with_name(target.name+'.partial');pending.mkdir(parents=True,exist_ok=True)
            for n,module in [('encoder.pth',enc),('atomic_preprocessor.pth',pre.atomic_preprocessor),('geometric_preprocessor.pth',pre.geometric_preprocessor),('cosmo_head.pth',model.head)]:
                torch.save(module.state_dict(),pending/n)
            torch.save(dict(epoch=epoch,global_step=global_step,optimizer=optimizer.state_dict(),scheduler=scheduler.state_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(),dataset_revision=revision,config=cfg),pending/'training_state.pt')
            write(pending/'metrics.json',metrics)
            shutil.copyfile(data_root/'label_transform.json',pending/'label_transform.json')
            architecture=yaml.safe_load((ORIGINAL/'post_training_architecture_config.yaml').read_text())
            architecture.pop('decoder_config',None);architecture['kind']='encoder_only'
            architecture['mace_config']['model_path']='MACE-POLAR-1-M.model'
            (pending/'post_training_architecture_config.yaml').write_text(yaml.safe_dump(architecture,sort_keys=False))
            write(pending/'cosmo_decoder_config.json',dict(input_features=640,hidden_features=256,output_features=155,activation='SiLU'))
            write(pending/'manifest.json',dict(epoch=epoch,dataset_revision=revision,config=cfg,
                files={p.name:dict(bytes=p.stat().st_size,sha256=sha(p)) for p in pending.iterdir() if p.is_file()}))
            pending.rename(target)
            status('uploading_checkpoint',epoch=epoch,metrics=metrics,wandb_url=run.url)
            for attempt in range(5):
                try:
                    result=upload_epoch(api,cfg,target);break
                except Exception:
                    if attempt==4:raise
                    print('UPLOAD_RETRY',traceback.format_exc(),flush=True);time.sleep(min(10*2**attempt,60))
            run.log(dict(epoch=epoch,**{'validation/'+k.removeprefix('validation_'):v for k,v in metrics.items() if k.startswith('validation_')},
                **{'epoch/'+k:v for k,v in metrics.items()},**{'backup/revision':result['revision'],'backup/verified_epoch':epoch}))
            write(WORK/'latest_epoch.json',dict(**metrics,hub=result,wandb_url=run.url))
        write(WORK/'training_complete.json',dict(epochs=cfg['epochs'],wandb_url=run.url,completed_unix=time.time()))
        run.summary['training_complete']=True;run.finish();status('complete',epochs=cfg['epochs'])
    except Exception:
        status('failed',error=traceback.format_exc());run.finish(exit_code=1);raise
    finally:
        hook.remove()


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