import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载 tokenizer 和 模型
model_name = "/path/to/Qwen-7B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 输入文本
input_text = "Hello, world"

# Tokenize 输入文本（不添加特殊token）
input_ids = tokenizer(input_text, return_tensors="pt").input_ids
attention_mask = inputs["attention_mask"]

# 前向传播（logits）
with torch.no_grad():
    outputs = model(input_ids=input_ids)
    logits = outputs.logits  # shape: (batch_size, seq_len, vocab_size)

# 计算 log softmax 得到 log probabilities
log_probs = torch.nn.functional.log_softmax(logits, dim=-1)

# 获取每个 token 的 log probs
input_ids = input_ids[0]  # 去掉 batch 维度 (seq_len,)
log_probs = log_probs[0]  # 同样去掉 batch 维度 (seq_len, vocab_size)

# 打印每个 token 及其对应的 log probability
total_log_prob = 0.0
print(f"{'Token':<15} | {'Log Prob'}")
print("-" * 25)
for i in range(len(input_ids)):
    token_id = input_ids[i]
    token = tokenizer.decode([token_id])
    log_prob = log_probs[i][token_id].item()
    total_log_prob += log_prob
    print(f"{token:<15} | {log_prob:.4f}")

print("\nTotal log probability of 'Hello, world': {:.4f}".format(total_log_prob))

# ========== 1. Rollout Generation Phase ==========
prompt = "Question: Which is bigger? 9.11 or 9.9?"

# Generate multiple completions through parallel sampling
completions = rollout_function(
    model=current_policy_model,
    prompt=prompt,
    num_generations=8,  # Hyperparameter: number of samples per prompt
    temperature=1.0     # Hyperparameter: sampling diversity
)
"""
completions = [
    (completion 1) "The larger number is 9.11...",
    (completion 2) "9.9 is bigger than...",
    ...
    (completion 8) "After calculation, 9.11..."
]
"""

# ========== 2. Reward Calculation Phase ==========
# Evaluate generated completions using reward model
rewards = reward_function(
    completions=completions,
    ground_truth="9.11"  # Expected correct answer
)
"""
rewards = [
    (reward 1) 1.0,  # Correct answer
    (reward 2) 0.0,  # Incorrect
    ...
    (reward 8) 1.0   # Correct
]
"""

# Normalize rewards to advantages
rewards_mean = mean(rewards)  # μ = 0.5
rewards_std = std(rewards)    # σ = 0.25
advantages = (rewards - rewards_mean) / (rewards_std + 1e-8)  # Standardization
"""
advantages = [
    (advantage 1)  2.0,  # (1.0 - 0.5)/0.25
    (advantage 2) -2.0,
    ...
    (advantage 8)  2.0
]
"""

# ========== 3. Policy Optimization Phase ==========
# Get token-level log probabilities from different models
current_logps = get_per_token_logps(current_policy_model, prompt, completions)  # π_θ
old_logps = get_per_token_logps(old_policy_model, prompt, completions)          # π_θ_old
ref_logps = get_per_token_logps(reference_model, prompt, completions)           # π_ref

# PPO Clipped Objective
is_ratio = exp(current_logps - old_logps)  # Importance sampling ratio: e^(π_θ - π_θ_old)
clipped_ratio = clip(is_ratio, 1-ε, 1+ε)   # ε=0.2 typically

# Policy gradient term (dual form)
policy_loss = -mean(
    minimum(is_ratio * advantages,       # Unclipped objective
           clipped_ratio * advantages)  # Clipped objective
)

# KL Divergence Penalty (K3 estimator)
# KL(π_θ||π_ref) ≈ e^(logπ_ref - logπ_θ) - (logπ_ref - logπ_θ) - 1
kl_penalty = beta * mean(
    exp(ref_logps - current_logps) -
    (ref_logps - current_logps) - 1
)

# Total Loss = Policy Loss + KL Penalty
total_loss = policy_loss + kl_penalty

# ========== 4. Update Rule ==========
# Apply gradient descent to minimize total_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()