"""Matched COSMO epoch 10 / Monroe + Kumo Large on official OpenADMET splits."""
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).resolve().parent))
import eval_common as common
from eval_common import WORK, sha, write
import concurrent.futures
import fcntl
import gc
import hashlib
import importlib.metadata
import json
import multiprocessing as mp
import os
import shutil
import subprocess
import time
import traceback

import numpy as np
import pandas as pd
import torch
import wandb
from huggingface_hub import HfApi, hf_hub_download, CommitOperationAdd
from rdkit import Chem

OUT = WORK/'openadmet_evaluation'
REPOSITORY = 'Queimo/rem3di-pcqm-epoch2-cosmo'
DATASET = 'openadmet/openadmet-expansionrx-challenge-data'
DATA_REVISION = '6b898ccc43d10d25b230fb09e22a6e30c30022b5'
MONROE = WORK/'external/Monroe'
MONROE_REVISION = 'de075eabe98e4adbbedbc86ca9e877f4bfd9b629'
MONROE_SHA = 'df996e8e98b12a3cda5fe5acbef3e1e9a4ea31946d3736461cd07e53a5bf205a'
SEED = 42
PROTOCOL = dict(
    dataset=DATASET, dataset_revision=DATA_REVISION,
    reference_notebook=f'https://github.com/blazejba/Monroe/blob/{MONROE_REVISION}/examples/openadmet.ipynb',
    split='Original organizer train/test CSVs: 5326/2282; no new CV or test-based selection',
    representations=['cosmo_epoch10','monroe'], methods=['kumo_large'], expected_endpoint_results=18,
    labels='Raw numeric endpoint values, matching the reference notebook; only missing measurements excluded per endpoint',
    geometry='One author Monroe ETKDG/MMFF94s/UFF conformer per unique molecule, cached and shared exactly by both encoders',
    encoder='Frozen released Monroe encoder (720 dimensions) and frozen COSMO epoch 10 backbone (640 dimensions); strict state loading',
    monroe_revision=MONROE_REVISION, monroe_weights_sha256=MONROE_SHA,
    kumo=dict(size='large',num_estimators=4,seed=SEED,max_columns=640,estimator_batch_size=1,
        precision='float16 autocast',query_chunk_size=256,
        preprocessing='Default recipe fitted only on measured training rows; processed context shapes recorded'),
    metrics=['rmse','mae','r2','spearman'], comparison='Same molecules, conformers, endpoint masks, seeds and Kumo recipe; only representation differs')


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


def prepare_data():
    status('preparing_official_dataset')
    parts = []
    for split in ['train','test']:
        name = f'expansion_data_{split}.csv'
        source = Path(hf_hub_download(DATASET,name,repo_type='dataset',revision=DATA_REVISION))
        target = OUT/'data'/name
        if not target.exists(): shutil.copyfile(source,target)
        assert sha(source)==sha(target)
        parts.append(pd.read_csv(target))
    train,test = parts
    assert len(train)==5326 and len(test)==2282
    endpoints = [c for c in train if c not in ['Molecule Name','SMILES']]
    assert len(endpoints)==9 and list(train.columns)==list(test.columns)
    smiles = [str(s) for s in pd.unique(pd.concat([train.SMILES,test.SMILES]))]
    assert len(smiles)==7608
    canonical = [{Chem.MolToSmiles(Chem.MolFromSmiles(s),isomericSmiles=True) for s in df.SMILES} for df in parts]
    assert not canonical[0]&canonical[1]
    assert all(np.isfinite(df[e].dropna().to_numpy(dtype=float)).all() for df in parts for e in endpoints)
    PROTOCOL['endpoints'] = endpoints
    PROTOCOL['dataset_files_sha256'] = {p.name:sha(p) for p in (OUT/'data').glob('*.csv')}
    PROTOCOL['checkpoint_sha256'] = {n:sha(common.checkpoint(10)/n) for n in
        ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth']}
    manifest=json.loads((common.checkpoint(10)/'manifest.json').read_text())
    for name,digest in PROTOCOL['checkpoint_sha256'].items():assert digest==manifest['files'][name]['sha256']
    write(OUT/'protocol.json',PROTOCOL)
    write(OUT/'data/measurement_counts.json', {e:dict(train=int(train[e].notna().sum()),test=int(test[e].notna().sum())) for e in endpoints})
    return train,test,endpoints,smiles


def restore_monroe():
    env={**os.environ,'GIT_LFS_SKIP_SMUDGE':'1'}
    MONROE.parent.mkdir(exist_ok=True)
    if not (MONROE/'.git').exists():
        subprocess.run(['git','clone','--depth','1','https://github.com/blazejba/Monroe.git',str(MONROE)],check=True,env=env)
    revision=subprocess.check_output(['git','-C',str(MONROE),'rev-parse','HEAD'],text=True).strip()
    if revision!=MONROE_REVISION:
        subprocess.run(['git','-C',str(MONROE),'fetch','--depth','1','origin',MONROE_REVISION],check=True,env=env)
        subprocess.run(['git','-C',str(MONROE),'checkout',MONROE_REVISION],check=True,env=env)
    weights=MONROE/'checkpoint/weights.pt'
    if not weights.exists() or sha(weights)!=MONROE_SHA:
        temp=weights.with_suffix('.download')
        subprocess.run(['curl','-fL','--retry','4','-o',str(temp),f'https://media.githubusercontent.com/media/blazejba/Monroe/{MONROE_REVISION}/checkpoint/weights.pt'],check=True)
        assert sha(temp)==MONROE_SHA
        temp.replace(weights)
    sys.path.insert(0,str(MONROE))


def graph_job(item):
    index,smiles=item
    from monroe.eval.embed import _featurize_one
    from rdkit import rdBase
    rdBase.SeedRandomNumberGenerator(SEED+index)
    return _featurize_one(smiles)


def prepare_graphs(smiles):
    import pyarrow as pa
    import pyarrow.parquet as pq
    path=OUT/'features/shared_graphs.parquet'
    if path.exists():
        records=pq.read_table(path).to_pylist()
        assert [r['smiles'] for r in records]==smiles
        return [(r.pop('smiles'),r) for r in records]
    status('generating_shared_conformers',total=len(smiles))
    with concurrent.futures.ProcessPoolExecutor(max_workers=8,mp_context=mp.get_context('fork')) as pool:
        built=list(pool.map(graph_job,enumerate(smiles),chunksize=8))
    missing=[smiles[i] for i,v in enumerate(built) if v is None]
    write(OUT/'data/coverage.json',dict(requested=len(smiles),missing=missing))
    assert not missing, f'Do not silently shrink the official split: {missing[:8]}'
    assert [s for s,_ in built]==smiles
    records=[dict(smiles=s,**{k:np.asarray(g[k]).tolist() for k in
        ['node_float','node_codes','edge_index','edge_codes','pos_rdkit']}) for s,g in built]
    pq.write_table(pa.Table.from_pylist(records),path,compression='zstd')
    write(OUT/'features/shared_geometry_manifest.json',dict(rows=len(smiles),sha256=sha(path),
        all_z_zero_layouts=sum(bool(np.all(np.asarray(g['pos_rdkit'])[:,2]==0)) for _,g in built),
        source_revision=MONROE_REVISION))
    return built


def monroe_embeddings(built,smiles):
    path=OUT/'features/monroe.npz'
    if path.exists():
        z=np.load(path);assert z['smiles'].tolist()==smiles;return z['X']
    status('embedding_monroe',total=len(built))
    from monroe.eval.embed import to_pyg
    from monroe.model.ckpt import _build_encoder,_load_state_dict
    from torch_geometric.data import Batch
    hp=json.loads((MONROE/'checkpoint/config.json').read_text())
    encoder=_build_encoder(hp).eval()
    state=_load_state_dict(MONROE/'checkpoint/weights.pt',device='cpu')
    encoder.load_state_dict({k[len('encoder.'):]:v for k,v in state.items() if k.startswith('encoder.')},strict=True)
    del state
    encoder=encoder.cuda().requires_grad_(False)
    X=np.empty((len(built),720),dtype=np.float32)
    with torch.inference_mode():
        for start in range(0,len(built),8):
            batch=Batch.from_data_list([to_pyg(g) for _,g in built[start:start+8]]).to('cuda')
            emb,_=encoder(batch);X[start:start+len(emb)]=emb.float().cpu().numpy()
            if start%512==0:status('embedding_monroe',completed=start+len(emb),total=len(built))
    assert np.isfinite(X).all()
    np.savez_compressed(path,X=X,smiles=np.asarray(smiles),weights_sha256=MONROE_SHA)
    del encoder;gc.collect();torch.cuda.empty_cache()
    return X


def cosmo_embeddings(built,smiles):
    path=OUT/'features/cosmo_epoch10.npz'
    if path.exists():
        z=np.load(path);assert z['smiles'].tolist()==smiles
        assert json.loads(str(z['checkpoint_sha256']))==PROTOCOL['checkpoint_sha256']
        return z['X']
    from collections import Counter
    from monroe.model.constants import NODE_FEAT_LIST_ONE_HOT
    from remedi.data_handling.sample import Sample,yield_molecules_collate_fn
    from build_author_mace_runtime import build_author_mace_runtime
    from threedscriptors.data_handling.sample import Sample as LegacySample
    columns=list(NODE_FEAT_LIST_ONE_HOT);j=columns.index('atomic_number')
    vocab=NODE_FEAT_LIST_ONE_HOT['atomic_number'];ptable=Chem.GetPeriodicTable()
    atoms=[];charge_counts={}
    for s,g in built:
        codes=np.asarray(g['node_codes'])[:,j]
        assert (codes<len(vocab)).all()
        z=np.asarray([ptable.GetAtomicNumber(vocab[int(code)]) for code in codes],dtype=np.int64)
        pos=np.asarray(g['pos_rdkit'],dtype=np.float32)
        charge=int(round(float(np.asarray(g['node_float'])[:,5].sum())))
        original=Chem.AddHs(Chem.MolFromSmiles(s))
        assert Counter(z.tolist())==Counter(a.GetAtomicNum() for a in original.GetAtoms()),s
        assert charge==Chem.GetFormalCharge(original),s
        assert int(np.asarray(g['node_float'])[:,0].sum())==0 and (int(z.sum())-charge)%2==0,s
        assert pos.shape==(len(z),3) and np.isfinite(pos).all()
        atoms.append((torch.from_numpy(z),torch.from_numpy(pos),charge))
        charge_counts[str(charge)]=charge_counts.get(str(charge),0)+1
    status('embedding_cosmo_epoch10',completed=0,total=len(built))
    pre,encoder=common.load_backbone(10)
    pre.eval().requires_grad_(False);encoder.eval().requires_grad_(False)
    mace=build_author_mace_runtime().cuda().eval().requires_grad_(False)
    captures=[]
    hook=mace.preprocessor.atomic_preprocessor.register_forward_pre_hook(lambda module,args:captures.append(args[0].detach()))
    lengths=[len(a[0]) for a in atoms]
    X=np.empty((len(built),640),dtype=np.float32);completed=0
    with torch.inference_mode():
        for ids in common.batch_indices(np.arange(len(built)),lengths,max_size=24,atom_cap=2400):
            samples=[Sample(atomic_numbers=atoms[i][0],atomic_positions=atoms[i][1],total_charge=float(atoms[i][2]),multiplicity=1.) for i in ids]
            batch=yield_molecules_collate_fn(samples).to_('cuda')
            captures.clear();mace.preprocessor(batch)
            assert len(captures)==1
            result=encoder(pre(LegacySample(embeddings=captures.pop().float(),
                atomic_positions=batch.atomic_positions,padding_mask=batch.padding_mask)))
            if not torch.is_tensor(result):result=result.flat
            X[ids]=result.float().cpu().numpy();completed+=len(ids)
            if completed%512<24:status('embedding_cosmo_epoch10',completed=completed,total=len(built))
    hook.remove();assert np.isfinite(X).all()
    np.savez_compressed(path,X=X,smiles=np.asarray(smiles),checkpoint_sha256=json.dumps(PROTOCOL['checkpoint_sha256']))
    write(OUT/'features/cosmo_input_validation.json',dict(passed=True,rows=len(built),
        atomic_compositions_validated=True,formal_charges_validated=True,charge_counts=charge_counts,
        geometry_shared_with_monroe=True,encoder_frozen=True))
    del pre,encoder,mace,atoms;gc.collect();torch.cuda.empty_cache()
    return X


def metrics(y,p):
    from scipy.stats import spearmanr
    return dict(**common.metric(y,p),spearman=float(spearmanr(y,p).statistic))


def fit_kumo(model,X,y,tr,te):
    import sdm
    import sdm.processing as sp
    from sdm.models.callback import Callback
    def table(x):return sdm.TableTensor.from_tensor(torch.as_tensor(np.ascontiguousarray(x),dtype=torch.float32,device='cuda'))
    recipe=model.default_recipe()
    selectors=[m for m in recipe.features.modules() if isinstance(m,sp.SelectColumns)]
    assert len(selectors)==1;selectors[0].max_columns=640
    shapes=[]
    class Audit(Callback):
        def on_context_preprocessing_end(self,model,x,y,related_tables):
            shapes.append(list(x.shape));return x,y,related_tables
    values=[]
    with torch.inference_mode(),torch.autocast('cuda',dtype=torch.float16):
        model.fit(x=table(X[tr]),y=table(y[tr,None]),num_estimators=4,estimator_batch_size=1,
            generator=torch.Generator(device='cuda').manual_seed(SEED),recipe=recipe,callbacks=[Audit()])
        for start in range(0,len(te),256):
            q=model.predict(table(X[te[start:start+256]])).numerical.float().cpu().numpy()
            assert q.shape[0]==len(te[start:start+256])
            values.append(q.mean(1,dtype=np.float64))
    pred=np.concatenate(values);assert np.isfinite(pred).all()
    assert shapes and max(s[-1] for s in shapes)<=640
    model.clear();torch.cuda.empty_cache()
    return pred,dict(processed_context_shapes=shapes,num_estimators=4,raw_columns=X.shape[1])


def run_endpoints(run,features,train,test,endpoints,smiles):
    import sdm
    path=OUT/'results.csv'
    records=pd.read_csv(path).to_dict('records') if path.exists() else []
    lookup={s:i for i,s in enumerate(smiles)}
    model=sdm.models.KumoTabular(task='regression',size='large',device='cuda')
    for representation in ['cosmo_epoch10','monroe']:
        X=features[representation]
        for ei,endpoint in enumerate(endpoints):
            if any(r['representation']==representation and r['endpoint']==endpoint for r in records):continue
            train_rows=np.flatnonzero(train[endpoint].notna().to_numpy())
            test_rows=np.flatnonzero(test[endpoint].notna().to_numpy())
            tr=np.asarray([lookup[s] for s in train.iloc[train_rows].SMILES])
            te=np.asarray([lookup[s] for s in test.iloc[test_rows].SMILES])
            y=np.full(len(smiles),np.nan,dtype=float)
            y[tr]=train.iloc[train_rows][endpoint].to_numpy(dtype=float)
            y[te]=test.iloc[test_rows][endpoint].to_numpy(dtype=float)
            assert not set(tr)&set(te)
            status('predicting_endpoints',representation=representation,endpoint=endpoint,completed=len(records),total=18,n_train=len(tr),n_test=len(te))
            started=time.perf_counter();pred,extras=fit_kumo(model,X,y,tr,te)
            result=dict(representation=representation,endpoint=endpoint,method='kumo_large',seed=SEED,
                n_train=len(tr),n_test=len(te),**metrics(y[te],pred),seconds=time.perf_counter()-started,
                raw_columns=extras['raw_columns'])
            p=OUT/'predictions'/f'{representation}_endpoint_{ei:02d}.npz'
            np.savez_compressed(p,train_csv_rows=train_rows,test_csv_rows=test_rows,
                train_feature_indices=tr,test_feature_indices=te,y_true=y[te],y_pred=pred,
                test_smiles=test.iloc[test_rows].SMILES.to_numpy(dtype=str))
            write(p.with_suffix('.json'),{**result,**extras})
            records.append(result)
            tmp=path.with_suffix('.partial.csv');pd.DataFrame(records).to_csv(tmp,index=False);tmp.replace(path)
            write(OUT/'progress.json',dict(completed=len(records),total=18,latest=result))
            run.log({'endpoint/'+k:v for k,v in result.items()})
            artifact=wandb.Artifact('cosmo-openadmet-endpoint-backups',type='evaluation',metadata=result)
            for file in [p,p.with_suffix('.json'),path,OUT/'progress.json']:
                artifact.add_file(str(file),name=str(file.relative_to(OUT)))
            run.log_artifact(artifact,aliases=['latest',f'endpoint-{len(records):02d}']).wait(timeout=60)
            print('OPENADMET_RESULT',json.dumps(result),flush=True)
    assert len(records)==18
    del model;gc.collect();torch.cuda.empty_cache()


def publish(run):
    artifact=wandb.Artifact('cosmo-openadmet-evaluation-results',type='evaluation',metadata=dict(endpoint_results=18,protocol=PROTOCOL))
    files=[p for p in OUT.rglob('*') if p.is_file() and 'wandb' not in p.relative_to(OUT).parts and p.suffix in ['.json','.csv','.npz','.parquet','.png','.svg','.md']]
    for p in files:artifact.add_file(str(p),name='openadmet_evaluation/'+str(p.relative_to(OUT)))
    sources=[Path(__file__),Path(__file__).with_name('openadmet_report.py')]
    for p in sources:artifact.add_file(str(p),name='source/'+p.name)
    combined=WORK/'expanded_evaluation/analysis/REPORT.md'
    artifact.add_file(str(combined),name='expanded_evaluation/analysis/REPORT.md')
    run.log_artifact(artifact,aliases=['latest','complete']).wait(timeout=60)
    operations=[CommitOperationAdd(path_in_repo='openadmet_evaluation/'+str(p.relative_to(OUT)),path_or_fileobj=str(p)) for p in files]
    operations += [CommitOperationAdd(path_in_repo='source/'+p.name,path_or_fileobj=str(p)) for p in sources]
    operations.append(CommitOperationAdd(path_in_repo='expanded_evaluation/analysis/REPORT.md',path_or_fileobj=str(combined)))
    revision=HfApi().create_commit(repo_id=REPOSITORY,operations=operations,
        commit_message='Add matched COSMO epoch 10 and Monroe Kumo Large results on nine OpenADMET endpoints').oid
    write(OUT/'hub_verified.json',dict(completed=True,repository=REPOSITORY,revision=revision))
    return revision


def main():
    OUT.mkdir(exist_ok=True)
    with (OUT/'driver.lock').open('a') as lock:
        try:fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
        except BlockingIOError:print('OPENADMET_ALREADY_RUNNING',flush=True);return
        for folder in ['data','features','predictions','analysis']:(OUT/folder).mkdir(exist_ok=True)
        PROTOCOL['versions']={n:importlib.metadata.version(n) for n in ['torch','numpy','rdkit','scikit-learn','structured-data-models','torch-geometric','datamol','wandb']}
        train,test,endpoints,smiles=prepare_data()
        restore_monroe()
        built=prepare_graphs(smiles)
        # Featurize without labels before any endpoint model is fitted.
        features={'monroe':monroe_embeddings(built,smiles),'cosmo_epoch10':cosmo_embeddings(built,smiles)}
        del built
        write(OUT/'features/manifest.json',{n:dict(rows=len(X),columns=X.shape[1],sha256=sha(OUT/'features'/f'{n}.npz')) for n,X in features.items()})
        saved=json.loads((OUT/'wandb.json').read_text()) if (OUT/'wandb.json').exists() else {}
        run=wandb.init(entity='queimo',project='rem3di-reproduction',id=saved.get('run_id') or wandb.util.generate_id(),resume='allow',
            name='openadmet-cosmo-epoch10-vs-monroe-kumo-large',group='cosmo-openadmet',job_type='official-split-comparison',
            config=PROTOCOL,dir=str(OUT),settings=wandb.Settings(init_timeout=90))
        write(OUT/'wandb.json',dict(run_id=run.id,url=run.url))
        try:
            run_endpoints(run,features,train,test,endpoints,smiles)
            from openadmet_report import build_report
            build_report(OUT,train,test,endpoints,run)
            write(OUT/'complete.json',dict(completed=True,endpoint_results=18,endpoints=9,representations=2,official_train_rows=5326,official_test_rows=2282,wandb_url=run.url))
            revision=publish(run)
            run.summary.update(dict(all_evaluations_complete=True,endpoint_results=18,hub_revision=revision));run.finish()
            status('complete',endpoint_results=18,hub_revision=revision)
            HfApi().create_commit(repo_id=REPOSITORY,operations=[
                CommitOperationAdd(path_in_repo='openadmet_evaluation/'+p.name,path_or_fileobj=str(p)) for p in [OUT/'status.json',OUT/'hub_verified.json']],
                commit_message='Verify completed OpenADMET evaluation publication')
        except Exception:
            status('failed',error=traceback.format_exc());run.finish(exit_code=1);raise


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