"""Sequential, matched three-epoch component ablations; original run is immutable."""
from pathlib import Path
import os,sys,json,subprocess,fcntl,shutil,time,traceback
BASE=Path(__file__).resolve().parent.parent
ORIGINAL=BASE/'multitask_combisolv'
ROOT=ORIGINAL/'ablations'
VARIANTS=['no_denoising','no_decorrelation']


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


def make_runtime(variant):
    out=ROOT/variant;source=out/'source';source.mkdir(parents=True,exist_ok=True)
    for name in ['multitask_data.py','multitask_driver.py','eval_common.py','train_cosmo.py','bigsoldb_evaluate.py','openadmet_evaluate.py']:
        target=source/name
        if not target.exists():target.symlink_to(BASE/'source'/name)
    model=(BASE/'source/multitask_model.py').read_text()
    model=model.replace('CONFIG=dict(epochs=5,','CONFIG=dict(epochs=3,scheduler_horizon=5,ablation='+repr(variant)+',denoising_enabled='+repr(variant!='no_denoising')+',')
    if variant=='no_decorrelation':model=model.replace('decorrelation_weight=8.','decorrelation_weight=0.')
    # Preserve identical stochastic forward kernels; detach the diagnostic
    # denoising scalar so its entire graph receives no backward/optimizer update.
    model=model.replace('        return torch.stack',
        '        if not CONFIG["denoising_enabled"]:denoising=denoising.detach()\n        return torch.stack',1)
    train=(BASE/'source/multitask_train.py').read_text()
    train=train.replace("math.pi*e/CONFIG['epochs']","math.pi*e/CONFIG.get('scheduler_horizon',CONFIG['epochs'])")
    if variant=='no_denoising':
        train=train.replace('uncertainty_objective(losses,model.log_variance[:2])','uncertainty_objective(losses[1:],model.log_variance[1:2])')
    train=train.replace('pcqm2-cosmo10-combisolv-three-task-5epochs','ablation-'+variant+'-3epochs')
    train=train.replace("group='rem3di-cosmo-combisolv'","group='rem3di-combisolv-ablations'")
    evaluation=(BASE/'source/multitask_evaluate.py').read_text()
    evaluation=evaluation.replace('rem3di-cosmo10-combisolv-multitask-report','rem3di-combisolv-ablation-'+variant+'-report')
    evaluation=evaluation.replace('pcqm2-cosmo10-combisolv-every-epoch-bigsoldb-kumo-large','ablation-'+variant+'-bigsoldb-kumo-large')
    evaluation=evaluation.replace("group='rem3di-cosmo-combisolv'","group='rem3di-combisolv-ablations'")
    for name,code in [('multitask_model.py',model),('multitask_train.py',train),('multitask_evaluate.py',evaluation)]:
        compile(code,str(source/name),'exec');(source/name).write_text(code)
    for name in ['initialization.json','prepared_data_hub.json','preparation_complete.json','loss_implementation_validation.json','independent_dataset_validation.json']:
        if (ORIGINAL/name).exists() and not (out/name).exists():shutil.copyfile(ORIGINAL/name,out/name)
    return out,source


def configure(variant):
    out,source=make_runtime(variant)
    # Load shared helpers first: their path bootstrap otherwise shadows the
    # generated ablation modules while the first model is being imported.
    sys.path.insert(0,str(BASE/'source'))
    import multitask_data,train_cosmo,bigsoldb_evaluate,openadmet_evaluate
    sys.path.insert(0,str(source))
    import multitask_model as model,multitask_train as train,multitask_evaluate as evaluate
    assert all(Path(module.__file__).parent==source for module in [model,train,evaluate])
    repo='Queimo/rem3di-combisolv-ablation-'+variant.replace('_','-')
    for module in [model,train,evaluate]:module.OUT=out;module.MODEL_REPO=repo
    def status(phase,**kw):
        row=dict(phase=phase,variant=variant,pid=os.getpid(),updated_unix=time.time(),**kw)
        write(out/'status.json',row);write(ROOT/'status.json',row)
        print('ABLATION_STATUS',json.dumps(row),flush=True)
    for module in [model,train,evaluate]:module.status=status
    evaluate.EVAL=out/'evaluation'
    # initialize() returns exact control weights without rewriting its metadata.
    def initial():
        from huggingface_hub import hf_hub_download
        p=BASE/'checkpoints/epoch_10'
        for n,h in model.INITIAL_HASHES.items():assert model.sha(p/n)==h
        return p,Path(hf_hub_download(model.DECODER_REPO,model.DECODER_PATH,revision=model.DECODER_REV))
    model.initialize=initial
    model.CONFIG.update(control_run_repository='Queimo/rem3di-cosmo10-combisolv-multitask',
        active_tasks=['cosmo','combisolv'] if variant=='no_denoising' else ['rem3di_denoising','cosmo','combisolv'],
        control_epochs=[1,2,3],schedule_note='First three epochs of the original five-epoch cosine schedule, matched exactly',
        stochastic_control='Detached diagnostic denoising forward preserves identical shared dropout RNG; decoder receives no gradients or updates and its entire UW term is excluded' if variant=='no_denoising' else 'Same stochastic forwards and task losses as control; covariance coefficient alone changes from 8 to 0')
    return out,model,train,evaluate


def smoke(out,model,train):
    torch=model.torch;np=model.np
    torch.manual_seed(42026);torch.cuda.manual_seed_all(42026);np.random.seed(42026)
    torch.use_deterministic_algorithms(True)
    data=train.TrainingData();net=model.MultiTask();opt=model.optimizer_for(net)
    net.eval()
    with torch.no_grad():
        cached=model.RawCache('bigsoldb');ref=np.load(BASE/'bigsoldb_evaluation/features/cosmo_epoch10.npz')['X'];ours=np.empty_like(ref)
        for ids in model.common.batch_indices(np.arange(len(ref)),cached.lengths,max_size=24,atom_cap=2400):
            _,_,z=net.encode(cached.sample(ids));ours[ids]=z.cpu().numpy()
        assert np.allclose(ours,ref,rtol=2e-4,atol=2e-4)
        initial_error=float(np.abs(ours-ref).max())
    scheduler=torch.optim.lr_scheduler.LambdaLR(opt,lambda e:1.)
    cp,qp,_=data.plans(1);ci=cp[10];qi=qp[10]
    net.train()
    if not model.CONFIG['denoising_enabled']:
        rng=train.rng_state()
        net.cosmo_denoising(data.c_cache.sample(ci),data.c_target[ci]);inactive_rng=train.rng_state()
        train.restore_rng(rng);model.CONFIG['denoising_enabled']=True
        net.cosmo_denoising(data.c_cache.sample(ci),data.c_target[ci]);active_rng=train.rng_state()
        model.CONFIG['denoising_enabled']=False;train.restore_rng(rng)
        assert all(torch.equal(a,b) for a,b in zip(inactive_rng['cuda_rng'],active_rng['cuda_rng'])),'Inactive denoising forward changed shared RNG sequence'
    before=next(net.encoder.parameters()).detach().clone();decoder=next(net.denoiser.parameters()).detach().clone()
    net.train();loss=train.update(net,opt,data,ci,qi)
    assert not torch.equal(before,next(net.encoder.parameters()).detach())
    active=['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth','cosmo_head.pth','combisolv_head.pth']
    if model.CONFIG['denoising_enabled']:active.append('denoising_decoder.pth')
    assert all(any(p.grad is not None and bool((p.grad!=0).any()) for p in net.parts[n].parameters()) for n in active)
    if not model.CONFIG['denoising_enabled']:
        assert all(p.grad is None for p in net.denoiser.parameters())
        assert torch.equal(decoder,next(net.denoiser.parameters()).detach())
        assert float(net.log_variance[0].detach())==-.5 and float(net.log_variance.grad[0])==0.
    temp=out/'smoke_checkpoint';train.save_checkpoint(temp,net,opt,scheduler,0,1,loss)
    expected=train.update(net,opt,data,ci,qi);expected_state={n:v.detach().cpu().clone() for n,v in net.state_dict().items()}
    net.load_parts(temp);state=torch.load(temp/'training_state.pt',map_location='cpu',weights_only=False)
    opt.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']);train.restore_rng(state)
    actual=train.update(net,opt,data,ci,qi)
    assert all(np.isclose(actual[k],expected[k],rtol=1e-7,atol=1e-7) for k in expected)
    assert all(torch.equal(v.detach().cpu(),expected_state[n]) for n,v in net.state_dict().items())
    shutil.rmtree(temp)
    write(out/'smoke_validation.json',dict(passed=True,config=model.CONFIG,active_head_gradients=True,
        initial_bigsoldb_embedding_max_abs_error=initial_error,
        disabled_objective_has_no_gradient=True,optimizer_rng_next_step_replay_exact=True,
        denoising_enabled=model.CONFIG['denoising_enabled'],decorrelation_weight=model.CONFIG['decorrelation_weight'],losses=loss))
    print('ABLATION_SMOKE_PASSED',model.CONFIG['ablation'],flush=True)


def report(out,model,evaluate,run):
    import pandas as pd
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    current=pd.read_csv(evaluate.EVAL/'results.csv')
    control=pd.read_csv(ORIGINAL/'evaluation/results.csv').query('epoch <= 3')
    comparison=pd.concat([control.assign(variant='full_model'),current.assign(variant=model.CONFIG['ablation'])],ignore_index=True)
    comparison.to_csv(evaluate.EVAL/'analysis/comparison.csv',index=False)
    fig,axes=plt.subplots(1,2,figsize=(9,3.5),layout='constrained')
    for name,frame in comparison.groupby('variant'):
        for ax,metric in zip(axes,['mae','rmse']):
            ax.plot(frame.epoch,frame[metric],'-o',label=name);ax.set_xlabel('Epoch');ax.set_ylabel(metric.upper());ax.set_xticks([1,2,3])
    axes[0].legend();fig.savefig(evaluate.EVAL/'analysis/epoch_curve.png',dpi=160);plt.close(fig)
    lines=['# Three-epoch component ablation', '',f"Variant: **{model.CONFIG['ablation']}**.", '',
        'Each run starts from the exact evaluated PCQMep2→COSMOep10 checkpoint. Data, geometry, order/shuffle seeds, initial task heads, optimizers and first-three-epoch learning rates match the full control. Training stops after three epochs. BigSolDB uses the identical 9686 test rows, 75731 training contexts, all1281 columns and Kumo Large four estimators, seed42.','',
        'The denoising objective corrupts atomic embeddings, not coordinates. For no_denoising, the decoder receives no gradients or updates; its loss and complete uncertainty term (including log variance regularizer) are excluded. A detached diagnostic forward preserves the original dropout/noise RNG sequence. For no_decorrelation, all three UW task terms remain, and the covariance coefficient alone is zero.','',
        '| Variant | Epoch | MAE | RMSE | R² |','|---|---:|---:|---:|---:|']
    for r in comparison.itertuples():lines.append(f'| {r.variant} | {r.epoch} | {r.mae:.6f} | {r.rmse:.6f} | {r.r2:.6f} |')
    lines+=['','One matched training seed; checkpoint3 is the predeclared final comparison, with all epoch results retained. No test-based hyperparameter changes.','',f'[W&B]({run.url})']
    text='\n'.join(lines)+'\n';(evaluate.EVAL/'analysis/REPORT.md').write_text(text);(out/'README.md').write_text(text)
    return comparison


def child(variant,stage):
    out,model,train,evaluate=configure(variant)
    sys.argv=[sys.argv[0]]
    if stage=='smoke':smoke(out,model,train);return
    if stage=='train':
        from huggingface_hub import HfApi,CommitOperationAdd
        api=HfApi();api.create_repo(model.MODEL_REPO,exist_ok=True)
        sources=list((out/'source').glob('*.py'))+[Path(__file__)]
        api.create_commit(repo_id=model.MODEL_REPO,operations=[CommitOperationAdd(path_in_repo='source/'+p.name,path_or_fileobj=str(p)) for p in sources]+[
            CommitOperationAdd(path_in_repo='smoke_validation.json',path_or_fileobj=str(out/'smoke_validation.json'))],commit_message='Pin matched three-epoch ablation implementation and gradient/resume validation')
        train.main();return
    evaluate.build_report=lambda run:report(out,model,evaluate,run)
    evaluate.main()


def combined():
    import numpy as np,pandas as pd
    from huggingface_hub import HfApi,CommitOperationAdd
    records=[pd.read_csv(ORIGINAL/'evaluation/results.csv').query('epoch <= 3').assign(variant='full_model')]
    records += [pd.read_csv(ROOT/v/'evaluation/results.csv').assign(variant=v) for v in VARIANTS]
    result=pd.concat(records,ignore_index=True);result.to_csv(ROOT/'comparison.csv',index=False)
    test=pd.read_csv(BASE/'bigsoldb_evaluation/data/test.csv')
    groups=test.SMILES_Solute.to_numpy();unique=np.unique(groups);ids=[np.flatnonzero(groups==g) for g in unique]
    reference=np.load(ORIGINAL/'evaluation/predictions/epoch_03.npz');y=reference['y_true'];p0=reference['y_pred']
    intervals={}
    for v in VARIANTS:
        z=np.load(ROOT/v/'evaluation/predictions/epoch_03.npz');assert np.array_equal(z['y_true'],y)
        assert np.array_equal(z['test_clean_rows'],reference['test_clean_rows'])
        p=z['y_pred'];rng=np.random.default_rng(42026);delta=[]
        for _ in range(2000):
            idx=np.concatenate([ids[j] for j in rng.integers(0,len(ids),len(ids))])
            delta.append([np.abs(y[idx]-p[idx]).mean()-np.abs(y[idx]-p0[idx]).mean(),
                np.sqrt(np.mean((y[idx]-p[idx])**2))-np.sqrt(np.mean((y[idx]-p0[idx])**2))])
        intervals[v]=dict(mae_delta_95pct=np.quantile(np.asarray(delta)[:,0],[.025,.975]).tolist(),rmse_delta_95pct=np.quantile(np.asarray(delta)[:,1],[.025,.975]).tolist())
    write(ROOT/'paired_solute_bootstrap.json',dict(seed=42026,replicates=2000,test_solutes=len(unique),difference='ablation minus full control at epoch3; negative favors removal',intervals=intervals))
    lines=['# Sequential three-epoch component ablations','',
        'Full-model control versus independent removals of embedding denoising and covariance decorrelation. Same PCQMep2→COSMOep10 start, data, first-three-epoch learning rates, seeds and BigSolDB Kumo settings. Each ablation runs for three full epochs; neither starts from the other ablation.','',
        '| Variant | Epoch | MAE | RMSE | R² |','|---|---:|---:|---:|---:|']
    for r in result.itertuples():lines.append(f'| {r.variant} | {r.epoch} | {r.mae:.6f} | {r.rmse:.6f} | {r.r2:.6f} |')
    lines+=['','Final comparison is epoch3, predeclared. Negative paired differences favor removal. Intervals resample the126 held-out solutes (2000 paired bootstrap samples); they measure test-cohort uncertainty, not training-seed variability.','',json.dumps(intervals,indent=2)]
    (ROOT/'REPORT.md').write_text('\n'.join(lines)+'\n')
    api=HfApi();c=api.create_commit(repo_id='Queimo/rem3di-cosmo10-combisolv-multitask',operations=[CommitOperationAdd(path_in_repo='ablations/'+n,path_or_fileobj=str(ROOT/n)) for n in ['REPORT.md','comparison.csv','paired_solute_bootstrap.json']]+[CommitOperationAdd(path_in_repo='source/ablation_runner.py',path_or_fileobj=str(Path(__file__)))],commit_message='Add sequential matched three-epoch denoising and decorrelation ablations')
    write(ROOT/'complete.json',dict(completed=True,variants=VARIANTS,epochs_each=3,probes_each=3,revision=c.oid,report='https://huggingface.co/Queimo/rem3di-cosmo10-combisolv-multitask/blob/'+c.oid+'/ablations/REPORT.md'))


def main():
    ROOT.mkdir(exist_ok=True);lock=(ROOT/'driver.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    if (ROOT/'complete.json').exists():return
    assert json.loads((ORIGINAL/'complete.json').read_text())['completed']
    for v in VARIANTS:
        for stage,marker in [('smoke','smoke_validation.json'),('train','training_complete.json'),('evaluate','complete.json')]:
            if (ROOT/v/marker).exists():continue
            write(ROOT/'status.json',dict(phase=stage,variant=v,pid=os.getpid(),updated_unix=time.time()))
            proc=subprocess.Popen([sys.executable,'-u',str(Path(__file__)),v,stage],env={**os.environ,'TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD':'1','CUBLAS_WORKSPACE_CONFIG':':4096:8'})
            write(ROOT/'active_process.json',dict(pid=proc.pid,variant=v,stage=stage))
            if proc.wait():raise RuntimeError(f'{v} {stage} failed; see ablation driver.log')
    combined();write(ROOT/'status.json',dict(phase='complete',updated_unix=time.time()))


if __name__=='__main__':
    try:
        if len(sys.argv)==3:child(sys.argv[1],sys.argv[2])
        else:main()
    except Exception:
        write(ROOT/'failure.json',dict(error=traceback.format_exc(),time=time.time()));raise