import torch
import sys
import math
from modelscope import snapshot_download
from transformers import AutoTokenizer
from transformers import AutoConfig

from modeling_qwen import QWenLMHeadModel
from modeling_qwen import QwenRunner
import numpy as np

import torch.nn.functional as F

from qwen_generation_utils import (
    make_context,
    decode_tokens,
)

sys.path.append("..")
from tools import show

seed = 4321
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

model_dir = snapshot_download("qwen/Qwen-1_8B-Chat")
# model_dir = "/home/colin/.cache/modelscope/hub/qwen/Qwen-1_8B-Chat"

config, kwargs = AutoConfig.from_pretrained(
    "./",
    return_unused_kwargs=True,
    trust_remote_code=True,
    code_revision=None,
    _commit_hash=None,
)
model = QWenLMHeadModel(config)

print(model)

tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

model = model.from_pretrained(model_dir)
if torch.cuda.device_count() > 0:
    model = model.cuda()
model = model.eval()

index = 0


class ResearchRunner(QwenRunner):
    def __init__(self, model):
        super().__init__(model)

    def prepareInput(self, tokenizer, query, query_assistant, history, system):
        start_to = [151644]
        n_to = [198]
        end_to = [151645]
        system_str = "system\nYou are a helpful assistant."
        user_str = "user\n" + query
        aassistant_str = "assistant\n" + query_assistant

        system_token = start_to + tokenizer.encode(system_str, allowed_special=set()) + end_to + n_to
        user_token = start_to + tokenizer.encode(user_str, allowed_special=set()) + end_to + n_to
        aassistant_token = start_to + tokenizer.encode(aassistant_str, allowed_special=set())

        tokens = system_token + user_token + aassistant_token
        tokens = user_token + aassistant_token
        tokens = start_to + tokenizer.encode("user\nHi你好\nassistant\n我是", allowed_special=set())

        return "", tokens

    def forwardQWenBlock(
        self,
        block,
        hidden_states,
        rotary_pos_emb_list=None,
    ):
        layernorm_output = block.ln_1(hidden_states)

        attn_outputs = self.forwardAttention(block.attn, layernorm_output, rotary_pos_emb_list)
        attn_output = attn_outputs[0]
        layernorm_input = attn_output + hidden_states

        layernorm_output = block.ln_2(layernorm_input)
        a1 = block.mlp.w1(layernorm_output)
        a2 = block.mlp.w2(layernorm_output)

        activation = (F.relu(a2) > 0).to(float)
        act_mean = torch.mean(activation, 2)
        print("Layer:" + str(block.index))
        print(act_mean.cpu())
        global index
        if index == 0:
            activation = activation.reshape(activation.shape[1], 64, -1)
            show.DumpTensorToImage(activation, "./temp/activation_layer_" + str(block.index) + ".png")

        intermediate_parallel = a1 * F.silu(a2)
        mlp_output = block.mlp.c_proj(intermediate_parallel)

        hidden_states = layernorm_input + mlp_output
        return hidden_states

    def isFinish(self, next_tokens):
        global index
        index = index + 1
        finish, next = super().isFinish(next_tokens)
        return finish, next


para = list(model.parameters())
runner = ResearchRunner(model)

output_ids, history, decoded = runner.Chat(tokenizer, "你好!!", "")
print(decoded)


tokens = []
for i, token in enumerate(output_ids):
    de = tokenizer.decode([token])
    de = str(i + 1).zfill(3) + " : " + repr(de)
    tokens.append(de)

show.DumpListToFile(tokens, "./temp/token_decode_list.txt")