#!/usr/bin/env python
"""
提取 GLUE 领域的模型表示（限制 5000 条样本）
"""

import os
os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"

import gc
import torch
import json
from pathlib import Path
from typing import List
from transformers import AutoTokenizer, AutoModel
import warnings

warnings.filterwarnings("ignore")


def load_texts(filepath: str, limit: int = 5000) -> List[str]:
    """从 JSONL 文件加载文本，限制数量"""
    texts = []
    with open(filepath, "r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            if i >= limit:
                break
            record = json.loads(line.strip())
            texts.append(record.get("text", ""))
    return texts


def extract_representations(
    model_path: str,
    texts: List[str],
    batch_size: int = 8,
    max_length: int = 512
) -> torch.Tensor:
    """提取模型表示"""
    print(f"  加载模型：{model_path}")

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModel.from_pretrained(
        model_path,
        torch_dtype=torch.float32,
        device_map="auto",
        trust_remote_code=True,
        output_hidden_states=True
    )
    model.eval()

    print(f"  模型设备：{model.device}")
    print(f"  文本数量：{len(texts)}")

    all_embeddings = []

    with torch.no_grad():
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            inputs = tokenizer(
                batch_texts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=max_length
            )
            inputs = inputs.to(model.device)
            outputs = model(**inputs)
            hidden = outputs.last_hidden_state
            attention_mask = inputs["attention_mask"].unsqueeze(-1).float()
            sum_embeddings = (hidden * attention_mask).sum(dim=1)
            valid_tokens = attention_mask.sum(dim=1)
            mean_embeddings = sum_embeddings / (valid_tokens + 1e-9)
            all_embeddings.append(mean_embeddings.cpu())

            if (i // batch_size + 1) % 100 == 0:
                print(f"    进度：{min(i + batch_size, len(texts))}/{len(texts)}")

    embeddings = torch.cat(all_embeddings, dim=0)
    H = embeddings.T

    del model, tokenizer
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return H


def save_representation(H: torch.Tensor, model_name: str, domain: str, output_dir: str = "model/representations") -> str:
    """保存表示矩阵"""
    output_path = Path(output_dir) / f"{model_name}_{domain}.pt"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(H, output_path)

    meta_path = output_path.with_suffix(".json")
    meta = {
        "model_name": model_name,
        "domain": domain,
        "d": int(H.shape[0]),
        "N": int(H.shape[1]),
        "dtype": str(H.dtype)
    }
    with open(meta_path, "w") as f:
        json.dump(meta, f, indent=2)

    return str(output_path)


def get_model_name(model_path: str) -> str:
    """从路径提取模型简称"""
    path = Path(model_path)
    name = path.name.lower()
    if "qwen" in name:
        return "qwen_7b"
    elif "mistral" in name:
        return "mistral_7b"
    elif "llama3" in name or "llama-3" in name:
        return "llama3_8b"
    return path.name.replace("-", "_")


def main():
    print("=" * 70)
    print("GLUE 领域模型表示提取（限制 5000 条样本）")
    print("=" * 70)
    print()

    models = [
        "model/weights/Qwen2.5-7B-Instruct",
        "model/weights/Mistral-7B-v0.3",
        "model/weights/llama3-8b",
    ]

    # 仅 GLUE 领域
    glue_domains = [
        ("glue_mnli", "database/corpus/glue_mnli/texts.jsonl"),
        ("glue_qnli", "database/corpus/glue_qnli/texts.jsonl"),
        ("glue_sst2", "database/corpus/glue_sst2/texts.jsonl"),
        ("glue_stsb", "database/corpus/glue_stsb/texts.jsonl"),
        ("glue_cola", "database/corpus/glue_cola/texts.jsonl"),
    ]

    limit = 5000

    for model_path in models:
        if not Path(model_path).exists():
            print(f"跳过（模型不存在）：{model_path}")
            continue

        model_name = get_model_name(model_path)
        print(f"\n[模型：{model_name}]")

        for domain, texts_path in glue_domains:
            print(f"\n  [{domain}]")

            texts = load_texts(texts_path, limit=limit)
            print(f"    加载了 {len(texts)} 条文本（限制：{limit}）")

            if len(texts) == 0:
                print("    跳过：没有文本")
                continue

            H = extract_representations(model_path, texts, batch_size=8)
            print(f"    表示矩阵形状：{H.shape}")

            output_path = save_representation(H, model_name, domain)
            print(f"    已保存：{output_path}")

            del H
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    print()
    print("=" * 70)
    print("GLUE 领域提取完成!")
    print("=" * 70)


if __name__ == "__main__":
    main()