# %%
# 全局常量

DIR = "../models/HuggingFace/Qwen3-0.6B"

# %%
# 处理tokenize逻辑

from tokenizers import Tokenizer

tokenizer = Tokenizer.from_file(DIR + "/tokenizer.json")
output = tokenizer.encode("测试一下tokenizer的效果")
print(output.ids)
tokenizer.decode(output.ids)

# %%
# 处理参数和超参数

import json

import torch
from safetensors import safe_open

# 常数和超参数
eos_token_id: int  # 用于判断模型输出完的token id
rms_norm_eps: float  # RMSNorm防止归一化分母为0的eps
rope_theta: int  # RoPE的theta的底数部分
num_hidden_layers: int  # transformer层数
num_attention_heads: int  # Q的注意力头数
num_key_value_heads: int  # KV的注意力头数
head_dim: int  # 单头注意力维数
max_position_embeddings: int  # 允许位置编码的最大序列长度

# 参数
E: torch.Tensor  # 嵌入层
Gamma_finalRMSNorm: torch.Tensor  # 最终RMSNorm的缩放系数
W_output: torch.Tensor  # 最终投影到词表大小的线性层
Gamma_inputRMSNorm: list[torch.Tensor] = []  # transformer层中对输入的RMSNorm的缩放系数
Gamma_postATTNRMSNorm: list[torch.Tensor] = []  # transformer层中对自注意力结果的RMSNorm的缩放系数
W_FFNGate: list[torch.Tensor] = []  # transformer层中FFN的门控层
W_FFNUp: list[torch.Tensor] = []  # transformer层中FFN的升维层
W_FFNDown: list[torch.Tensor] = []  # transformer层中FFN的降维层
Gamma_QRMSNorm: list[torch.Tensor] = []  # transformer层中Q投影后、RoPE前的每个头的RMSNorm的缩放系数
Gamma_KRMSNorm: list[torch.Tensor] = []  # transformer层中K投影后、RoPE前的每个头的RMSNorm的缩放系数
W_Q: list[torch.Tensor] = []  # transformer层中Q投影矩阵
W_K: list[torch.Tensor] = []  # transformer层中K投影矩阵
W_V: list[torch.Tensor] = []  # transformer层中V投影矩阵
W_O: list[torch.Tensor] = []  # transformer层中O投影矩阵

# 读取常数和超参数
with open(DIR + "/config.json", "r", encoding="utf-8") as f:
    data = json.load(f)

    def readParam(name, t):
        d = data.get(name)
        if not isinstance(d, t):
            raise ValueError(f"get {name} error")
        return d

    eos_token_id = readParam("eos_token_id", int)  # 151645
    rms_norm_eps = readParam("rms_norm_eps", float)  # 1e-06
    rope_theta = readParam("rope_theta", int)  # 1000000
    num_hidden_layers = readParam("num_hidden_layers", int)  # 28
    num_attention_heads = readParam("num_attention_heads", int)  # 16
    num_key_value_heads = readParam("num_key_value_heads", int)  # 8
    head_dim = readParam("head_dim", int)  # 128
    max_position_embeddings = readParam("max_position_embeddings", int)  # 40960

# 读取参数
with safe_open(DIR + "/model.safetensors", framework="pt", device=0) as f:
    E = f.get_tensor("model.embed_tokens.weight")  # torch.Size([151936, 1024])
    Gamma_finalRMSNorm = f.get_tensor("model.norm.weight")  # torch.Size([1024])
    W_output = f.get_tensor("lm_head.weight")  # torch.Size([151936, 1024])
    for i in range(num_hidden_layers):
        Gamma_inputRMSNorm.append(f.get_tensor(f"model.layers.{i}.input_layernorm.weight"))  # torch.Size([1024])
        Gamma_postATTNRMSNorm.append(f.get_tensor(f"model.layers.{i}.post_attention_layernorm.weight"))  # torch.Size([1024])
        W_FFNGate.append(f.get_tensor(f"model.layers.{i}.mlp.gate_proj.weight"))  # torch.Size([3072, 1024])
        W_FFNUp.append(f.get_tensor(f"model.layers.{i}.mlp.up_proj.weight"))  # torch.Size([3072, 1024])
        W_FFNDown.append(f.get_tensor(f"model.layers.{i}.mlp.down_proj.weight"))  # torch.Size([1024, 3072])
        Gamma_QRMSNorm.append(f.get_tensor(f"model.layers.{i}.self_attn.q_norm.weight"))  # torch.Size([128])
        Gamma_KRMSNorm.append(f.get_tensor(f"model.layers.{i}.self_attn.k_norm.weight"))  # torch.Size([128])
        W_Q.append(f.get_tensor(f"model.layers.{i}.self_attn.q_proj.weight"))  # torch.Size([2048, 1024])
        W_K.append(f.get_tensor(f"model.layers.{i}.self_attn.k_proj.weight"))  # torch.Size([1024, 1024])
        W_V.append(f.get_tensor(f"model.layers.{i}.self_attn.v_proj.weight"))  # torch.Size([1024, 1024])
        W_O.append(f.get_tensor(f"model.layers.{i}.self_attn.o_proj.weight"))  # torch.Size([1024, 2048])

model_dtype = E.dtype

# %%
# 工具函数


def RMSNorm(X, gamma):
    return X / torch.sqrt(torch.mean(X[..., :] ** 2, -1, keepdim=True) + rms_norm_eps) * gamma


# interleaved: 交错配对theta[0,0,1,1,2,2,3,3]
# rope_theta_repeat2 = torch.pow(rope_theta, -torch.arange(head_dim // 2) * 2 / head_dim).repeat_interleave(2)
# rotate_half: 前后两半配对theta_0,1,2,3,0,1,2,3
rope_theta_repeat2 = torch.pow(rope_theta, -torch.arange(head_dim // 2) * 2 / head_dim).repeat(2)
rope_m_theta = torch.arange(max_position_embeddings).reshape(-1, 1) * rope_theta_repeat2.reshape(1, -1)
rope_sin_m_theta = torch.sin(rope_m_theta).cuda()  # torch.Size([40960, 128])
rope_cos_m_theta = torch.cos(rope_m_theta).cuda()  # torch.Size([40960, 128])


def RoPE(X, offset=0):
    Ls = X.shape[-2]
    d = X.shape[-1]
    cos, sin = rope_cos_m_theta[offset : offset + Ls], rope_sin_m_theta[offset : offset + Ls]
    # interleaved: 交错配对x[-1,0,-3,2]
    # X_rot = torch.stack((-X[..., 1::2], X[..., 0::2]), dim=-1).flatten(-2)
    # rotate_half: 前后两半配对x[-0,-1,-2,-3,0,1,2,3]
    X_rot = torch.cat((-X[..., d // 2 :], X[..., : d // 2]), dim=-1)
    return (X * cos + X_rot * sin).to(model_dtype)


# RoPE(torch.arange(2 * 3 * head_dim).reshape(2, 3, head_dim).cuda(), 0)


def SiLU(X):
    return X / (1 + torch.exp(-X))


# SiLU(torch.arange(-10, 10).reshape(2, 10))


def softmax(X, temperature=1):
    X_shifted = X - torch.max(X / temperature, dim=-1, keepdims=True).values
    exp_X = torch.exp(X_shifted)
    return exp_X / torch.sum(exp_X, dim=-1, keepdims=True)


# softmax(torch.arange(-10, 10).reshape(2, 10))


# 右上方三角被mask
def masked_softmax(X):
    # X: (batch_size, num_heads, num_steps, num_steps)
    maxlen = X.size(-1)
    mask = (torch.arange(maxlen).reshape(-1, 1) >= torch.arange(maxlen)).to(X.device)
    mask_not = ~mask
    return softmax(X * mask + (mask_not * -1e6).to(X.dtype))


# masked_softmax(torch.arange(96).reshape(2, 3, 4, 4))

# %%
# 核心计算函数与Transformer层


def FFN(X, gate, up, down):
    return (SiLU(X @ gate.T) * (X @ up.T)) @ down.T


# 通过转置实现Q/K/V张量按注意力头的拆分与恢复，便于一次性计算所有注意力头
def transpose_qkv(X, num_heads):
    # X_h: (batch_size, num_steps, num_heads, num_hiddens/num_heads)
    X_h = X.reshape(*X.shape[:-1], num_heads, -1)
    return X_h.transpose(-2, -3)


def transpose_output(X):
    # X_t: (batch_size, num_steps, num_heads, num_hiddens/num_heads)
    X_t = X.transpose(-2, -3)
    return X_t.reshape(*X_t.shape[:-2], -1)


# print(transpose_qkv(torch.arange(80).reshape(2, 5, 8), 2).shape)
# print(transpose_output(transpose_qkv(torch.arange(80).reshape(2, 5, 8), 2)).shape)


def DotProductAttention(Q, K, V, offset=0):
    KVScale = num_attention_heads // num_key_value_heads
    # Q/K/V: (batch_size, num_heads, num_steps, num_hiddens/num_heads)
    # scores/A: (batch_size, num_heads, num_steps, num_steps)
    scores = Q @ torch.repeat_interleave(K.transpose(-1, -2), KVScale, -3) / torch.sqrt(torch.tensor(Q.shape[-1]))
    A = softmax(scores) if offset else masked_softmax(scores)
    return A @ torch.repeat_interleave(V, KVScale, -3)


KVCache: list[list[torch.Tensor]]


def ClearKVCache():
    global KVCache
    KVCache = [None] * num_hidden_layers


# KVCache中已经prefill或predict了多少个字符
def GetOffset():
    global KVCache
    return KVCache[0][0].size(-2)


# offset为0代表prefill阶段
def TransfomerLayer(layer, X, Gamma_inputRMSNorm, Gamma_postATTNRMSNorm, W_FFNGate, W_FFNUp, W_FFNDown, Gamma_QRMSNorm, Gamma_KRMSNorm, W_Q, W_K, W_V, W_O, offset=0):
    # X: (batch_size, num_steps, num_hiddens)
    X_inputRMSNorm = RMSNorm(X, Gamma_inputRMSNorm)
    # 投影 -> QK RMSNorm -> RoPE -> attention
    # Q/K/V: (batch_size, num_steps, num_heads, num_hiddens) -> (batch_size, num_heads, num_steps, num_hiddens/num_heads)
    Q = transpose_qkv(X_inputRMSNorm @ W_Q.T, num_attention_heads)
    K = transpose_qkv(X_inputRMSNorm @ W_K.T, num_key_value_heads)
    V = transpose_qkv(X_inputRMSNorm @ W_V.T, num_key_value_heads)
    Q_RMSNorm = RMSNorm(Q, Gamma_QRMSNorm)
    K_RMSNorm = RMSNorm(K, Gamma_KRMSNorm)
    Q_RoPE = RoPE(Q_RMSNorm, offset)
    K_RoPE = RoPE(K_RMSNorm, offset)
    # prefill阶段初始化KVCache
    global KVCache
    if offset == 0:
        KVCache[layer] = [K_RoPE, V]
    else:
        KVCache[layer][0] = torch.cat((KVCache[layer][0], K_RoPE), dim=-2)
        KVCache[layer][1] = torch.cat((KVCache[layer][1], V), dim=-2)
    # V_attn: (batch_size, num_heads, num_steps, num_hiddens/num_heads)
    V_attn = DotProductAttention(Q_RoPE, KVCache[layer][0], KVCache[layer][1], offset)
    X2 = X + transpose_output(V_attn) @ W_O.T
    # FFN前的RMSNorm -> FFN
    # O: (batch_size, num_steps, num_hiddens)
    O = X2 + FFN(RMSNorm(X2, Gamma_postATTNRMSNorm), W_FFNGate, W_FFNUp, W_FFNDown)
    return O

# %%
# prefill与predict


def prefill(X):
    ClearKVCache()
    # X: (batch_size, num_steps)
    # O: (batch_size, num_steps, num_hiddens)
    O = E[X]
    for i in range(num_hidden_layers):
        O = TransfomerLayer(
            i, O, Gamma_inputRMSNorm[i], Gamma_postATTNRMSNorm[i], W_FFNGate[i], W_FFNUp[i], W_FFNDown[i], Gamma_QRMSNorm[i], Gamma_KRMSNorm[i], W_Q[i], W_K[i], W_V[i], W_O[i]
        )
    return O


# 多层Transformer得到的结果经过最终的RMSNorm、线性层和softmax得到概率
def lastOutputToken(O, temperature=1):
    O = RMSNorm(O, Gamma_finalRMSNorm)
    Y = O @ W_output.T
    P = softmax(Y, temperature)
    return torch.argmax(P, dim=-1)[..., -1:].to(torch.long)


def predict(X, temperature=1):
    offset = GetOffset()
    # O: (batch_size, 1)
    res = X.clone().cpu()
    while offset < max_position_embeddings:
        token_id = res[0][-1].numpy()
        print(tokenizer.decode([token_id], skip_special_tokens=False), end="")
        if token_id == eos_token_id:
            break
        O = E[X]
        for i in range(num_hidden_layers):
            O = TransfomerLayer(
                i, O, Gamma_inputRMSNorm[i], Gamma_postATTNRMSNorm[i], W_FFNGate[i], W_FFNUp[i], W_FFNDown[i], Gamma_QRMSNorm[i], Gamma_KRMSNorm[i], W_Q[i], W_K[i], W_V[i], W_O[i], offset
            )
        X = lastOutputToken(O, temperature)
        res = torch.cat((res, X.cpu()), dim=-1)
        offset += 1
    print()
    return res.numpy()


prefill(torch.ones(1, 5, dtype=torch.long))

# %%
# 构造输入

import jinja2

with open(DIR + "/tokenizer_config.json", "r", encoding="utf-8") as f:
    data = json.load(f)
    chat_template = data["chat_template"]

tpl = jinja2.Environment().from_string(chat_template)

messages = [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "steam现在最火的游戏是什么"}]

rendered = tpl.render(
    messages=messages,
    add_generation_prompt=True,
    enable_thinking=True,
)
tokens = tokenizer.encode(rendered).ids

print(rendered)
print(tokens)

# %%
# 发起预测

import time

batch_size = 50

start = time.perf_counter()
O = prefill(torch.tensor(tokens, dtype=torch.long).reshape(1, -1).repeat(batch_size, 1))
end = time.perf_counter()
print(f"prefill耗时: {end - start:.2f} 秒")

start = time.perf_counter()
X = lastOutputToken(O)
res = predict(X)
end = time.perf_counter()
print(f"predict: {len(res[0])*batch_size/(end - start):.2f} tokens/s")
print(res)