"""Train one fifth-epoch, save an exact-resume checkpoint, then release the GPU."""
from pathlib import Path
import sys
sys.path.insert(0,str(Path(__file__).resolve().parent))
from scratch_model import *
import random,math,gc,fcntl,shutil,traceback,argparse,importlib.metadata,platform
from huggingface_hub import HfApi,CommitOperationAdd


def status(phase,**extra):
    write(OUT/'status.json',dict(phase=phase,pid=os.getpid(),updated_unix=time.time(),**extra))
    print('SCRATCH_TRAIN',phase,json.dumps(extra),flush=True)


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 seed():
    random.seed(CONFIG['seed']);np.random.seed(CONFIG['seed']);torch.manual_seed(CONFIG['seed']);torch.cuda.manual_seed_all(CONFIG['seed'])
    torch.use_deterministic_algorithms(True)


class TrainingData:
    def __init__(self):
        label_columns=[f'p_{i:03d}' for i in range(153)]+['area_A2','volume_A3']
        self.pcqm=pd.read_parquet(DATA/'pcqm.parquet',columns=['split'])
        self.cosmo=pd.read_parquet(DATA/'cosmo.parquet',columns=['split',*label_columns])
        self.p_cache=RawCache('pcqm');self.c_cache=RawCache('cosmo')
        self.pt=np.flatnonzero(self.pcqm.split.to_numpy()=='train');self.pv=np.flatnonzero(self.pcqm.split.to_numpy()=='validation')
        self.ct=np.flatnonzero(self.cosmo.split.to_numpy()=='train');self.cv=np.flatnonzero(self.cosmo.split.to_numpy()=='validation')
        transform=json.loads((DATA/'cosmo_label_transform.json').read_text())
        y=self.cosmo[[f'p_{i:03d}' for i in range(153)]+['area_A2','volume_A3']].to_numpy(dtype=np.float64)
        y=np.c_[np.log1p(y[:,:153]),np.log(y[:,153:])]
        self.labels=torch.tensor((y-np.asarray(transform['mean']))/np.asarray(transform['scale']),dtype=torch.float32,device='cuda')
        assert torch.isfinite(self.labels).all()
        self.epoch_plans={e:self.plan(e) for e in range(1,CONFIG['epochs']+1)}
        self.total_steps=sum(v[2] for v in self.epoch_plans.values())
        schedule=[];offset=0
        for epoch,(pp,cp,n) in self.epoch_plans.items():
            for fifth in range(1,6):
                end=(fifth*n+4)//5;start=((fifth-1)*n+4)//5
                schedule.append(dict(point=len(schedule)+1,epoch=epoch,fifth=fifth,batch_start=start,batch_end=end,
                    batches=n,epoch_position=(epoch-1)+end/n,nominal_epoch=(epoch-1)+fifth/5,global_step=offset+end))
            offset+=n
        assert len(schedule)==50 and schedule[-1]['global_step']==self.total_steps
        write(OUT/'probe_schedule.json',dict(total_steps=self.total_steps,probes=50,points=schedule))
        self.schedule=schedule

    def plan(self,epoch):
        rng=np.random.default_rng(CONFIG['seed']+epoch)
        pp=list(batches(self.pt,self.p_cache.lengths,rng,max_size=CONFIG['pcqm_batch_size'],atom_cap=CONFIG['pcqm_atom_cap']))
        cp=list(batches(self.ct,self.c_cache.lengths,rng,max_size=CONFIG['cosmo_batch_size'],atom_cap=CONFIG['cosmo_atom_cap']))
        assert np.array_equal(np.sort(np.concatenate(pp)),self.pt)
        assert np.array_equal(np.sort(np.concatenate(cp)),self.ct)
        return pp,cp,max(len(pp),len(cp))


def update(model,optimizer,data,pi,ci):
    optimizer.zero_grad(set_to_none=True)
    raw,z,target=model.denoising(data.p_cache.sample(pi))
    first=uncertainty(raw[None],model.log_variance[:1]);first.backward()
    log=dict(denoising_raw=float(raw.detach()),denoising_per_channel=float(raw.detach())/1088,
        descriptor_std_pcqm=float(z.detach().std(dim=0,unbiased=False).mean()),weighted_denoising=float(first.detach()))
    del raw,z,target,first
    loss,z,pred,parts=model.cosmo(data.c_cache.sample(ci),data.labels[ci])
    second=uncertainty(loss[None],model.log_variance[1:]);second.backward()
    log.update(cosmo=float(loss.detach()),descriptor_std_cosmo=float(z.detach().std(dim=0,unbiased=False).mean()),
        weighted_cosmo=float(second.detach()),**{'cosmo_'+k:v for k,v in parts.items()})
    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.update(gradient_norm=float(norm),total=log['weighted_denoising']+log['weighted_cosmo'],
        log_variance_denoising=float(model.log_variance[0].detach()),log_variance_cosmo=float(model.log_variance[1].detach()))
    return log


@torch.no_grad()
def validate(model,data):
    old=rng_state();torch.manual_seed(83026);torch.cuda.manual_seed_all(83026);model.eval()
    result={}
    try:
        for name,ids,cache in [('pcqm',data.pv,data.p_cache),('cosmo',data.cv,data.c_cache)]:
            total=0.;count=0;zs=np.zeros(640);zss=zs.copy();std=0.;atoms=0
            for group in batches(ids,cache.lengths,max_size=64,atom_cap=3600):
                if name=='pcqm':
                    loss,z,target=model.denoising(cache.sample(group))
                    valid=target[~(torch.arange(target.shape[1],device='cuda')[None]>=torch.tensor(cache.lengths[group],device='cuda')[:,None])]
                    total+=float(loss)*len(valid);atoms+=len(valid)
                    std+=float(valid.std(dim=0).mean())*len(valid)
                else:
                    loss,z,_,_=model.cosmo(cache.sample(group),data.labels[group]);total+=float(loss)*len(group)
                values=z.double().cpu().numpy();zs+=values.sum(0);zss+=np.square(values).sum(0);count+=len(group)
            assert count==len(ids)
            result[name+'_rows']=count
            result[name+'_loss']=total/(atoms if name=='pcqm' else count)
            result[name+'_descriptor_channel_std']=float(np.sqrt(np.maximum(zss/count-(zs/count)**2,0)).mean())
            if name=='pcqm':result['pcqm_target_batch_std']=std/atoms;result['pcqm_loss_per_channel']=result['pcqm_loss']/1088
        assert all(math.isfinite(v) for v in result.values()),result
        return result
    finally:restore_rng(old)


def checkpoint(point,model,optimizer,scheduler,global_step,metrics):
    path=OUT/'checkpoints'/f'probe_{point:03d}';pending=path.with_name(path.name+'.partial')
    if path.exists():raise RuntimeError(f'Checkpoint already exists: {path}')
    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(point=point,global_step=global_step,optimizer=optimizer.state_dict(),scheduler=scheduler.state_dict(),
        config=CONFIG,**rng_state()),pending/'training_state.pt')
    write(pending/'metrics.json',metrics);write(pending/'config.json',CONFIG)
    for source,name in [(OUT/'architecture.yaml','post_training_architecture_config.yaml'),
        (DATA/'normalization.pt','normalization.pt'),(DATA/'cosmo_label_transform.json','cosmo_label_transform.json'),
        (OUT/'preparation_complete.json','data_manifest.json')]:shutil.copyfile(source,pending/name)
    write(pending/'manifest.json',dict(point=point,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'}))
    pending.rename(path);return path


def upload(path):
    if (path/'hub_verified.json').exists():return json.loads((path/'hub_verified.json').read_text())
    api=HfApi();api.create_repo(MODEL_REPO,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'Save exact-resume scratch checkpoint {path.name} before its BigSolDB probe')
            entries={p.rfilename:p for p in api.model_info(MODEL_REPO,revision=commit.oid,files_metadata=True).siblings}
            for p in files:
                if entries[path.name+'/'+p.name].lfs:assert entries[path.name+'/'+p.name].lfs.sha256==sha(p)
            result=dict(verified=True,repository=MODEL_REPO,revision=commit.oid,point=int(path.name.split('_')[-1]))
            write(path/'hub_verified.json',result);return result
        except Exception:
            if attempt==4:raise
            time.sleep(min(10*2**attempt,50))


def scheduler_for(optimizer,data):
    return torch.optim.lr_scheduler.OneCycleLR(optimizer,max_lr=[CONFIG['learning_rate_max'],CONFIG['learning_rate_uncertainty']],total_steps=data.total_steps)


def restore(path,model,optimizer,scheduler):
    manifest=json.loads((path/'manifest.json').read_text())
    assert sha(path/'training_state.pt')==manifest['files']['training_state.pt']['sha256']
    state=torch.load(path/'training_state.pt',map_location='cpu',weights_only=False)
    assert state['config']==CONFIG
    optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']);restore_rng(state)
    return state


def smoke(data):
    import copy
    seed();loads=[];original_load=torch.load
    def record_load(path,*args,**kwargs):loads.append(str(path));return original_load(path,*args,**kwargs)
    torch.load=record_load
    try:model=ScratchJoint()
    finally:torch.load=original_load
    assert loads==[str(DATA/'normalization.pt')],loads
    opt=optimizer_for(model);scheduler=scheduler_for(opt,data)
    state=copy.deepcopy(model.state_dict());seeded_rng=rng_state()
    pp,cp,_=data.epoch_plans[1];model.train()
    update(model,opt,data,pp[0],cp[0]);scheduler.step()
    first_gradients={name:any(p.grad is not None and p.grad.abs().sum()>0 for p in module.parameters()) for name,module in model.parts.items()}
    # Author pair2bias weights start at zero: geometry receives its first gradient
    # after the first optimizer update has opened this attention path.
    gates=[v for k,v in state.items() if k.endswith('pair2bias.weight')]
    assert gates and all(torch.count_nonzero(v)==0 for v in gates)
    assert all(value for name,value in first_gradients.items() if name!='geometric_preprocessor.pth')
    assert not torch.equal(model.encoder.state_dict()[next(iter(model.encoder.state_dict()))],state['encoder.'+next(iter(model.encoder.state_dict()))])
    # Save after one step, then require an exact next-step replay including optimizer and RNG.
    path=OUT/'checkpoints/probe_999';path.parent.mkdir(exist_ok=True)
    if path.exists():shutil.rmtree(path)
    saved=checkpoint(999,model,opt,scheduler,1,dict(smoke=True))
    expected=update(model,opt,data,pp[1],cp[1]);scheduler.step()
    gradients={name:any(p.grad is not None and p.grad.abs().sum()>0 for p in module.parameters()) for name,module in model.parts.items()}
    assert all(gradients.values()),gradients
    after={k:v.detach().cpu().clone() for k,v in model.state_dict().items()}
    del model,opt,scheduler;gc.collect();torch.cuda.empty_cache()
    replay=ScratchJoint(saved);opt=optimizer_for(replay);scheduler=scheduler_for(opt,data);restore(saved,replay,opt,scheduler)
    replay.train();actual=update(replay,opt,data,pp[1],cp[1]);scheduler.step()
    assert expected==actual and all(torch.equal(v.cpu(),after[k]) for k,v in replay.state_dict().items())
    del replay,opt,scheduler,state,after;gc.collect();torch.cuda.empty_cache();shutil.rmtree(saved)
    seed();model=ScratchJoint();opt=optimizer_for(model);scheduler=scheduler_for(opt,data)
    initial=validate(model,data)
    path=OUT/'checkpoints/probe_000'
    if path.exists():
        existing=ScratchJoint(path)
        assert all(torch.equal(v,existing.state_dict()[k]) for k,v in model.state_dict().items())
        del existing;gc.collect();torch.cuda.empty_cache()
    else:path=checkpoint(0,model,opt,scheduler,0,dict(point=0,epoch_position=0,validation=initial))
    upload(path)
    write(OUT/'smoke_validation.json',dict(passed=True,all_trainable_head_gradients=gradients,
        first_step_gradients=first_gradients,author_zero_initialized_pair_bias_preserved=True,
        initialization_loaded_only_training_statistics=loads,no_pretrained_rem3di_weights_loaded=True,
        exact_optimizer_rng_next_step_replay=True,initial_validation=initial))
    source_names=['scratch_prepare.py','scratch_data.py','scratch_model.py','scratch_train.py','scratch_probe.py','scratch_driver.py','scratch_cache_initial.py','scratch_fast_cache.py',
        'eval_common.py','train_cosmo.py','multitask_data.py','bigsoldb_evaluate.py']
    write(OUT/'provenance.json',dict(initial_rem3di_checkpoint=None,initialization='Constructor initialization of all learned components',
        architecture_sha256=sha(OUT/'architecture.yaml'),
        foundation_sha256=sha(Path('/marimo/rem3di_training/MACE-POLAR-1-M.model')),
        python=platform.python_version(),cuda=torch.version.cuda,gpu=torch.cuda.get_device_name(0),
        versions={n:importlib.metadata.version(n) for n in ['torch','numpy','pandas','scikit-learn','rdkit','structured-data-models','mace-torch','e3nn']},
        source_sha256={n:sha(Path(__file__).with_name(n)) for n in source_names},
        legacy_source_sha256={str(p.relative_to(AUDIT/'author_legacy_source')):sha(p)
            for p in sorted((AUDIT/'author_legacy_source/threedscriptors').rglob('*.py'))}))
    ops=[CommitOperationAdd(path_in_repo='source/'+n,path_or_fileobj=str(Path(__file__).with_name(n))) for n in source_names]
    ops += [CommitOperationAdd(path_in_repo=n,path_or_fileobj=str(OUT/n)) for n in ['smoke_validation.json','config.json','preparation_complete.json','probe_schedule.json','provenance.json']]
    HfApi().create_commit(repo_id=MODEL_REPO,operations=ops,commit_message='Verify random initialization, real-data gradients and exact next-step resume before scratch training')


def train_point(point,data):
    import wandb
    path=OUT/'checkpoints'/f'probe_{point:03d}'
    if path.exists():upload(path);return
    previous=OUT/'checkpoints'/f'probe_{point-1:03d}'
    seed();model=ScratchJoint(previous);opt=optimizer_for(model);scheduler=scheduler_for(opt,data)
    state=restore(previous,model,opt,scheduler);global_step=state['global_step']
    info=data.schedule[point-1];pp,cp,n=data.epoch_plans[info['epoch']]
    assert state['point']==point-1
    assert global_step==info['global_step']-(info['batch_end']-info['batch_start'])
    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',name='scratch-pcqm-cosmo-ten-epochs',
        group='rem3di-scratch-pcqm-cosmo',job_type='scratch-pretraining',config=CONFIG,dir=str(OUT),settings=wandb.Settings(init_timeout=90))
    write(wb_path,dict(run_id=run.id,url=run.url));restore_rng(state)
    model.train();logs=[];start=time.perf_counter()
    try:
        for step in range(info['batch_start'],info['batch_end']):
            log=update(model,opt,data,pp[step%len(pp)],cp[step%len(cp)]);scheduler.step();global_step+=1;logs.append(log)
            if (step+1)%20==0 or step+1==info['batch_end']:
                status('training',point=point,epoch=info['epoch'],batch=step+1,batches=n,wandb_url=run.url)
                run.log(dict(optimizer_step=global_step,epoch_position=info['epoch']-1+(step+1)/n,**{'train/'+k:v for k,v in log.items()},lr=scheduler.get_last_lr()[0]))
        assert global_step==info['global_step']
        validation=validate(model,data)
        metrics={**info,'global_step':global_step,'train':{k:float(np.mean([r[k] for r in logs])) for k in logs[0]},
                 'validation':validation,'seconds':time.perf_counter()-start,
                 'epoch_coverage':dict(pcqm_unique=len(data.pt),cosmo_unique=len(data.ct)) if info['fifth']==5 else None}
        saved=checkpoint(point,model,opt,scheduler,global_step,metrics)
        status('uploading_checkpoint',point=point,epoch=info['epoch'])
        hub=upload(saved);write(OUT/'latest_point.json',dict(**metrics,hub=hub))
        run.log(dict(point=point,epoch_position=info['epoch_position'],**{'validation/'+k:v for k,v in validation.items()},hub_revision=hub['revision']))
        run.finish()
    except Exception:run.finish(exit_code=1);raise


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--point',type=int);parser.add_argument('--smoke',action='store_true');args=parser.parse_args()
    lock=(OUT/'training.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    assert json.loads((OUT/'preparation_complete.json').read_text())['verified']
    seed();data=TrainingData();write(OUT/'config.json',CONFIG)
    if args.smoke:
        if not (OUT/'smoke_validation.json').exists():smoke(data)
    else:
        assert json.loads((OUT/'smoke_validation.json').read_text())['passed']
        assert 1<=args.point<=50;train_point(args.point,data)


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