# ============================================================================
# GEOLIP-BERTENSTEIN STAGE 1: MULTI-EXPERT PRECOMPUTE (REFACTORED)
#
# BERT is the shared text spine.
#
# Pipeline per expert pair:
#   1. Load dataset / stream
#   2. CPU preprocess text + expert input
#   3. GPU encode text with BERT + expert with expert encoder
#   4. Shard-safe Arrow write
#   5. Merge shards -> final save_to_disk
#   6. Unload expert, keep BERT
#
# Experts:
#   image   : DINOv2-large   + COCO-Caption
#   audio   : Whisper-large  + LibriSpeech ASR (streaming)
#   protein : ESM-2-650M     + Protein2Text-QA (streaming)
#   code    : CodeBERT-base  + CodeSearchNet python
# ============================================================================

import subprocess
import sys

try:
    import sympy
    _ = sympy.core
except (ImportError, AttributeError):
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "--upgrade", "sympy", "--break-system-packages", "-q"]
    )

import gc
import os
import shutil
import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from datasets import (
    Audio,
    Dataset as HFDataset,
    Features,
    Sequence,
    Value,
    Array2D,
    concatenate_datasets,
    load_dataset,
    load_from_disk,
)

# ============================================================================
# BASE CONFIG
# ============================================================================

@dataclass
class BaseConfig:
    cache_dir: str = "/home/claude/geo_cache"
    max_text_len: int = 32

    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    amp_enabled: bool = torch.cuda.is_available()

    bert_model_name: str = "google-bert/bert-large-uncased"
    bert_hidden_dim: int = 1024

    batch_size: int = 256
    num_workers: int = 8
    prefetch_factor: int = 2
    pin_memory: bool = torch.cuda.is_available()

    shard_size_default: int = 2048

    # expert-specific max samples
    max_audio_samples: int = 10000
    max_protein_samples: int = 15000
    max_code_samples: int = 50000

    cleanup_hf_cache_between_experts: bool = True


CFG = BaseConfig()
DEVICE = torch.device(CFG.device)


# ============================================================================
# HF CACHE CLEANUP
# ============================================================================

def cleanup_hf_cache() -> None:
    """Delete HF datasets/hub cache to free disk after encoding an expert."""
    hf_cache = os.path.expanduser("~/.cache/huggingface")
    for subdir in ["datasets", "hub"]:
        p = os.path.join(hf_cache, subdir)
        if not os.path.exists(p):
            continue

        size_gb = 0.0
        for dp, _, files in os.walk(p):
            for f in files:
                fp = os.path.join(dp, f)
                try:
                    size_gb += os.path.getsize(fp)
                except OSError:
                    pass
        size_gb /= 1e9

        print(f"    Cleaning {p} ({size_gb:.1f} GB)...")
        shutil.rmtree(p, ignore_errors=True)
        os.makedirs(p, exist_ok=True)


def cleanup_cuda() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


# ============================================================================
# SHARED BERT
# ============================================================================

_bert_tokenizer = None

def get_bert_tokenizer():
    global _bert_tokenizer
    if _bert_tokenizer is None:
        from transformers import BertTokenizer
        _bert_tokenizer = BertTokenizer.from_pretrained(CFG.bert_model_name)
    return _bert_tokenizer


def load_shared_bert():
    from transformers import BertModel
    print("Loading shared BERT-large...")
    bert = BertModel.from_pretrained(
        CFG.bert_model_name,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    ).to(DEVICE).eval()
    print("  BERT ready.")
    return bert


# ============================================================================
# COMMON HELPERS
# ============================================================================

def ensure_dir(path: str) -> None:
    os.makedirs(path, exist_ok=True)


def make_loader(ds: Dataset, batch_size: int, num_workers: int) -> DataLoader:
    kwargs = dict(
        dataset=ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=CFG.pin_memory,
        persistent_workers=num_workers > 0,
    )
    if num_workers > 0:
        kwargs["prefetch_factor"] = CFG.prefetch_factor
    return DataLoader(**kwargs)


def masked_text_tokenize(text: str, tokenizer) -> Tuple[torch.Tensor, torch.Tensor]:
    tok = tokenizer(
        text,
        padding="max_length",
        truncation=True,
        max_length=CFG.max_text_len,
        return_tensors="pt",
    )
    return tok["input_ids"].squeeze(0), tok["attention_mask"].squeeze(0)


def extract_first_text(sample: Dict[str, Any], keys: List[str]) -> str:
    for key in keys:
        if key not in sample:
            continue
        value = sample[key]

        if isinstance(value, str):
            value = value.strip()
            if value:
                return value

        if isinstance(value, list) and value:
            first = value[0]
            if isinstance(first, str):
                first = first.strip()
                if first:
                    return first
            if isinstance(first, dict):
                txt = str(first.get("raw", first.get("text", ""))).strip()
                if txt:
                    return txt
            txt = str(first).strip()
            if txt:
                return txt

    return ""


# ============================================================================
# SHARD WRITER
# ============================================================================

class ShardWriter:
    def __init__(
        self,
        cache_dir: str,
        tag: str,
        features: Features,
        shard_size: int,
        row_keys: List[str],
    ):
        self.cache_dir = cache_dir
        self.tag = tag
        self.features = features
        self.shard_size = shard_size
        self.row_keys = row_keys

        self.cache_path = os.path.join(cache_dir, tag)
        self.shard_root = os.path.join(cache_dir, f"{tag}__shards")

        self.rows = {k: [] for k in row_keys}
        self.shard_paths: List[str] = []
        self.shard_idx = 0
        self.n_written = 0

    @property
    def exists(self) -> bool:
        return os.path.exists(self.cache_path)

    def add_row(self, row: Dict[str, Any]) -> None:
        for k in self.row_keys:
            self.rows[k].append(row[k])

        if len(self.rows[self.row_keys[0]]) >= self.shard_size:
            self.flush()

    def flush(self) -> None:
        n_rows = len(self.rows[self.row_keys[0]])
        if n_rows == 0:
            return

        ensure_dir(self.shard_root)
        shard_path = os.path.join(self.shard_root, f"shard_{self.shard_idx:05d}")
        ds = HFDataset.from_dict(self.rows, features=self.features)
        ds.save_to_disk(shard_path)

        self.shard_paths.append(shard_path)
        self.shard_idx += 1
        self.n_written += n_rows
        self.rows = {k: [] for k in self.row_keys}

    def finalize(self) -> str:
        self.flush()

        print(f"    Merging {len(self.shard_paths)} shards...")
        merged = concatenate_datasets([load_from_disk(p) for p in self.shard_paths])
        merged.save_to_disk(self.cache_path)
        print(f"    Saved {len(merged)} pairs to {self.cache_path}")

        if os.path.exists(self.shard_root):
            shutil.rmtree(self.shard_root, ignore_errors=True)

        return self.cache_path


# ============================================================================
# MAP-STYLE DATASETS (NON-STREAMING)
# ============================================================================

class ImageTextDataset(Dataset):
    def __init__(self, hf_ds, bert_tokenizer, image_processor):
        self.ds = hf_ds
        self.tok = bert_tokenizer
        self.proc = image_processor
        self.fallback_shape = (3, 518, 518)

    def __len__(self):
        return len(self.ds)

    def __getitem__(self, idx):
        sample = self.ds[idx]

        caption = extract_first_text(
            sample,
            ["answer", "caption", "captions", "text", "original_alt_text"],
        )
        ids, mask = masked_text_tokenize(caption, self.tok)

        image = sample.get("image", None)
        valid = True

        if image is not None and hasattr(image, "convert"):
            try:
                expert_input = self.proc(
                    images=image.convert("RGB"),
                    return_tensors="pt",
                )["pixel_values"].squeeze(0)
            except Exception:
                expert_input = torch.zeros(self.fallback_shape, dtype=torch.float32)
                valid = False
        else:
            expert_input = torch.zeros(self.fallback_shape, dtype=torch.float32)
            valid = False

        return ids, mask, expert_input, valid


class CodeTextDataset(Dataset):
    def __init__(self, hf_ds, bert_tokenizer, code_tokenizer):
        self.ds = hf_ds
        self.tok = bert_tokenizer
        self.code_tok = code_tokenizer

    def __len__(self):
        return len(self.ds)

    def __getitem__(self, idx):
        sample = self.ds[idx]

        doc = sample.get("func_documentation_string", "")
        if not doc or not doc.strip():
            doc = str(sample.get("whole_func_string", ""))[:200]
        doc = str(doc).strip()[:500]

        ids, mask = masked_text_tokenize(doc, self.tok)

        code = str(sample.get("func_code_string", sample.get("whole_func_string", ""))).strip()[:512]
        valid = len(code) > 5 and len(doc) > 5

        if valid:
            try:
                tok = self.code_tok(
                    code,
                    padding="max_length",
                    truncation=True,
                    max_length=256,
                    return_tensors="pt",
                )
                code_ids = tok["input_ids"].squeeze(0)
                code_mask = tok["attention_mask"].squeeze(0)
            except Exception:
                code_ids = torch.zeros(256, dtype=torch.long)
                code_mask = torch.zeros(256, dtype=torch.long)
                valid = False
        else:
            code_ids = torch.zeros(256, dtype=torch.long)
            code_mask = torch.zeros(256, dtype=torch.long)

        return ids, mask, torch.stack([code_ids, code_mask]), valid


# ============================================================================
# SHARED NON-STREAM ENCODER
# ============================================================================

@torch.no_grad()
def encode_map_dataset(
    *,
    tag: str,
    loader: DataLoader,
    bert,
    expert_name: str,
    expert_hidden_shape: Tuple[int, int],
    expert_forward: Callable[[torch.Tensor], torch.Tensor],
    shard_size: int,
    max_samples: Optional[int] = None,
) -> str:
    cache_path = os.path.join(CFG.cache_dir, tag)
    if os.path.exists(cache_path):
        print(f"    Cache exists: {cache_path}")
        return cache_path

    features = Features({
        "text_hidden": Array2D(shape=(CFG.max_text_len, CFG.bert_hidden_dim), dtype="float16"),
        "text_mask": Sequence(Value("bool"), length=CFG.max_text_len),
        f"{expert_name}_hidden": Array2D(shape=expert_hidden_shape, dtype="float16"),
    })

    writer = ShardWriter(
        cache_dir=CFG.cache_dir,
        tag=tag,
        features=features,
        shard_size=shard_size,
        row_keys=["text_hidden", "text_mask", f"{expert_name}_hidden"],
    )

    t0 = time.time()
    n = 0

    for batch in loader:
        text_ids, text_mask, expert_input, valid = batch
        valid_b = valid.bool()

        if not valid_b.any():
            continue

        text_ids = text_ids[valid_b].to(DEVICE, non_blocking=True)
        text_mask_gpu = text_mask[valid_b].to(DEVICE, non_blocking=True)
        expert_input = expert_input[valid_b].to(DEVICE, non_blocking=True)

        text_hidden = bert(
            input_ids=text_ids,
            attention_mask=text_mask_gpu,
        ).last_hidden_state.detach().to(dtype=torch.float16).cpu().numpy()

        text_mask_np = text_mask_gpu.bool().cpu().numpy()
        expert_hidden = expert_forward(expert_input).detach().to(dtype=torch.float16).cpu().numpy()

        for i in range(text_hidden.shape[0]):
            writer.add_row({
                "text_hidden": text_hidden[i],
                "text_mask": text_mask_np[i].tolist(),
                f"{expert_name}_hidden": expert_hidden[i],
            })

        n += text_hidden.shape[0]
        if n % 1000 < CFG.batch_size or n <= CFG.batch_size:
            rate = n / max(time.time() - t0, 1e-6)
            print(f"    {n}" + (f"/{max_samples}" if max_samples else "") + f" ({rate:.0f}/s)")

        if max_samples is not None and n >= max_samples:
            break

    final_path = writer.finalize()
    print(f"    Completed {n} samples in {time.time() - t0:.0f}s")
    return final_path


# ============================================================================
# STREAMING HELPERS
# ============================================================================

def decode_audio_obj(audio_obj) -> Tuple[np.ndarray, int]:
    if hasattr(audio_obj, "get_all_samples"):
        samples = audio_obj.get_all_samples()
        arr = samples.data.numpy().squeeze()
        sr = samples.sample_rate
        return arr, sr

    if isinstance(audio_obj, dict):
        return audio_obj["array"], audio_obj.get("sampling_rate", 16000)

    raise TypeError(f"Unsupported audio object type: {type(audio_obj)}")


def stream_librispeech_batches(
    stream,
    bert_tokenizer,
    whisper_processor,
    batch_size: int,
) -> Iterable[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
    batch_ids = []
    batch_masks = []
    batch_mels = []

    for sample in stream:
        text = sample.get("text", sample.get("transcription", ""))
        audio_obj = sample.get("audio")
        if not text or audio_obj is None:
            continue

        try:
            audio_array, sr = decode_audio_obj(audio_obj)
        except Exception:
            continue

        ids, mask = masked_text_tokenize(str(text), bert_tokenizer)

        try:
            mel = whisper_processor(
                audio_array,
                sampling_rate=sr,
                return_tensors="pt",
            ).input_features.squeeze(0)
        except Exception:
            continue

        batch_ids.append(ids)
        batch_masks.append(mask)
        batch_mels.append(mel)

        if len(batch_ids) >= batch_size:
            yield (
                torch.stack(batch_ids),
                torch.stack(batch_masks),
                torch.stack(batch_mels),
            )
            batch_ids, batch_masks, batch_mels = [], [], []

    if batch_ids:
        yield (
            torch.stack(batch_ids),
            torch.stack(batch_masks),
            torch.stack(batch_mels),
        )


def extract_protein_caption(sample: Dict[str, Any]) -> str:
    convos = sample.get("conversations", [])
    if isinstance(convos, list):
        for c in convos:
            if isinstance(c, dict) and c.get("from") == "gpt":
                v = str(c.get("value", "")).strip()
                if v:
                    return v[:500]
        for c in convos:
            if isinstance(c, dict) and "value" in c:
                v = str(c["value"]).strip()
                if v:
                    return v[:500]

    return str(sample.get("protein", "")).strip()[:500]


def stream_protein_batches(
    stream,
    bert_tokenizer,
    esm_tokenizer,
    batch_size: int,
) -> Iterable[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]:
    batch_ids = []
    batch_masks = []
    batch_esm_ids = []
    batch_esm_masks = []

    for sample in stream:
        caption = extract_protein_caption(sample)
        seq = str(sample.get("amino_seq", sample.get("protein_sequence", ""))).strip()

        if len(caption) < 5 or len(seq) < 5:
            continue

        ids, mask = masked_text_tokenize(caption, bert_tokenizer)

        try:
            esm_t = esm_tokenizer(
                seq,
                padding="max_length",
                truncation=True,
                max_length=512,
                return_tensors="pt",
            )
        except Exception:
            continue

        batch_ids.append(ids)
        batch_masks.append(mask)
        batch_esm_ids.append(esm_t["input_ids"].squeeze(0))
        batch_esm_masks.append(esm_t["attention_mask"].squeeze(0))

        if len(batch_ids) >= batch_size:
            yield (
                torch.stack(batch_ids),
                torch.stack(batch_masks),
                torch.stack(batch_esm_ids),
                torch.stack(batch_esm_masks),
            )
            batch_ids, batch_masks, batch_esm_ids, batch_esm_masks = [], [], [], []

    if batch_ids:
        yield (
            torch.stack(batch_ids),
            torch.stack(batch_masks),
            torch.stack(batch_esm_ids),
            torch.stack(batch_esm_masks),
        )


@torch.no_grad()
def encode_streaming_batches(
    *,
    tag: str,
    expert_name: str,
    expert_hidden_shape: Tuple[int, int],
    batch_iter: Iterable,
    bert,
    expert_batch_forward: Callable[..., torch.Tensor],
    shard_size: int,
    row_keys: List[str],
    max_samples: Optional[int] = None,
) -> str:
    cache_path = os.path.join(CFG.cache_dir, tag)
    if os.path.exists(cache_path):
        print(f"    Cache exists: {cache_path}")
        return cache_path

    features = Features({
        "text_hidden": Array2D(shape=(CFG.max_text_len, CFG.bert_hidden_dim), dtype="float16"),
        "text_mask": Sequence(Value("bool"), length=CFG.max_text_len),
        f"{expert_name}_hidden": Array2D(shape=expert_hidden_shape, dtype="float16"),
    })

    writer = ShardWriter(
        cache_dir=CFG.cache_dir,
        tag=tag,
        features=features,
        shard_size=shard_size,
        row_keys=row_keys,
    )

    t0 = time.time()
    n = 0

    for packed in batch_iter:
        # first two are always bert ids/masks
        text_ids = packed[0].to(DEVICE, non_blocking=True)
        text_mask = packed[1].to(DEVICE, non_blocking=True)

        text_hidden = bert(
            input_ids=text_ids,
            attention_mask=text_mask,
        ).last_hidden_state.detach().to(dtype=torch.float16).cpu().numpy()

        text_mask_np = text_mask.bool().cpu().numpy()

        expert_hidden = expert_batch_forward(*[p.to(DEVICE, non_blocking=True) for p in packed[2:]])
        expert_hidden = expert_hidden.detach().to(dtype=torch.float16).cpu().numpy()

        for i in range(text_hidden.shape[0]):
            writer.add_row({
                "text_hidden": text_hidden[i],
                "text_mask": text_mask_np[i].tolist(),
                f"{expert_name}_hidden": expert_hidden[i],
            })

        n += text_hidden.shape[0]
        batch_size = text_hidden.shape[0]
        if n % 1000 < batch_size or n <= batch_size:
            rate = n / max(time.time() - t0, 1e-6)
            print(f"    {n}" + (f"/{max_samples}" if max_samples else "") + f" ({rate:.0f}/s)")

        if max_samples is not None and n >= max_samples:
            break

    final_path = writer.finalize()
    print(f"    Completed {n} samples in {time.time() - t0:.0f}s")
    return final_path


# ============================================================================
# EXPERT RUNNERS
# ============================================================================

def encode_image_expert(bert, split: str, tag: str, max_samples: Optional[int] = None) -> str:
    from transformers import Dinov2Model, AutoImageProcessor

    print(f"\n  [IMAGE] Loading DINOv2-large + COCO-Caption ({split})...")
    dino = Dinov2Model.from_pretrained(
        "facebook/dinov2-large",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    ).to(DEVICE).eval()
    proc = AutoImageProcessor.from_pretrained("facebook/dinov2-large")
    tok = get_bert_tokenizer()

    hf_ds = load_dataset("lmms-lab/COCO-Caption", split=split)
    if max_samples is not None:
        hf_ds = hf_ds.select(range(min(max_samples, len(hf_ds))))
    print(f"    Dataset: {len(hf_ds)} samples")

    torch_ds = ImageTextDataset(hf_ds, tok, proc)
    loader = make_loader(torch_ds, batch_size=CFG.batch_size, num_workers=CFG.num_workers)

    def expert_forward(pixel_values):
        return dino(pixel_values=pixel_values).last_hidden_state

    path = encode_map_dataset(
        tag=tag,
        loader=loader,
        bert=bert,
        expert_name="image",
        expert_hidden_shape=(257, 1024),
        expert_forward=expert_forward,
        shard_size=CFG.shard_size_default,
        max_samples=max_samples,
    )

    del dino, proc, hf_ds, torch_ds, loader
    cleanup_cuda()
    return path


def encode_code_expert(bert, max_samples: Optional[int] = None) -> str:
    from transformers import RobertaModel, RobertaTokenizer

    print("\n  [CODE] Loading CodeBERT + CodeSearchNet python...")
    codebert = RobertaModel.from_pretrained(
        "microsoft/codebert-base",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    ).to(DEVICE).eval()
    code_tok = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
    tok = get_bert_tokenizer()

    hf_ds = load_dataset("code-search-net/code_search_net", "python", split="train")
    if max_samples is not None:
        hf_ds = hf_ds.select(range(min(max_samples, len(hf_ds))))

    hf_ds = hf_ds.filter(
        lambda x: bool(x.get("func_documentation_string", "").strip()),
        num_proc=4,
    )
    print(f"    Dataset: {len(hf_ds)} samples (after filtering)")

    torch_ds = CodeTextDataset(hf_ds, tok, code_tok)
    loader = make_loader(torch_ds, batch_size=CFG.batch_size, num_workers=CFG.num_workers)

    def expert_forward(packed):
        code_ids = packed[:, 0].long()
        code_mask = packed[:, 1].long()
        return codebert(input_ids=code_ids, attention_mask=code_mask).last_hidden_state

    path = encode_map_dataset(
        tag="code_csn",
        loader=loader,
        bert=bert,
        expert_name="code",
        expert_hidden_shape=(256, 768),
        expert_forward=expert_forward,
        shard_size=CFG.shard_size_default,
        max_samples=max_samples,
    )

    del codebert, code_tok, hf_ds, torch_ds, loader
    cleanup_cuda()
    return path


def encode_audio_expert(bert, max_samples: Optional[int] = None) -> str:
    from transformers import WhisperModel, WhisperFeatureExtractor

    print("\n  [AUDIO] Loading Whisper-large-v3 + LibriSpeech ASR (streaming)...")
    whisper = WhisperModel.from_pretrained(
        "openai/whisper-large-v3",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    ).to(DEVICE).eval()
    proc = WhisperFeatureExtractor.from_pretrained("openai/whisper-large-v3")
    tok = get_bert_tokenizer()

    max_n = max_samples or CFG.max_audio_samples
    audio_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

    # probe shape
    probe_stream = load_dataset("openslr/librispeech_asr", "clean", split="train.100", streaming=True)
    probe_stream = probe_stream.cast_column("audio", Audio(sampling_rate=16000))
    first = next(iter(probe_stream))
    arr, sr = decode_audio_obj(first["audio"])

    mel = proc(arr, sampling_rate=sr, return_tensors="pt").input_features
    mel = mel.to(device=DEVICE, dtype=audio_dtype)

    with torch.no_grad():
        probe_hidden = whisper.encoder(mel).last_hidden_state

    seq_len, hidden_dim = probe_hidden.shape[1], probe_hidden.shape[2]
    print(f"    Whisper encoder output: ({seq_len}, {hidden_dim})")
    del mel, probe_hidden

    stream = load_dataset("openslr/librispeech_asr", "clean", split="train.100", streaming=True)
    stream = stream.cast_column("audio", Audio(sampling_rate=16000))

    batch_iter = stream_librispeech_batches(
        stream=stream,
        bert_tokenizer=tok,
        whisper_processor=proc,
        batch_size=16,
    )

    def expert_batch_forward(mels: torch.Tensor) -> torch.Tensor:
        mels = mels.to(dtype=audio_dtype)
        return whisper.encoder(mels).last_hidden_state

    path = encode_streaming_batches(
        tag="audio_librispeech",
        expert_name="audio",
        expert_hidden_shape=(seq_len, hidden_dim),
        batch_iter=batch_iter,
        bert=bert,
        expert_batch_forward=expert_batch_forward,
        shard_size=256,  # large hidden size; keep shards smaller
        row_keys=["text_hidden", "text_mask", "audio_hidden"],
        max_samples=max_n,
    )

    del whisper, proc
    cleanup_cuda()
    return path


def encode_protein_expert(bert, max_samples: Optional[int] = None) -> str:
    from transformers import EsmModel, EsmTokenizer

    print("\n  [PROTEIN] Loading ESM-2-650M + Protein2Text-QA (streaming)...")
    esm = EsmModel.from_pretrained(
        "facebook/esm2_t33_650M_UR50D",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    ).to(DEVICE).eval()
    esm_tok = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
    tok = get_bert_tokenizer()

    max_n = max_samples or CFG.max_protein_samples
    stream = load_dataset("tumorailab/Protein2Text-QA", split="test", streaming=True)

    batch_iter = stream_protein_batches(
        stream=stream,
        bert_tokenizer=tok,
        esm_tokenizer=esm_tok,
        batch_size=32,
    )

    def expert_batch_forward(esm_ids: torch.Tensor, esm_mask: torch.Tensor) -> torch.Tensor:
        return esm(input_ids=esm_ids.long(), attention_mask=esm_mask.long()).last_hidden_state

    path = encode_streaming_batches(
        tag="protein_p2t",
        expert_name="protein",
        expert_hidden_shape=(512, 1280),
        batch_iter=batch_iter,
        bert=bert,
        expert_batch_forward=expert_batch_forward,
        shard_size=512,
        row_keys=["text_hidden", "text_mask", "protein_hidden"],
        max_samples=max_n,
    )

    del esm, esm_tok
    cleanup_cuda()
    return path


# ============================================================================
# MAIN
# ============================================================================

def main():
    ensure_dir(CFG.cache_dir)

    print("=" * 70)
    print("STAGE 1: MULTI-EXPERT PRECOMPUTE")
    print("=" * 70)

    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name()}")
        print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    print(f"Cache: {CFG.cache_dir}")

    required_tags = [
        "image_coco",
        "image_coco_test",
        "audio_librispeech",
        "protein_p2t",
        "code_csn",
    ]
    missing = [t for t in required_tags if not os.path.exists(os.path.join(CFG.cache_dir, t))]

    if not missing:
        print("\nAll caches exist. Nothing to encode.")
        bert = None
    else:
        print(f"\nMissing caches: {missing}")
        if CFG.cleanup_hf_cache_between_experts:
            cleanup_hf_cache()
        bert = load_shared_bert()

    paths: Dict[str, Optional[str]] = {}

    # IMAGE TRAIN
    print(f"\n{'─' * 50}")
    print("[1/4] IMAGE — DINOv2 + COCO-Caption")
    if os.path.exists(os.path.join(CFG.cache_dir, "image_coco")):
        print("  [IMAGE] Cache exists, skipping.")
        paths["image"] = os.path.join(CFG.cache_dir, "image_coco")
    else:
        paths["image"] = encode_image_expert(bert, split="val", tag="image_coco")
        if CFG.cleanup_hf_cache_between_experts:
            cleanup_hf_cache()

    # IMAGE TEST
    if os.path.exists(os.path.join(CFG.cache_dir, "image_coco_test")):
        print("\n  [IMAGE-TEST] Cache exists, skipping.")
        paths["image_test"] = os.path.join(CFG.cache_dir, "image_coco_test")
    else:
        print("\n  [IMAGE-TEST] COCO test split...")
        paths["image_test"] = encode_image_expert(bert, split="test", tag="image_coco_test")
        if CFG.cleanup_hf_cache_between_experts:
            cleanup_hf_cache()

    # AUDIO
    print(f"\n{'─' * 50}")
    print("[2/4] AUDIO — Whisper + LibriSpeech")
    if os.path.exists(os.path.join(CFG.cache_dir, "audio_librispeech")):
        print("  [AUDIO] Cache exists, skipping.")
        paths["audio"] = os.path.join(CFG.cache_dir, "audio_librispeech")
    else:
        try:
            paths["audio"] = encode_audio_expert(bert, max_samples=CFG.max_audio_samples)
        except Exception as e:
            print(f"  AUDIO failed: {e}")
            paths["audio"] = None
        if CFG.cleanup_hf_cache_between_experts:
            cleanup_hf_cache()

    # PROTEIN
    print(f"\n{'─' * 50}")
    print("[3/4] PROTEIN — ESM-2 + Protein2Text-QA")
    if os.path.exists(os.path.join(CFG.cache_dir, "protein_p2t")):
        print("  [PROTEIN] Cache exists, skipping.")
        paths["protein"] = os.path.join(CFG.cache_dir, "protein_p2t")
    else:
        try:
            paths["protein"] = encode_protein_expert(bert, max_samples=CFG.max_protein_samples)
        except Exception as e:
            print(f"  PROTEIN failed: {e}")
            paths["protein"] = None
        if CFG.cleanup_hf_cache_between_experts:
            cleanup_hf_cache()

    # CODE
    print(f"\n{'─' * 50}")
    print("[4/4] CODE — CodeBERT + CodeSearchNet Python")
    if os.path.exists(os.path.join(CFG.cache_dir, "code_csn")):
        print("  [CODE] Cache exists, skipping.")
        paths["code"] = os.path.join(CFG.cache_dir, "code_csn")
    else:
        try:
            paths["code"] = encode_code_expert(bert, max_samples=CFG.max_code_samples)
        except Exception as e:
            print(f"  CODE failed: {e}")
            paths["code"] = None
        if CFG.cleanup_hf_cache_between_experts:
            cleanup_hf_cache()

    if bert is not None:
        del bert
        cleanup_cuda()

    flickr_path = os.path.join(CFG.cache_dir, "flickr30k")
    if os.path.exists(flickr_path):
        paths["flickr"] = flickr_path

    print(f"\n{'=' * 70}")
    print("CACHE SUMMARY")
    print(f"{'=' * 70}")
    for name, path in paths.items():
        if path and os.path.exists(path):
            ds = load_from_disk(path)
            print(f"  {name:15s}: {len(ds):6d} pairs  [{path}]")

    print("\nReady for Stage 2 multi-expert training.")
    print("Done.")


if __name__ == "__main__":
    main()