"""Training-only identity splits, fresh calibration, and frozen MACE caches."""
from pathlib import Path
import sys
sys.path.insert(0,str(Path(__file__).resolve().parent))
import multitask_data as shared
from eval_common import *
import pandas as pd
import concurrent.futures,multiprocessing as mp,collections,fcntl
from huggingface_hub import HfApi,CommitOperationAdd

OUT=WORK/'scratch_pcqm_cosmo';DATA=OUT/'data'
MODEL_REPO='Queimo/rem3di-scratch-pcqm-cosmo'
DATA_REPO='Queimo/rem3di-scratch-pcqm-cosmo-data'
shared.OUT=OUT;shared.DATA=DATA
RawCache=shared.RawCache


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


def identity(item):
    smiles,numbers,charge=item
    from rdkit import RDLogger
    RDLogger.DisableLog('rdApp.warning')
    mol=Chem.MolFromSmiles(smiles)
    if mol is None:return None,'invalid_smiles'
    m=Chem.AddHs(mol)
    if Chem.GetFormalCharge(m)!=charge:return None,'source_smiles_geometry_charge_mismatch'
    if collections.Counter(a.GetAtomicNum() for a in m.GetAtoms())!=collections.Counter(numbers):
        return None,'source_smiles_geometry_composition_mismatch'
    key=Chem.MolToInchiKey(mol)
    if not key:return None,'missing_connectivity_identity'
    return key.split('-')[0],''


def pcqm_frame():
    path=DATA/'pcqm.parquet'
    if path.exists():return pd.read_parquet(path)
    from remedi.data_handling.dataset.molecule_dataset import MoleculeDataset
    ds=MoleculeDataset.open_existing_dataset_from_dir(Path('/marimo/rem3di_training/pcqm500k_std'))
    ptr=np.asarray(ds.ptr[:]);positions=np.asarray(ds.positions[:],dtype=np.float32)
    numbers=np.asarray(ds.atomic_numbers[:],dtype=np.int16)
    # Isomeric strings are indexed by stereoisomer ID, not connectivity/molecule ID.
    strings=np.asarray(ds.isomeric_smiles.to_list(),dtype=object)
    smiles=strings[np.asarray(ds.isomer_ids[:],dtype=np.int64)].tolist()
    assert len(ptr)==500001 and len(smiles)==500000 and ptr[-1]==len(numbers)==len(positions)
    assert np.isfinite(positions).all()
    charges=np.asarray(ds.total_charge[:]);mult=np.asarray(ds.multiplicity[:])
    nums=[numbers[a:b].tolist() for a,b in zip(ptr[:-1],ptr[1:])]
    with concurrent.futures.ProcessPoolExecutor(8,mp_context=mp.get_context('fork')) as pool:
        identities=list(pool.map(identity,zip(smiles,nums,charges),chunksize=1000))
    groups=[v[0] for v in identities];issues=[v[1] for v in identities]
    frame=pd.DataFrame(dict(source_row=np.arange(500000),smiles=smiles,molecule_group=groups,
        atomic_numbers=nums,positions=[positions[a:b].tolist() for a,b in zip(ptr[:-1],ptr[1:])],
        charge=charges,multiplicity=mult,identity_issue=issues,identity_valid=[not v for v in issues]))
    frame.to_parquet(path,index=False);return frame


def prepare_splits():
    marker=OUT/'split_validation.json'
    if marker.exists():return
    status('validating_pcqm_identities_and_splits')
    pcqm=pcqm_frame();cosmo=pd.concat([pd.read_parquet(p) for p in sorted(DATA.glob('cosmo-[0-9]*.parquet'))],ignore_index=True)
    if 'molecule_group' not in cosmo:
        cosmo['molecule_group']=[shared.connectivity(s) for s in cosmo.smiles]
    test=pd.read_csv(WORK/'bigsoldb_evaluation/data/test.csv')
    protected={shared.connectivity(s) for s in test.SMILES_Solute.unique()}
    # Keep original seed-zero 98/2 PCQM partition, expanding within connectivity groups.
    val_rows=np.random.default_rng(0).permutation(len(pcqm))[:10000]
    pval=set(pcqm.molecule_group.iloc[val_rows]);cval=set(cosmo.loc[cosmo.joint_split=='validation','molecule_group'])
    pcqm['split']='train'
    pcqm.loc[pcqm.molecule_group.isin(pval),'split']='validation'
    pcqm.loc[(pcqm.split=='train')&pcqm.molecule_group.isin(cval),'split']='cross_objective_holdout'
    pcqm.loc[pcqm.molecule_group.isin(protected),'split']='benchmark_overlap'
    pcqm.loc[~pcqm.identity_valid,'split']='identity_quarantine'
    cosmo['split']=cosmo.joint_split
    cosmo.loc[(cosmo.split=='train')&cosmo.molecule_group.isin(pval|cval),'split']='cross_objective_holdout'
    cosmo.loc[cosmo.molecule_group.isin(protected),'split']='benchmark_overlap'
    ptr=set(pcqm.loc[pcqm.split=='train','molecule_group']);ctr=set(cosmo.loc[cosmo.split=='train','molecule_group'])
    validation=set(pcqm.loc[pcqm.split=='validation','molecule_group'])|set(cosmo.loc[cosmo.split=='validation','molecule_group'])
    assert not (ptr|ctr)&(validation|protected)
    pcqm.to_parquet(DATA/'pcqm.parquet',index=False)
    pcqm.loc[~pcqm.identity_valid,['source_row','smiles','identity_issue']].to_csv(DATA/'pcqm_identity_quarantine.csv',index=False)
    cosmo.to_parquet(DATA/'cosmo.parquet',index=False)
    cols=[f'p_{i:03d}' for i in range(153)]+['area_A2','volume_A3']
    y=cosmo[cols].to_numpy(dtype=np.float64)
    assert np.isfinite(y).all() and (y[:,:153]>=0).all() and (y[:,153:]>0).all()
    transformed=np.c_[np.log1p(y[:,:153]),np.log(y[:,153:])]
    fit=transformed[cosmo.split.to_numpy()=='train']
    mean=fit.mean(0);scale=fit.std(0);scale[scale<1e-9]=1.
    write(DATA/'cosmo_label_transform.json',dict(mean=mean.tolist(),scale=scale.tolist(),
        train_rows=len(fit),training_only=True,transform='log1p(sigma bin areas), ln(A), ln(V), population standardization'))
    write(marker,dict(passed=True,pcqm_rows=len(pcqm),pcqm_splits=pcqm.split.value_counts().to_dict(),
        cosmo_rows=len(cosmo),cosmo_splits=cosmo.split.value_counts().to_dict(),
        benchmark_test_solute_groups=len(protected),no_pretraining_train_benchmark_overlap=True,
        no_cross_objective_train_validation_group_overlap=True,pcqm_composition_validated_all_rows=True,
        pcqm_invalid_identity_rows=int((~pcqm.identity_valid).sum()),all_train_compositions_and_charges_validated=True,
        geometry='Original PCQM SDF coordinates and original COSMO PM6 coordinates; no regeneration',
        initial_rem3di_checkpoint=None,combisolv_training=False))


def calibrate():
    target=DATA/'normalization.pt'
    if target.exists():return
    from e3nn.o3 import Irreps
    from threedscriptors.utils.model_utils import get_invariant_indices
    indices,_=get_invariant_indices(Irreps('512x0e+512x1o+512x0e+512x1o'))
    count=0;total=np.zeros(1024,dtype=np.float64);squares=total.copy();counts={}
    for name in ['pcqm','cosmo']:
        frame=pd.read_parquet(DATA/f'{name}.parquet',columns=['split']);cache=RawCache(name)
        keep=np.repeat(frame.split.to_numpy()=='train',cache.lengths);counts[name]=int(keep.sum())
        for start in range(0,len(keep),8192):
            stop=min(start+8192,len(keep));mask=keep[start:stop]
            if not mask.any():continue
            # Select invariant columns before copying retained rows; avoid reading
            # and copying the unused 3072 equivariant channels from the mmap.
            x=np.asfortranarray(np.asarray(cache.features[start:stop])[:,indices][mask],dtype=np.float64)
            count+=len(x);total+=x.sum(0);squares+=np.square(x).sum(0)
            if start%(8192*100)==0:status('fitting_training_only_normalization',corpus=name,atoms_done=start,atoms_total=len(keep))
    mean=total/count;std=np.sqrt(np.maximum(squares/count-mean**2,1e-18))
    assert np.isfinite(mean).all() and np.isfinite(std).all() and np.all(std>0)
    torch.save(dict(mean=torch.tensor(mean),std=torch.tensor(std),indices=indices,training_only=True),target)
    write(DATA/'normalization.json',dict(training_only=True,atoms=count,atoms_per_corpus=counts,
        invariant_channels=1024,population_std=True,std_floor=1e-9,sha256=sha(target),
        mean_std=float(std.mean()),minimum_std=float(std.min()),maximum_std=float(std.max())))


def publish_geometry():
    marker=OUT/'geometry_hub.json'
    if marker.exists():return json.loads(marker.read_text())
    files=[DATA/n for n in ['pcqm.parquet','cosmo.parquet','bigsoldb_fixed_geometries.parquet','cosmo_label_transform.json','pcqm_identity_quarantine.csv']]
    files += [OUT/'architecture.yaml',OUT/'split_validation.json']
    api=HfApi();api.create_repo(DATA_REPO,repo_type='dataset',exist_ok=True)
    commit=api.create_commit(repo_id=DATA_REPO,repo_type='dataset',
        operations=[CommitOperationAdd(path_in_repo=p.name,path_or_fileobj=str(p)) for p in files],
        commit_message='Back up complete scratch-training structures, identities, splits and COSMO transforms before MACE caching finishes')
    entries={s.rfilename:s for s in api.dataset_info(DATA_REPO,revision=commit.oid,files_metadata=True).siblings}
    for p in files:
        if entries[p.name].lfs:assert entries[p.name].lfs.sha256==sha(p)
    result=dict(verified=True,repository=DATA_REPO,revision=commit.oid,stage='complete_geometries_and_memberships',normalization='pending frozen features')
    write(marker,result);return result


def main():
    DATA.mkdir(parents=True,exist_ok=True)
    if (OUT/'preparation_complete.json').exists():return
    prepare_splits();publish_geometry();shared.status=status
    from scratch_fast_cache import cache_mace_fast
    with (OUT/'feature_cache.lock').open('a') as lock:
        fcntl.flock(lock,fcntl.LOCK_EX)
        for name in ['pcqm','cosmo','bigsoldb']:
            if (OUT/'cache'/name/'complete.json').exists():continue
            frame=pd.read_parquet(DATA/('bigsoldb_fixed_geometries.parquet' if name=='bigsoldb' else name+'.parquet'))
            if name=='bigsoldb':shared.cache_mace(name,frame)
            else:cache_mace_fast(name,frame)
    calibrate()
    card='''---
configs:
- config_name: pcqm
  data_files:
  - split: structures
    path: pcqm.parquet
- config_name: cosmo
  data_files:
  - split: profiles
    path: cosmo.parquet
- config_name: bigsoldb_geometries
  data_files:
  - split: molecules
    path: bigsoldb_fixed_geometries.parquet
---

# Inputs for random-initialized Rem3Di + COSMO training

PCQM contains the first 500,000 structures accepted by the original standardization,
deduplication, element and atom-count filters from the verified OGB PCQM4Mv2 training SDF
(archive MD5 fd72bce606e7ddf36c2a832badeec6ab). SDF coordinates are retained.
Isomeric SMILES use the stored stereoisomer IDs and all compositions and charges are validated.
Source annotation/geometry mismatches are retained with identity_quarantine membership;
their original coordinates are not changed and they never enter training.
COSMO profiles retain their original joined PubChemQC PM6 geometries.

These viewer splits expose complete source tables. The **split column inside each
PCQM/COSMO table** records actual train, validation and quarantine memberships;
filter it before fitting anything. BigSolDB test-solute connectivity groups are excluded
from both training streams, and validation groups are excluded across objectives.
All invariant normalization and COSMO log1p/ln standardization statistics use training rows only.
MACE-POLAR is frozen. No pretrained Rem3Di weights define the new model initialization.

[PCQM4Mv2](https://ogb.stanford.edu/docs/lsc/pcqm4mv2/) ·
[Original COSMO data](https://huggingface.co/datasets/Queimo/pubchemqc-cosmo-pm6) ·
[Model checkpoints](https://huggingface.co/Queimo/rem3di-scratch-pcqm-cosmo).
'''
    (DATA/'README.md').write_text(card)
    files=[DATA/n for n in ['pcqm.parquet','cosmo.parquet','bigsoldb_fixed_geometries.parquet','cosmo_label_transform.json','normalization.pt','normalization.json']]
    files.append(DATA/'README.md')
    files.append(DATA/'pcqm_identity_quarantine.csv')
    files += [OUT/'architecture.yaml',OUT/'split_validation.json']
    api=HfApi();api.create_repo(DATA_REPO,repo_type='dataset',exist_ok=True)
    commit=api.create_commit(repo_id=DATA_REPO,repo_type='dataset',
        operations=[CommitOperationAdd(path_in_repo=p.name,path_or_fileobj=str(p)) for p in files],
        commit_message='Preserve exact scratch-training structures, memberships and training-only transforms')
    entries={s.rfilename:s for s in api.dataset_info(DATA_REPO,revision=commit.oid,files_metadata=True).siblings}
    for p in files:
        if entries[p.name].lfs:assert entries[p.name].lfs.sha256==sha(p)
    write(OUT/'preparation_complete.json',dict(verified=True,repository=DATA_REPO,revision=commit.oid,
        files_sha256={p.name:sha(p) for p in files},**json.loads((OUT/'split_validation.json').read_text())))
    status('complete')


if __name__=='__main__':main()