"""Embed every train + validation protein with ESM-2 on an Apple-silicon GPU (MPS).

Run on the Mac:
    python3 esm_embed.py --model 35M      (~10-20 min)
    python3 esm_embed.py --model 650M     (~2-4 h, optional)

Writes to ~/enzyme-esm/:
    ids_train.txt, ids_val.txt            sequence ids in row order
    emb_<model>_train.npy, emb_<model>_val.npy   float16, mean-pooled final layer
The ESM-2 weights are sequence-only (masked-LM pretraining, no EC labels).
"""
import argparse, os, sys, time
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, EsmModel

MODELS = {'8M': 'facebook/esm2_t6_8M_UR50D', '35M': 'facebook/esm2_t12_35M_UR50D',
          '150M': 'facebook/esm2_t30_150M_UR50D', '650M': 'facebook/esm2_t33_650M_UR50D'}
OUT = Path.home() / 'enzyme-esm'
MAX_LEN = 1022
PART = 10000


def load_data():
    d = OUT / 'data'
    d.mkdir(parents=True, exist_ok=True)
    paths = {}
    for f in ('train.parquet', 'val_inputs.parquet'):
        paths[f] = hf_hub_download('starkAhmed43/leak-curber-enzyme-classification', f,
                                   repo_type='dataset', local_dir=d)
    tr = pd.read_parquet(paths['train.parquet']).drop_duplicates('sequence')
    va = pd.read_parquet(paths['val_inputs.parquet']).drop_duplicates('sequence')
    return tr[['id', 'sequence']], va[['id', 'sequence']]


@torch.no_grad()
def embed(model, tok, seqs, device, token_budget):
    order = np.argsort([len(s) for s in seqs])[::-1]
    out = np.zeros((len(seqs), model.config.hidden_size), dtype=np.float16)
    i, t0, done = 0, time.time(), 0
    while i < len(order):
        L = min(len(seqs[order[i]]), MAX_LEN) + 2
        bs = max(1, token_budget // L)
        idx = order[i:i + bs]
        batch = tok([seqs[j][:MAX_LEN] for j in idx], return_tensors='pt', padding=True)
        batch = {k: v.to(device) for k, v in batch.items()}
        h = model(**batch).last_hidden_state
        m = batch['attention_mask'].clone()
        m[:, 0] = 0                                           # drop <cls>
        m[torch.arange(len(idx)), batch['attention_mask'].sum(1) - 1] = 0   # drop <eos>
        e = (h * m.unsqueeze(-1)).sum(1) / m.sum(1, keepdim=True).clamp(min=1)
        out[idx] = e.float().cpu().numpy().astype(np.float16)
        i += bs
        done += len(idx)
        if done // 2000 != (done - len(idx)) // 2000:
            rate = done / (time.time() - t0)
            print(f'  {done}/{len(seqs)}  {rate:.0f} seq/s  eta {(len(seqs) - done) / rate / 60:.1f} min', flush=True)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='35M', choices=MODELS)
    ap.add_argument('--tokens', type=int, default=0, help='tokens per batch (default by model size)')
    ap.add_argument('--max-train', type=int, default=0, help='embed only a random sample of N training proteins')
    a = ap.parse_args()
    device = 'mps' if torch.backends.mps.is_available() else 'cpu'
    budget = a.tokens or {'8M': 20000, '35M': 8000, '150M': 4000, '650M': 2000}[a.model]
    print(f'device={device} model={MODELS[a.model]} tokens/batch={budget}', flush=True)
    tr, va = load_data()
    sfx = ''
    if a.max_train and a.max_train < len(tr):
        tr = tr.sample(a.max_train, random_state=0)
        sfx = f'_s{a.max_train}'
    (OUT / f'ids_train{sfx}.txt').write_text('\n'.join(tr.id))
    (OUT / 'ids_val.txt').write_text('\n'.join(va.id))
    tok = AutoTokenizer.from_pretrained(MODELS[a.model])
    model = EsmModel.from_pretrained(MODELS[a.model], add_pooling_layer=False).to(device).eval()
    for name, df in (('val', va), ('train' + sfx, tr)):
        path = OUT / f'emb_{a.model}_{name}.npy'
        if path.exists():
            print(f'skip {name}: {path} already exists', flush=True)
            continue
        # Checkpointed: embed in parts of PART sequences, each saved as soon as it is done,
        # so a crash or a dead battery only loses the part in progress. Re-running resumes.
        seqs = list(df.sequence)
        parts = []
        nparts = (len(seqs) + PART - 1) // PART
        for k in range(nparts):
            pp = OUT / f'emb_{a.model}_{name}.part{k:03d}.npy'
            parts.append(pp)
            if pp.exists():
                print(f'  part {k + 1}/{nparts} already done', flush=True)
                continue
            print(f'embedding {name} part {k + 1}/{nparts} ({len(seqs[k * PART:(k + 1) * PART])} sequences)', flush=True)
            e = embed(model, tok, seqs[k * PART:(k + 1) * PART], device, budget)
            np.save(pp, e)
            if device == 'mps':
                torch.mps.empty_cache()
        e = np.concatenate([np.load(pp) for pp in parts])
        np.save(path, e)
        for pp in parts:
            pp.unlink()
        print(f'saved {path} {e.shape}', flush=True)
    print('DONE', flush=True)


if __name__ == '__main__':
    sys.exit(main())