#!/usr/bin/env python
"""Pre-compute frozen base model embeddings for belief texts.

For each record with high-quality beliefs (GPT-4o or annotation),
encodes the belief text through the frozen Qwen3-VL-4B language model
and saves the mean-pooled hidden state from layer 28.

Output: data/belief_targets_v6.pt
  - "embeddings": [N, max_beliefs, 2560] float16
  - "ids": list of record IDs
  - "valid": [N, max_beliefs] bool

Usage:
    python tools/precompute_belief_targets.py
"""
import json, sys, torch, logging
from pathlib import Path
from tqdm import tqdm

ROOT = Path("PROJECT_ROOT")
sys.path.insert(0, str(ROOT))

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
log = logging.getLogger("targets")

TRAIN_JSONL = ROOT / "data/cot_corpus_v3/v6_sft_train.jsonl"
OUTPUT = ROOT / "data/belief_targets_v6.pt"
BASE_MODEL = ROOT / "models/Qwen3-VL-4B-Instruct"
TARGET_LAYER = 28
BATCH_SIZE = 64
MAX_BELIEFS = 8


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Load records with high-quality beliefs
    log.info("Loading training data...")
    lines = TRAIN_JSONL.read_text().strip().split("\n")
    records = []
    for l in lines:
        d = json.loads(l)
        bsrc = d.get("belief_source", "")
        if "gpt" in bsrc.lower() or "annotation" in bsrc.lower():
            records.append(d)
    log.info(f"  {len(records)} records with high-quality beliefs (out of {len(lines)})")

    # Load tokenizer only (not the full model with vision)
    log.info("Loading tokenizer + language model...")
    from transformers import AutoTokenizer, AutoModelForCausalLM

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    # Load just the language model part for text encoding
    # We use the full model but only process text (no images)
    from transformers import AutoModelForImageTextToText
    model = AutoModelForImageTextToText.from_pretrained(
        BASE_MODEL, torch_dtype=torch.bfloat16, trust_remote_code=True
    ).to(device).eval()

    log.info(f"  Model loaded on {device}")

    # Pre-compute embeddings
    all_embeddings = []
    all_ids = []
    all_valid = []

    for start in tqdm(range(0, len(records), BATCH_SIZE), desc="encoding"):
        batch = records[start:start + BATCH_SIZE]
        batch_beliefs = []
        batch_valid = []

        for rec in batch:
            beliefs = rec.get("beliefs_per_frame", [])
            n = min(len(beliefs), MAX_BELIEFS)
            # Pad to MAX_BELIEFS
            padded = beliefs[:MAX_BELIEFS] + [""] * (MAX_BELIEFS - n)
            valid = [True] * n + [False] * (MAX_BELIEFS - n)
            batch_beliefs.append(padded)
            batch_valid.append(valid)
            all_ids.append(rec["id"])

        # Flatten all belief texts for batch encoding
        flat_texts = []
        for beliefs in batch_beliefs:
            flat_texts.extend(beliefs)

        # Tokenize
        encoded = tokenizer(
            flat_texts, return_tensors="pt", padding=True,
            truncation=True, max_length=64
        ).to(device)

        with torch.no_grad():
            out = model(
                input_ids=encoded["input_ids"],
                attention_mask=encoded.get("attention_mask"),
                output_hidden_states=True,
                return_dict=True,
            )
            hs = out.hidden_states[TARGET_LAYER]  # [B*MAX_BELIEFS, L, D]
            mask = encoded["attention_mask"].unsqueeze(-1).to(hs.dtype)
            pooled = (hs * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
            pooled = pooled.to(torch.float16).cpu()
            del out

        # Reshape back to [batch, MAX_BELIEFS, D]
        D = pooled.shape[-1]
        pooled = pooled.view(len(batch), MAX_BELIEFS, D)
        all_embeddings.append(pooled)
        all_valid.extend(batch_valid)

    embeddings = torch.cat(all_embeddings, dim=0)
    valid = torch.tensor(all_valid, dtype=torch.bool)

    log.info(f"Embeddings: {embeddings.shape} ({embeddings.dtype})")
    log.info(f"Valid: {valid.shape}")

    torch.save({
        "embeddings": embeddings,
        "ids": all_ids,
        "valid": valid,
        "layer": TARGET_LAYER,
        "model": str(BASE_MODEL),
    }, OUTPUT)

    log.info(f"Saved → {OUTPUT}")


if __name__ == "__main__":
    main()