"""Extract frozen MACE inputs before unused learned preprocessing runs.

The batch plan, MACE call, coordinates and cache layout match multitask_data.
Raising after its existing atomic-input capture avoids downstream preprocessing
that cannot affect the captured raw input. Existing batch-progress files resume.
"""
from pathlib import Path
import sys,gc
sys.path.insert(0,str(Path(__file__).resolve().parent))
from scratch_data import *
common=shared.common


def cache_mace_fast(name,frame):
    root=OUT/'cache'/name;root.mkdir(parents=True,exist_ok=True)
    if (root/'complete.json').exists():return
    from remedi.data_handling.sample import Sample,yield_molecules_collate_fn
    from build_author_mace_runtime import build_author_mace_runtime
    from train_cosmo import batches
    numbers=[torch.tensor(np.asarray(z),dtype=torch.long) for z in frame.atomic_numbers]
    positions=[torch.tensor(np.stack(p),dtype=torch.float32) for p in frame.positions]
    lengths=np.asarray([len(z) for z in numbers]);ptr=np.r_[0,np.cumsum(lengths)]
    np.save(root/'offsets.npy',ptr)
    pos=np.lib.format.open_memmap(root/'positions.npy',mode='w+',dtype=np.float32,shape=(int(ptr[-1]),3))
    for i,p in enumerate(positions):pos[ptr[i]:ptr[i+1]]=p.numpy()
    pos.flush();del pos
    mode='r+' if (root/'features.npy').exists() else 'w+'
    features=np.lib.format.open_memmap(root/'features.npy',mode=mode,dtype=np.float32,shape=(int(ptr[-1]),4096))
    progress=json.loads((root/'progress.json').read_text()) if (root/'progress.json').exists() else {'batches':0}
    mace=build_author_mace_runtime().cuda().eval().requires_grad_(False)
    captures=[]
    class RawCaptured(Exception):
        pass
    def capture_raw(module, arguments):
        captures.append(arguments[0].detach())
        raise RawCaptured
    hook=mace.preprocessor.atomic_preprocessor.register_forward_pre_hook(capture_raw)
    started=time.time()
    if name=='bigsoldb':
        plan=list(common.batch_indices(np.arange(len(frame)),lengths,max_size=24,atom_cap=2400))
    else:plan=list(batches(np.arange(len(frame)),lengths,max_size=64,atom_cap=3600))
    status('caching_frozen_mace',corpus=name,total=len(frame),atoms=int(ptr[-1]),bytes=int(ptr[-1])*4096*4,batches=len(plan))
    with torch.inference_mode():
        for bi,ids in enumerate(plan,1):
            if bi<=progress['batches']:continue
            batch=yield_molecules_collate_fn([Sample(atomic_positions=positions[i],atomic_numbers=numbers[i],total_charge=float(frame.charge.iloc[i]) if 'charge' in frame else 0.,multiplicity=1.) for i in ids]).to_('cuda')
            captures.clear()
            try:
                mace.preprocessor(batch)
            except RawCaptured:
                pass
            assert len(captures)==1
            raw=captures.pop().float().cpu().numpy()
            for j,i in enumerate(ids):
                v=raw[j,:lengths[i]];assert v.shape==(lengths[i],4096) and np.isfinite(v).all()
                features[ptr[i]:ptr[i+1]]=v
            if bi%25==0 or bi==len(plan):
                features.flush();write(root/'progress.json',dict(batches=bi,total_batches=len(plan)))
                status('caching_frozen_mace',corpus=name,batch=bi,batches=len(plan),seconds=time.time()-started)
    features.flush();del features
    hook.remove();del mace,numbers,positions;gc.collect();torch.cuda.empty_cache()
    write(root/'complete.json',dict(rows=len(frame),atoms=int(ptr[-1]),width=4096,dtype='float32',
        mace_sha256='fab8b8713c832f31a2a853aaa22fd638be8a369cbf5095e6b3e982a18d10e93a',
        raw_capture_only=True,description='Frozen MACE only. Full-precision inputs to learned atomic/geometric preprocessors; no trainable activations cached.'))