from __future__ import annotations

import json
import time
from pathlib import Path

import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer


IN_PATH = Path("runs/decode_lab/mauve_export_lm1b_latest_1024.jsonl")
OUT_PATH = Path("runs/decode_lab/mauve_features_lm1b_latest_1024_gpt2large.npz")
MODEL_PATH = "/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard"
MAX_LEN = 256
BATCH_SIZE = 16


def load_texts() -> tuple[list[str], list[str], dict[str, list[str]], dict[str, list[str]]]:
    refs_raw: list[str] = []
    refs_stripped: list[str] = []
    gen_raw: dict[str, list[str]] = {}
    gen_stripped: dict[str, list[str]] = {}
    with IN_PATH.open(encoding="utf-8") as f:
        for line in f:
            obj = json.loads(line)
            typ = obj.get("type")
            if typ == "reference":
                refs_raw.append(obj["raw_text"])
                refs_stripped.append(obj["stripped_text"])
            elif typ == "generated":
                setting = obj["setting"]
                key = f"t{float(setting['endpoint_temp']):.2f}_{setting['final_from']}"
                gen_raw.setdefault(key, []).append(obj["raw_text"])
                gen_stripped.setdefault(key, []).append(obj["stripped_text"])
    return refs_raw, refs_stripped, gen_raw, gen_stripped


@torch.no_grad()
def featurize(
    texts: list[str],
    name: str,
    tokenizer: AutoTokenizer,
    model: AutoModel,
    device: torch.device,
) -> np.ndarray:
    chunks: list[np.ndarray] = []
    t0 = time.time()
    for start in range(0, len(texts), BATCH_SIZE):
        batch = texts[start : start + BATCH_SIZE]
        enc = tokenizer(
            batch,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=MAX_LEN,
            return_attention_mask=True,
        ).to(device)
        out = model(
            input_ids=enc["input_ids"],
            attention_mask=enc["attention_mask"],
            output_hidden_states=True,
            return_dict=True,
        )
        hidden = out.hidden_states[-1]
        last_idx = enc["attention_mask"].sum(dim=1) - 1
        feat = hidden[torch.arange(hidden.size(0), device=device), last_idx]
        chunks.append(feat.float().cpu().numpy())
        if (start // BATCH_SIZE) % 10 == 0:
            print(f"{name} {start + len(batch)}/{len(texts)}", flush=True)
    arr = np.concatenate(chunks, axis=0)
    print(f"{name} {arr.shape} time={time.time() - t0:.1f}s", flush=True)
    return arr


def main() -> None:
    refs_raw, refs_stripped, gen_raw, gen_stripped = load_texts()
    print(
        "loaded",
        len(refs_raw),
        {key: len(value) for key, value in gen_raw.items()},
        flush=True,
    )
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id
    model = AutoModel.from_pretrained(MODEL_PATH, pad_token_id=tokenizer.eos_token_id).to(device).eval()

    arrays: dict[str, np.ndarray] = {
        "ref_raw": featurize(refs_raw, "ref_raw", tokenizer, model, device),
        "ref_stripped": featurize(refs_stripped, "ref_stripped", tokenizer, model, device),
    }
    for key in sorted(gen_raw):
        arrays[f"gen_{key}_raw"] = featurize(gen_raw[key], f"gen_{key}_raw", tokenizer, model, device)
        arrays[f"gen_{key}_stripped"] = featurize(
            gen_stripped[key],
            f"gen_{key}_stripped",
            tokenizer,
            model,
            device,
        )
    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(OUT_PATH, **arrays)
    print(f"DONE {OUT_PATH}", flush=True)


if __name__ == "__main__":
    main()