experiments/depth_benchmark.py

"""
Depth Comparison Benchmark - ThinkingDial 核心验证实验

目标：验证 ThinkingDial 的不同 depth 在数学推理任务上的表现差异。
这是证明 ThinkingDial 机制有效性的关键实验。

实验设计：
1. 扩大模型配置到 ~10M 参数（hidden=512, layers=8, heads=8）
2. 用合成数学数据（四则运算）预训练所有 depth
3. 用 GRPO + GSM8K 奖励函数分别在不同 depth 下微调
4. 输出 depth 准确率对比表

Usage: python experiments/depth_benchmark.py
"""
import sys
import math
import random
import time
import torch
import torch.nn.functional as F
from pathlib import Path
from typing import Dict, List, Tuple

sys.path.insert(0, str(Path(__file__).parent.parent))

from models.fusion_model import FusionModel, FusionConfig
from models.thinking_dial import (
    ThinkingDialModel, ThinkingConfig, GRPOTrainer, GRPOConfig
)
from evaluation.gsm8k_reward import GSM8KEvaluator, extract_answer, normalize_answer

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[Benchmark] Device: {DEVICE}")
print()

# ─── 实验配置 ───────────────────────────────────────────────────────────────

EXP_CONFIG = {
    # ~5M 参数模型配置（CPU 可在合理时间内完成）
    "hidden_size": 256,
    "num_hidden_layers": 6,
    "num_attention_heads": 8,
    "intermediate_size": 512,
    "vocab_size": 1000,
    "max_position_embeddings": 128,
    "block_size": 32,
    "latent_dim": 16,
    # 训练参数
    "pretrain_epochs": 100,
    "pretrain_batch_size": 8,
    "pretrain_lr": 1e-3,
    "grpo_epochs": 30,
    "grpo_batch_size": 4,
    "grpo_lr": 1e-4,
    "grpo_num_samples": 3,
    "num_thinking_depths": 4,
    "test_batch_size": 50,
    "seed": 42,
}

def estimate_params(config: FusionConfig) -> int:
    """粗略估计模型参数量"""
    emb = config.vocab_size * config.hidden_size
    per_layer = (
        # attention: Q, K, V, O projections
        4 * config.hidden_size * config.hidden_size +
        # FFN: gate, up, down
        3 * config.hidden_size * config.intermediate_size +
        # SBLA latent projections (approx)
        2 * config.hidden_size * config.latent_dim +
        # norms
        2 * config.hidden_size
    )
    return emb + config.num_hidden_layers * per_layer + config.hidden_size * config.vocab_size

# ─── 合成数学数据 ──────────────────────────────────────────────────────────

def generate_math_data(n: int = 2000, ops: List[str] = None, max_val: int = 100,
                       seed: int = 42) -> List[Tuple[str, int, int, int]]:
    """
    生成合成数学运算数据。
    格式: (x, op_code, y, result)
    op_code: 0=+, 1=-, 2=*, 3=/
    """
    if ops is None:
        ops = ["+", "-"]  # 先只做加减法（乘除需要更大模型）
    op_map = {"+": 0, "-": 1, "*": 2, "/": 3}

    rng = random.Random(seed)
    data = []

    for _ in range(n):
        op = rng.choice(ops)
        code = op_map[op]

        if op == "+":
            x = rng.randint(1, max_val)
            y = rng.randint(1, max_val)
            result = x + y
        elif op == "-":
            x = rng.randint(1, max_val * 2)
            y = rng.randint(1, x)  # 保证结果 >= 0
            result = x - y
        elif op == "*":
            x = rng.randint(2, 20)
            y = rng.randint(2, 20)
            result = x * y
        elif op == "/":
            y = rng.randint(2, 20)
            result = rng.randint(1, 20)
            x = y * result  # 保证整除

        # 结果 clamp 到 vocab 范围
        result = max(0, min(result, EXP_CONFIG["vocab_size"] - 2))

        data.append((x, code, y, result))

    return data


def encode_example(x: int, op: int, y: int, result: int) -> Tuple[List[int], List[int]]:
    """
    编码数学题:
    input:  [2, x, op, y, 0, 0, 0]  (7 tokens, padded)
    labels: [-100, x, op, y, 99, result, 1]  (7 tokens, same length)
    Loss 计算: predict y->99, 99->result, result->1
    """
    input_ids = [2, x, op, y, 0, 0, 0]
    labels = [-100, x, op, y, 99, result, 1]
    return input_ids, labels


# ─── 梯度化奖励函数 ────────────────────────────────────────────────────────

def gradient_reward_fn(prompt: str, response: str) -> float:
    """
    梯度化奖励函数（用于合成数学任务）：
    - 能提取到数字 → +0.1
    - 数值接近正确值（误差 < 10%）→ +0.3
    - 完全正确 → +1.0
    """
    extracted = extract_answer(response)
    if extracted is None:
        return 0.0

    # 查找 ground truth
    gold = None
    try:
        # 从 prompt 中提取数字: CLS x OP y
        tokens = prompt.strip().split()
        if len(tokens) >= 3:
            x, op_code, y = int(tokens[0]), int(tokens[1]), int(tokens[2])
            if op_code == 0:
                gold = x + y
            elif op_code == 1:
                gold = x - y
            elif op_code == 2:
                gold = x * y
            elif op_code == 3:
                gold = x // y if y != 0 else 0
    except (ValueError, IndexError):
        pass

    if gold is None:
        return 0.1  # 至少输出了数字

    extracted_norm = normalize_answer(extracted)
    gold_norm = normalize_answer(gold)

    if extracted_norm == gold_norm:
        return 1.0

    # 检查是否接近（误差 < 10%）
    if gold != 0:
        rel_error = abs(extracted_norm - gold_norm) / abs(gold_norm)
        if rel_error < 0.1:
            return 0.3

    # 输出了数字但答案错误
    return 0.1


# ─── 实验核心 ───────────────────────────────────────────────────────────────

def pretrain_model(config: FusionConfig, data: list) -> ThinkingDialModel:
    """阶段 1：用合成数学数据预训练"""
    model = FusionModel(config)
    thinking_config = ThinkingConfig(
        num_thinking_depths=EXP_CONFIG["num_thinking_depths"]
    )
    td_model = ThinkingDialModel(model, thinking_config)
    td_model.train().to(DEVICE)

    optimizer = torch.optim.AdamW(td_model.parameters(), lr=EXP_CONFIG["pretrain_lr"])

    losses = []
    for epoch in range(EXP_CONFIG["pretrain_epochs"]):
        batch = random.sample(data, min(EXP_CONFIG["pretrain_batch_size"], len(data)))

        input_batch = []
        target_batch = []
        for x, op, y, result in batch:
            inp, tgt = encode_example(x, op, y, result)
            input_batch.append(inp)
            target_batch.append(tgt)

        input_ids = torch.tensor(input_batch, device=DEVICE, dtype=torch.long)
        labels = torch.tensor(target_batch, device=DEVICE, dtype=torch.long)

        optimizer.zero_grad()
        outputs = td_model(input_ids, labels=labels)
        loss = outputs.loss
        loss.backward()
        torch.nn.utils.clip_grad_norm_(td_model.parameters(), 1.0)
        optimizer.step()

        losses.append(loss.item())

        if (epoch + 1) % 10 == 0:
            avg = sum(losses[-10:]) / 10
            print(f"    Pretrain Epoch {epoch+1}/{EXP_CONFIG['pretrain_epochs']}: "
                  f"Loss = {avg:.4f}")

    final_loss = sum(losses[-10:]) / 10
    print(f"  Pretrain complete. Final loss: {final_loss:.4f}")
    return td_model


def grpo_finetune(td_model: ThinkingDialModel, data: list,
                  depth: int, gsm8k_eval: GSM8KEvaluator = None) -> Dict:
    """
    阶段 2：GRPO 微调（使用合成数学数据 + 可选 GSM8K 奖励）

    Returns:
        {losses, rewards, depth}
    """
    grpo_config = GRPOConfig(
        grpo_sample_size=EXP_CONFIG["grpo_num_samples"],
        kl_coef=0.05,
    )

    # 设置奖励函数
    def reward_fn(prompt, response):
        if gsm8k_eval is not None:
            return gsm8k_eval.reward(prompt, response)
        return gradient_reward_fn(prompt, response)

    td_model.train()
    optimizer = torch.optim.AdamW(
        [p for p in td_model.parameters() if p.requires_grad],
        lr=EXP_CONFIG["grpo_lr"]
    )

    epoch_losses = []
    epoch_rewards = []

    for epoch in range(EXP_CONFIG["grpo_epochs"]):
        batch = random.sample(data, min(EXP_CONFIG["grpo_batch_size"], len(data)))
        batch_losses = []
        batch_rewards = []

        for x, op, y, result in batch:
            inp, _ = encode_example(x, op, y, result)
            input_ids = torch.tensor([inp], device=DEVICE, dtype=torch.long)

            # 用指定 depth 生成多个样本
            gen_ids = []
            for _ in range(EXP_CONFIG["grpo_num_samples"]):
                with torch.no_grad():
                    out = td_model.generate(
                        input_ids,
                        max_new_tokens=8,
                        thinking_depth=depth,
                        do_sample=True,
                        temperature=1.0,
                        pad_token_id=0,
                    )
                gen_ids.append(out[0].tolist())

            # 计算奖励
            rewards = []
            for gen in gen_ids:
                new_tokens = gen[len(inp):]
                # Extract result: expected format [99, result, 1, ...]
                if len(new_tokens) >= 2 and new_tokens[0] == 99:
                    predicted = new_tokens[1]
                else:
                    predicted = None

                # Compute reward
                if predicted is not None and predicted == result:
                    r = 1.0
                elif predicted is not None:
                    # Gradient reward: check proximity
                    if result != 0 and abs(predicted - result) / abs(result) < 0.1:
                        r = 0.3
                    else:
                        r = 0.1
                else:
                    r = 0.0
                rewards.append(r)

            # GRPO: 基于奖励计算 loss
            if len(rewards) > 1 and sum(rewards) > 0:
                # 归一化奖励
                rewards_t = torch.tensor(rewards, device=DEVICE, dtype=torch.float)
                mean_r = rewards_t.mean()
                std_r = rewards_t.std() + 1e-8
                advantages = (rewards_t - mean_r) / std_r

                # 对每个样本计算 policy gradient loss
                total_loss = torch.tensor(0.0, device=DEVICE)
                for i, gen in enumerate(gen_ids):
                    gen_t = torch.tensor([gen], device=DEVICE, dtype=torch.long)
                    # Mask labels: only compute loss on generated tokens
                    gen_labels = gen_t[:, 1:].clone()
                    gen_labels[:, :len(inp)-1] = -100  # ignore input portion
                    # Forward to get log probs
                    outputs = td_model(gen_t[:, :-1], labels=gen_labels)
                    neg_log_prob = outputs.loss
                    # Policy gradient: -advantage * log_prob
                    total_loss += advantages[i] * neg_log_prob

                loss = total_loss / len(gen_ids)

                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(td_model.parameters(), 1.0)
                optimizer.step()

                batch_losses.append(loss.item())
                batch_rewards.append(mean_r.item())
            else:
                batch_losses.append(0.0)
                batch_rewards.append(0.0)

        epoch_losses.append(sum(batch_losses) / len(batch_losses))
        epoch_rewards.append(sum(batch_rewards) / len(batch_rewards))

        if (epoch + 1) % 10 == 0:
            print(f"    GRPO Epoch {epoch+1}/{EXP_CONFIG['grpo_epochs']}: "
                  f"Loss = {epoch_losses[-1]:.4f}, "
                  f"Mean Reward = {epoch_rewards[-1]:.4f}")

    print(f"  GRPO complete. Final reward: {epoch_rewards[-1]:.4f}")
    return {
        "losses": epoch_losses,
        "rewards": epoch_rewards,
        "depth": depth,
    }


def evaluate_accuracy(td_model: ThinkingDialModel, data: list,
                      depth: int, n_samples: int = None) -> Dict:
    """
    在测试集上评估指定 depth 的准确率。
    期望生成格式: [..., 99, result, 1, ...] -> 取 position 5 的 token 作为结果
    """
    if n_samples is None:
        n_samples = EXP_CONFIG["test_batch_size"]

    test_data = random.sample(data, min(n_samples, len(data)))
    correct = 0
    total = len(test_data)
    results_by_op = {"+": {"correct": 0, "total": 0},
                     "-": {"correct": 0, "total": 0},
                     "*": {"correct": 0, "total": 0},
                     "/": {"correct": 0, "total": 0}}

    td_model.eval()
    op_names = {0: "+", 1: "-", 2: "*", 3: "/"}

    for x, op, y, result in test_data:
        inp, _ = encode_example(x, op, y, result)
        input_ids = torch.tensor([inp], device=DEVICE, dtype=torch.long)

        with torch.no_grad():
            out = td_model.generate(
                input_ids,
                max_new_tokens=8,
                thinking_depth=depth,
                do_sample=False,
                pad_token_id=0,
            )

        gen_tokens = out[0, len(inp):].tolist()
        # Expected: model generates [99, result, 1, ...]
        gen_result = None
        if len(gen_tokens) >= 2 and gen_tokens[0] == 99:
            gen_result = gen_tokens[1]
        elif gen_tokens:
            gen_result = gen_tokens[0]

        op_name = op_names.get(op, "?")
        results_by_op[op_name]["total"] += 1

        if gen_result is not None:
            if normalize_answer(gen_result) == normalize_answer(result):
                correct += 1
                results_by_op[op_name]["correct"] += 1

    accuracy = correct / total if total > 0 else 0.0

    return {
        "depth": depth,
        "accuracy": accuracy,
        "correct": correct,
        "total": total,
        "by_op": results_by_op,
    }


# ─── 消融实验 ───────────────────────────────────────────────────────────────

def create_vanilla_model(config: FusionConfig):
    """
    创建无 ThinkingDial 的 vanilla Transformer 对照组。
    直接用 FusionModel（无 thinking depth 偏置）。
    """
    model = FusionModel(config)
    model.train().to(DEVICE)
    return model


def pretrain_vanilla(model: FusionModel, data: list) -> Dict:
    """Vanilla Transformer 预训练（无 ThinkingDial）"""
    optimizer = torch.optim.AdamW(model.parameters(), lr=EXP_CONFIG["pretrain_lr"])

    losses = []
    for epoch in range(EXP_CONFIG["pretrain_epochs"]):
        batch = random.sample(data, min(EXP_CONFIG["pretrain_batch_size"], len(data)))

        input_batch = []
        target_batch = []
        for x, op, y, result in batch:
            inp, tgt = encode_example(x, op, y, result)
            input_batch.append(inp)
            target_batch.append(tgt)

        input_ids = torch.tensor(input_batch, device=DEVICE, dtype=torch.long)
        labels = torch.tensor(target_batch, device=DEVICE, dtype=torch.long)

        optimizer.zero_grad()
        outputs = model(input_ids, labels=labels)
        loss = outputs.loss
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        losses.append(loss.item())

    final_loss = sum(losses[-10:]) / 10
    return {"losses": losses, "final_loss": final_loss}


def evaluate_vanilla(model: FusionModel, data: list, n_samples: int = None) -> Dict:
    """评估 vanilla Transformer 准确率"""
    if n_samples is None:
        n_samples = EXP_CONFIG["test_batch_size"]

    test_data = random.sample(data, min(n_samples, len(data)))
    correct = 0
    total = len(test_data)

    model.eval()

    for x, op, y, result in test_data:
        inp, _ = encode_example(x, op, y, result)
        input_ids = torch.tensor([inp], device=DEVICE, dtype=torch.long)

        with torch.no_grad():
            out = model.generate(
                input_ids,
                max_new_tokens=8,
                do_sample=False,
                pad_token_id=0,
            )

        gen_tokens = out[0, len(inp):].tolist()
        gen_result = None
        if len(gen_tokens) >= 2 and gen_tokens[0] == 99:
            gen_result = gen_tokens[1]
        else:
            gen_result = extract_answer(" ".join(str(t) for t in gen_tokens))

        if gen_result is not None:
            if normalize_answer(gen_result) == normalize_answer(result):
                correct += 1

    accuracy = correct / total if total > 0 else 0.0
    return {"accuracy": accuracy, "correct": correct, "total": total}


# ─── 主实验流程 ──────────────────────────────────────────────────────────────

def run_benchmark():
    print("=" * 70)
    print("ThinkingDial Depth Comparison Benchmark")
    print("=" * 70)
    print()

    random.seed(EXP_CONFIG["seed"])
    torch.manual_seed(EXP_CONFIG["seed"])

    # 配置模型
    config = FusionConfig(
        vocab_size=EXP_CONFIG["vocab_size"],
        hidden_size=EXP_CONFIG["hidden_size"],
        num_hidden_layers=EXP_CONFIG["num_hidden_layers"],
        num_attention_heads=EXP_CONFIG["num_attention_heads"],
        intermediate_size=EXP_CONFIG["intermediate_size"],
        max_position_embeddings=EXP_CONFIG["max_position_embeddings"],
        block_size=EXP_CONFIG["block_size"],
        latent_dim=EXP_CONFIG["latent_dim"],
    )

    n_params = estimate_params(config)
    print(f"Model config: {EXP_CONFIG['hidden_size']}d x {EXP_CONFIG['num_hidden_layers']}L x {EXP_CONFIG['num_attention_heads']}H")
    print(f"Estimated parameters: ~{n_params:,}")
    print()

    # 生成训练数据
    train_data = generate_math_data(n=500, ops=["+"], max_val=20, seed=42)
    test_data = generate_math_data(n=100, ops=["+"], max_val=20, seed=999)
    print(f"Training set: {len(train_data)} examples")
    print(f"Test set: {len(test_data)} examples")
    print()

    # ─── 实验 1: Depth 对比（相同模型，不同 depth） ─────────────────────
    print("=" * 60)
    print("EXPERIMENT 1: Depth Comparison (Same Model, Different Depths)")
    print("=" * 60)
    print()

    t0 = time.time()
    print("[1/4] Pretraining ThinkingDial model...")
    td_model = pretrain_model(config, train_data)
    print(f"  Time: {time.time() - t0:.1f}s")
    print()

    # 在预训练后立即评估所有 depth（零样本）
    print("[2/4] Zero-shot accuracy by depth (before GRPO)...")
    zero_shot_results = {}
    for depth in range(EXP_CONFIG["num_thinking_depths"]):
        result = evaluate_accuracy(td_model, test_data, depth)
        zero_shot_results[depth] = result
        print(f"  Depth {depth}: {result['accuracy']*100:.1f}% "
              f"({result['correct']}/{result['total']})")
    print()

    # 对每个 depth 进行 GRPO 微调（共享预训练权重）
    print("[3/4] GRPO fine-tuning per depth...")
    grpo_results = {}
    for depth in range(EXP_CONFIG["num_thinking_depths"]):
        print(f"  --- Depth {depth} ---")
        # 每次从预训练检查点开始
        td_model_copy = pretrain_model(config, train_data)  # 重新预训练保证公平
        grpo_res = grpo_finetune(td_model_copy, train_data, depth)
        grpo_results[depth] = grpo_res

        # 评估
        result = evaluate_accuracy(td_model_copy, test_data, depth)
        grpo_results[depth]["eval"] = result
        print(f"  Post-GRPO Accuracy: {result['accuracy']*100:.1f}%")
        print()

        # 清理
        del td_model_copy
        if DEVICE == "cuda":
            torch.cuda.empty_cache()

    print("[4/4] Post-GRPO accuracy by depth...")
    post_grpo_results = {}
    for depth in range(EXP_CONFIG["num_thinking_depths"]):
        result = grpo_results[depth]["eval"]
        post_grpo_results[depth] = result
        print(f"  Depth {depth}: {result['accuracy']*100:.1f}% "
              f"({result['correct']}/{result['total']})")
    print()

    # ─── 实验 2: 消融实验（Vanilla vs ThinkingDial） ──────────────────
    print("=" * 60)
    print("EXPERIMENT 2: Ablation (Vanilla Transformer vs ThinkingDial)")
    print("=" * 60)
    print()

    print("[1/2] Pretraining vanilla Transformer...")
    t1 = time.time()
    vanilla_model = create_vanilla_model(config)
    vanilla_train = pretrain_vanilla(vanilla_model, train_data)
    print(f"  Pretrain loss: {vanilla_train['final_loss']:.4f}")
    print(f"  Time: {time.time() - t1:.1f}s")
    print()

    print("[2/2] Evaluating vanilla Transformer...")
    vanilla_result = evaluate_vanilla(vanilla_model, test_data)
    print(f"  Vanilla accuracy: {vanilla_result['accuracy']*100:.1f}% "
          f"({vanilla_result['correct']}/{vanilla_result['total']})")
    print()

    # ─── 结果汇总 ──────────────────────────────────────────────────────
    print()
    print("=" * 70)
    print("RESULTS SUMMARY")
    print("=" * 70)
    print()

    print("Table 1: Zero-shot Accuracy by Thinking Depth")
    print("-" * 50)
    print(f"{'Depth':<10} {'Accuracy':<12} {'Correct/Total':<15}")
    print("-" * 50)
    for depth in range(EXP_CONFIG["num_thinking_depths"]):
        r = zero_shot_results[depth]
        print(f"{depth:<10} {r['accuracy']*100:>6.1f}%       {r['correct']}/{r['total']}")
    print()

    print("Table 2: Post-GRPO Accuracy by Thinking Depth")
    print("-" * 50)
    print(f"{'Depth':<10} {'Accuracy':<12} {'Correct/Total':<15} {'Final Reward':<15}")
    print("-" * 50)
    for depth in range(EXP_CONFIG["num_thinking_depths"]):
        r = grpo_results[depth]["eval"]
        reward = grpo_results[depth]["rewards"][-1]
        print(f"{depth:<10} {r['accuracy']*100:>6.1f}%       "
              f"{r['correct']}/{r['total']:<10} {reward:>8.4f}")
    print()

    print("Table 3: Ablation Comparison")
    print("-" * 50)
    print(f"{'Model':<25} {'Accuracy':<12} {'Correct/Total':<15}")
    print("-" * 50)
    # ThinkingDial best depth
    best_depth = max(range(EXP_CONFIG["num_thinking_depths"]),
                     key=lambda d: post_grpo_results[d]["accuracy"])
    td_acc = post_grpo_results[best_depth]
    print(f"ThinkingDial (depth={best_depth})   {td_acc['accuracy']*100:>6.1f}%       "
          f"{td_acc['correct']}/{td_acc['total']}")
    print(f"Vanilla Transformer          {vanilla_result['accuracy']*100:>6.1f}%       "
          f"{vanilla_result['correct']}/{vanilla_result['total']}")
    print()

    # 分析
    td_best_acc = post_grpo_results[best_depth]["accuracy"]
    vanilla_acc = vanilla_result["accuracy"]

    print("Analysis:")
    if td_best_acc > vanilla_acc:
        improvement = (td_best_acc - vanilla_acc) / vanilla_acc * 100
        print(f"  ThinkingDial (depth={best_depth}) outperforms vanilla by +{improvement:.1f}%")
    elif td_best_acc == vanilla_acc:
        print(f"  ThinkingDial matches vanilla performance")
    else:
        degradation = (vanilla_acc - td_best_acc) / vanilla_acc * 100
        print(f"  ThinkingDial underperforms vanilla by -{degradation:.1f}%")
        print(f"  Note: This may indicate architecture needs adjustment or training is insufficient")

    # Depth variance analysis
    accuracies = [post_grpo_results[d]["accuracy"] for d in range(EXP_CONFIG["num_thinking_depths"])]
    if max(accuracies) > 0 and min(accuracies) < max(accuracies):
        print(f"  Depth sensitivity: accuracy ranges from "
              f"{min(accuracies)*100:.1f}% (depth {accuracies.index(min(accuracies))}) "
              f"to {max(accuracies)*100:.1f}% (depth {accuracies.index(max(accuracies))})")
        print(f"  ThinkingDial produces different outcomes at different depths")

    print()
    total_time = time.time() - t0
    print(f"Total benchmark time: {total_time:.1f}s")
    print()

    return {
        "zero_shot": zero_shot_results,
        "post_grpo": post_grpo_results,
        "grpo": grpo_results,
        "vanilla": vanilla_result,
    }


if __name__ == "__main__":
    results = run_benchmark()