import numpy as np
from transformer import SimpleTransformer, InferenceEngine, create_padding_mask, create_look_ahead_mask


def demo_training_forward():
    """演示训练时的前向传播（无KV cache）"""
    print("=" * 60)
    print("训练前向传播演示（无KV cache）")
    print("=" * 60)
    
    # 模型超参数配置
    vocab_size = 1000
    d_model = 512
    num_heads = 8
    num_layers = 6
    d_ff = 2048
    max_seq_len = 100
    
    # 创建模型
    model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len)
    
    # 创建随机输入
    batch_size = 2
    seq_len = 10
    x = np.random.randint(1, vocab_size, (batch_size, seq_len))
    
    print(f"输入形状: {x.shape}")
    print(f"输入示例: {x[0]}")
    
    # 训练模式前向传播（不使用KV cache）
    output = model.forward(x)
    
    print(f"输出形状: {output.shape}")
    print(f"输出示例（前5个值）: {output[0, 0, :5]}")
    print()


def demo_prefill_decode():
    """演示推理时的prefill和decode阶段"""
    print("=" * 60)
    print("推理演示（Prefill + Decode）")
    print("=" * 60)
    
    # 模型超参数配置
    vocab_size = 1000
    d_model = 128  # 使用较小的模型以便快速演示
    num_heads = 4
    num_layers = 2
    d_ff = 512
    max_seq_len = 100
    
    # 创建模型和推理引擎
    model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len)
    engine = InferenceEngine(model)
    
    # 模拟输入序列（prompt）
    batch_size = 1
    prompt_len = 5
    prompt = np.random.randint(1, vocab_size, (batch_size, prompt_len))
    
    print(f"输入prompt: {prompt[0]}")
    print(f"Prompt长度: {prompt_len}")
    print()
    
    # Prefill阶段
    print("--- Prefill阶段 ---")
    print("处理完整prompt，初始化KV cache")
    logits = engine.prefill(prompt)
    print(f"Prefill输出logits形状: {logits.shape}")
    print(f"当前位置: {engine.current_position}")
    
    # 采样第一个token
    first_token = engine._sample(logits, temperature=0.8)
    print(f"采样的第一个token: {first_token[0, 0]}")
    print()
    
    # Decode阶段
    print("--- Decode阶段 ---")
    print("逐个生成新token（使用KV cache）")
    
    generated_tokens = [first_token[0, 0]]
    current_token = first_token
    
    # 生成5个token作为演示
    for i in range(5):
        logits = engine.decode_step(current_token)
        current_token = engine._sample(logits, temperature=0.8)
        generated_tokens.append(current_token[0, 0])
        print(f"Step {i+1}: 生成token {current_token[0, 0]}, 位置 {engine.current_position}")
    
    print()
    print(f"完整生成序列: {generated_tokens}")
    print()


def demo_auto_regressive():
    """演示自回归生成"""
    print("=" * 60)
    print("自回归生成演示")
    print("=" * 60)
    
    # 模型超参数配置
    vocab_size = 1000
    d_model = 128
    num_heads = 4
    num_layers = 2
    d_ff = 512
    max_seq_len = 100
    
    # 创建模型和推理引擎
    model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len)
    engine = InferenceEngine(model)
    
    # 模拟输入序列
    batch_size = 1
    prompt_len = 3
    prompt = np.random.randint(1, vocab_size, (batch_size, prompt_len))
    
    print(f"输入prompt: {prompt[0]}")
    
    # 使用generate方法进行自回归生成
    output = engine.generate(prompt, max_new_tokens=10, temperature=0.8)
    
    print(f"生成的完整序列: {output[0]}")
    print(f"序列长度: {len(output[0])} (原始{prompt_len} + 生成{len(output[0]) - prompt_len})")
    print()


def demo_kv_cache_comparison():
    """对比有无KV cache的计算差异"""
    print("=" * 60)
    print("KV Cache计算对比")
    print("=" * 60)
    
    print("无KV cache（训练模式）:")
    print("  - 每次forward处理完整序列")
    print("  - 计算量: O(n^2) 每次")
    print("  - 适用于训练")
    print()
    
    print("有KV cache（推理模式）:")
    print("  - Prefill: 处理完整prompt，缓存KV")
    print("  - Decode: 只处理新token，使用缓存")
    print("  - 计算量: Prefill O(n^2), Decode O(n) 每步")
    print("  - 适用于推理")
    print()


if __name__ == "__main__":
    # 演示训练前向传播
    demo_training_forward()
    
    # 演示prefill和decode
    demo_prefill_decode()
    
    # 演示自回归生成
    demo_auto_regressive()
    
    # 对比说明
    demo_kv_cache_comparison()