import torch
from transformers import AutoProcessor
from modeling_interns1_pro import InternS1ProForConditionalGeneration

# 加载模型
model_path = "."  # 当前目录
model = InternS1ProForConditionalGeneration.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)

# 简单的文本输入测试
input_ids = torch.tensor([[1, 2, 3, 4, 5]]).to(model.device)
attention_mask = torch.ones_like(input_ids)

# 测试 1: 不请求 router_logits
print("=" * 50)
print("测试 1: 不请求 output_router_logits")
seq_len = input_ids.shape[1]
cache_position = torch.arange(seq_len, device=model.device)
outputs = model(
    input_ids=input_ids,
    attention_mask=attention_mask,
    cache_position=cache_position,
)
print(f"outputs keys: {outputs.keys() if hasattr(outputs, 'keys') else dir(outputs)}")
print(f"outputs.aux_loss: {outputs.aux_loss}")

# 测试 2: 请求 router_logits
print("=" * 50)
print("测试 2: 请求 output_router_logits=True")
outputs = model(
    input_ids=input_ids,
    attention_mask=attention_mask,
    output_router_logits=True,
    cache_position=cache_position,
)
print(f"outputs.aux_loss: {outputs.aux_loss}")

# 检查是否有 router_logits 属性
if hasattr(outputs, 'router_logits'):
    router_logits = outputs.router_logits
    print(f"router_logits type: {type(router_logits)}")
    if router_logits is not None:
        if isinstance(router_logits, tuple):
            print(f"router_logits length: {len(router_logits)}")
            for i, rl in enumerate(router_logits):
                if rl is not None:
                    print(f"  layer {i}: shape={rl.shape}, dtype={rl.dtype}")
                    print(f"    min={rl.min().item():.4f}, max={rl.max().item():.4f}, mean={rl.mean().item():.4f}")
        else:
            print(f"router_logits shape: {router_logits.shape}")
    else:
        print("router_logits is None")
else:
    print("outputs 没有 router_logits 属性")

# 测试 3: 手动检查 MoE 层的 gate
print("=" * 50)
print("测试 3: 手动检查 MoE 层")
moe_layer_indices = []
for i, layer in enumerate(model.model.language_model.layers):
    mlp = layer.mlp
    mlp_class_name = mlp.__class__.__name__
    print(f"Layer {i}: mlp type = {mlp_class_name}")
    if "SparseMoe" in mlp_class_name or "Moe" in mlp_class_name:
        moe_layer_indices.append(i)
        if hasattr(mlp, 'gate'):
            print(f"  gate: {mlp.gate}")
            print(f"  gate type: {type(mlp.gate)}")

print(f"\nMoE layers: {moe_layer_indices}")
print(f"Total MoE layers: {len(moe_layer_indices)}")

# 测试 4: 手动 forward 一个 MoE 层并获取 router_logits
print("=" * 50)
print("测试 4: 手动获取 router_logits")
if moe_layer_indices:
    layer_idx = moe_layer_indices[0]
    moe_block = model.model.language_model.layers[layer_idx].mlp

    # 创建假输入
    hidden_size = model.config.text_config.hidden_size
    fake_hidden = torch.randn(1, 5, hidden_size, dtype=torch.bfloat16, device=model.device)

    # 手动计算 router_logits
    # hidden_flat = fake_hidden.reshape(-1, hidden_size)
    # router_logits_manual = moe_block.gate(hidden_flat)
    router_logits_manual, router_scores, router_indices = moe_block.gate(fake_hidden)
    print(f"手动计算的 router_logits shape: {router_logits_manual.shape}")
    print(f"  应该是 (batch*seq_len, num_experts) = (5, {model.config.text_config.num_experts})")
    print(f"router_scores shape: {router_scores.shape}")
    print(f"router_indices shape: {router_indices.shape}")