# Copyright 2025 NVIDIA CORPORATION & AFFILIATES
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0
# Modified from LLaDA repos: https://github.com/ML-GSAI/LLaDA

'''
This file is inspired by the code from https://github.com/ML-GSAI/SMDM
'''
import accelerate
import torch
import re
from pathlib import Path
import random
import numpy as np
import torch.nn.functional as F
from datasets import Dataset
from lm_eval.__main__ import cli_evaluate
from lm_eval.api.instance import Instance
from lm_eval.api.model import LM
from lm_eval.api.registry import register_model
from tqdm import tqdm
import os
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import json
import time
import types
import modeling_dfm

def set_seed(seed):
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


@register_model("dfm_qwen")
class DFMQwenEvalHarness(LM):
    def __init__(
        self,
        model_path='',
        device="cuda",
        is_instruct=True,
        show_speed=False,
        max_new_tokens=2048,
        temperature=0.0,
        do_sample=False,
        batch_size=32,
        mask_id=151665,
        use_block_cache=False,
        small_block_size=16,
        bd_size=32,
        mc_num=1,
        cfg=0.0,
        dynamic_threshold=False,
        factor=1,
        threshold=0.9,
        token_shift=True,
        **kwargs,
    ):

        super().__init__()

        accelerator = accelerate.Accelerator()
        if accelerator.num_processes > 1:
            self.accelerator = accelerator
        else:
            self.accelerator = None
        
        model_kwargs = {}
        if self.accelerator is not None:
            model_kwargs.update({'device_map': {'': f'{self.accelerator.device}'}})
        
        # 使用AutoModelForCausalLM替代AutoModel
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path, 
            trust_remote_code=True, 
            torch_dtype=torch.bfloat16, 
            **model_kwargs
        )
        self.model.eval()

        # 绑定mdm_sample方法，匹配chatbot逻辑
        if token_shift:
            self.model.mdm_sample = types.MethodType(modeling_dfm.DFM_QwenForCausalLM.batch_sample, self.model)
        else:
            self.model.mdm_sample = types.MethodType(modeling_dfm.DFM_QwenForCausalLM.batch_sample_no_token_shift, self.model)

        self.device = torch.device(device)
        if self.accelerator is not None:
            self.model = self.accelerator.prepare(self.model)
            self.device = torch.device(f'{self.accelerator.device}')
            self._rank = self.accelerator.local_process_index
            self._world_size = self.accelerator.num_processes
        else: 
            self.model = self.model.to(device)
            self._rank = 0
            self._world_size = 1

        self.tokenizer = AutoTokenizer.from_pretrained("models/Fast_dLLM_1.5B", trust_remote_code=True)

        
        # 添加缺失的属性
        self.is_instruct = is_instruct
        self.show_speed = show_speed
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.do_sample = do_sample
        self.batch_size = int(batch_size)
        self.mask_id = mask_id
        self.model_path = model_path
        self.use_block_cache = use_block_cache
        self.small_block_size = small_block_size
        self.dynamic_threshold = dynamic_threshold
        self.factor = factor
        self.threshold = threshold
        self.mc_num = mc_num
        self.cfg = cfg
        self.bd_size = bd_size
        self.token_shift = token_shift
    @property
    def rank(self):
        return self._rank
    
    @property
    def world_size(self):
        return self._world_size

    @property
    def tokenizer_name(self):
        return self.model_path
    
    def apply_chat_template(self, chat_history, add_generation_prompt=True):
        return self.tokenizer.apply_chat_template(chat_history, add_generation_prompt=add_generation_prompt, tokenize=False)
    
    def loglikelihood_rolling(self, requests):
        raise NotImplementedError
    
    def _encode_pair(self, context, continuation):
        whole_enc = self.tokenizer(context + continuation)["input_ids"]
        context_enc = self.tokenizer(context)["input_ids"]

        context_enc_len = len(context_enc)
        continuation_enc = whole_enc[context_enc_len:]

        return context_enc, continuation_enc


    def _forward_process(self, batch, prompt_index):
        b, l = batch.shape

        target_len = (l - prompt_index.sum()).item()
        # k = torch.randint(1, target_len + 1, (), device=batch.device)

        # x = torch.round(torch.linspace(float(k), k + (b - 1) * (target_len / b), steps=b, device=batch.device)).long()
        # x = ((x - 1) % target_len) + 1
        # assert x.min() >= 1 and x.max() <= target_len

        # indices = torch.arange(target_len, device=batch.device).repeat(b, 1)
        # is_mask = indices < x.unsqueeze(1)

        # for i in range(b):
        #     is_mask[i] = is_mask[i][torch.randperm(target_len)]
        # is_mask = torch.cat((torch.zeros(b, prompt_index.sum(), dtype=torch.bool, device=batch.device), is_mask), dim=1)

        is_mask = torch.zeros_like(batch, dtype=torch.bool, device=batch.device)
        is_mask[:, prompt_index.sum()] = True


        noisy_batch = torch.where(is_mask, self.mask_id, batch)
        noisy_batch = torch.cat([noisy_batch, torch.full((b, self.bd_size-noisy_batch.shape[1]%self.bd_size), self.mask_id, dtype=torch.long, device=batch.device)], dim=1)
        if noisy_batch.shape[1] > l:
            noisy_batch[:, l] = self.tokenizer.eos_token_id
        # if noisy_batch.shape[1] > l+1:
        #     noisy_batch[:, l+1] = 198
        # l = noisy_batch.shape[1]
        # import pdb; pdb.set_trace()
        # return noisy_batch, (x / target_len).unsqueeze(1).repeat(1, l)
        return noisy_batch

    @torch.no_grad()
    def get_logits(self, batch, prompt_index):
        if self.cfg > 0.:
            assert len(prompt_index) == batch.shape[1]
            prompt_index = prompt_index.unsqueeze(0).repeat(batch.shape[0], 1)
            un_batch = batch.clone()
            un_batch[prompt_index] = self.mask_id
            batch = torch.cat([batch, un_batch])

        logits = self.model(batch).logits

        if self.cfg > 0.:
            logits, un_logits = torch.chunk(logits, 2, dim=0)
            logits = un_logits + (self.cfg + 1) * (logits - un_logits)
        if self.token_shift:
            logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
        return logits[:, :batch.shape[1]]

    @torch.no_grad()
    def get_loglikelihood(self, prefix, target):
        seq = torch.concatenate([prefix, target])[None, :]
        seq = seq.repeat((self.batch_size, 1)).to(self.device)

        prompt_index = torch.arange(seq.shape[1], device=self.device) < len(prefix)

        loss_acc = []
        for _ in range(self.mc_num // self.batch_size):
            # perturbed_seq, p_mask = self._forward_process(seq, prompt_index)
            perturbed_seq = self._forward_process(seq, prompt_index)

            mask_indices = perturbed_seq == self.mask_id

            logits = self.get_logits(perturbed_seq, prompt_index)
            seq = torch.cat([seq, torch.full((seq.shape[0], self.bd_size-seq.shape[1]%self.bd_size), -100, dtype=torch.long, device=self.device)], dim=1)
            # loss = F.cross_entropy(logits[mask_indices], seq[mask_indices], reduction='none') / p_mask[mask_indices]
            loss = F.cross_entropy(logits[mask_indices], seq[mask_indices], reduction='none')
            loss = loss.sum() / self.batch_size
            loss_acc.append(loss.item())

        return - sum(loss_acc) / len(loss_acc)


    def loglikelihood(self, requests):
        def _tokenize(e):
            prefix, target = self._encode_pair(e["prefix"], e["target"])
            return {
                "prefix_text": e["prefix"],
                "target_text": e["target"],
                "prefix": prefix,
                "target": target,
            }

        ds = []
        ds = [{"prefix": req.args[0], "target": req.args[1]} for req in requests]
        ds = Dataset.from_list(ds)
        ds = ds.map(_tokenize)
        ds = ds.with_format("torch")
        prompt_len = [len(x["prefix"]) + len(x["target"]) for x in ds]

        assert max(prompt_len) <= 4096

        out = []
        with torch.no_grad():
            for elem in tqdm(ds, desc="Computing likelihood..."):
                prefix = elem["prefix"]
                target = elem["target"]
                # assert len(target) == 2, f"target {elem['target_text']} length is {len(target)}"

                ll = self.get_loglikelihood(prefix, target)
                out.append((ll, 0.0))
        torch.cuda.empty_cache()
        return out
    
    def generate_until(self, requests):
        output = [None] * len(requests)  # 预分配输出列表
        num_tokens = 0
        num_nfe = 0
        
        start_time = time.time()
        
        # 按输入长度排序，并记录原始索引
        requests_with_indices = [(i, req) for i, req in enumerate(requests)]
        requests_with_indices.sort(key=lambda x: len(x[1].args[0]))
        
        # 重新组织批处理
        batched_requests = []
        current_batch = []
        for i, req in requests_with_indices:
            current_batch.append((i, req))
            if len(current_batch) == self.batch_size:
                batched_requests.append(current_batch)
                current_batch = []
        
        if current_batch:  # 处理最后一个不完整的批次
            batched_requests.append(current_batch)

        for batch_idx, batch in enumerate(tqdm(batched_requests, desc="Generating...")):
            batched_input_ids = []
            max_len = 0
            min_len = 1e9
            seq_len = []
            
            for orig_idx, req in batch:
                # 根据chatbot逻辑处理输入
                question = req.args[0]
                # messages = [{"role": "user", "content": question}]
                # user_input = self.tokenizer.apply_chat_template(
                #     messages, 
                #     add_generation_prompt=True, 
                #     tokenize=False
                # )
                
                # 对输入进行tokenize
                if req.task_name.startswith('minerva_math'):
                    question = question.replace("Solution:", "Please reason step by step, and put your final answer within \\boxed{{}}.")
                elif req.task_name.startswith('gsm8k'):
                    question = question.replace("Answer:", "Please reason step by step, and put your final answer within \\boxed{{}}.")
                model_inputs = self.tokenizer([question], return_tensors="pt").to(self.device)
                batched_input_ids.append(model_inputs["input_ids"])
                max_len = max(max_len, model_inputs["input_ids"].shape[1])
                min_len = min(min_len, model_inputs["input_ids"].shape[1])
                seq_len.append(model_inputs["input_ids"].shape[1])
            
            # pad batched_input_ids to the same length
            batched_input_ids = [torch.cat([input_ids, torch.full((1, max_len - input_ids.shape[1]), self.mask_id, dtype=torch.long, device=self.device)], dim=1) for input_ids in batched_input_ids]
            batched_input_ids = torch.cat(batched_input_ids, dim=0)
            batched_input_ids = batched_input_ids.to(self.device)
            
            # 使用mdm_sample方法，匹配chatbot逻辑
            # breakpoint()
            with torch.no_grad():
                if self.accelerator is not None:
                    generated_ids, nfe = self.accelerator.unwrap_model(self.model).mdm_sample(
                        batched_input_ids,
                        tokenizer=self.tokenizer,
                        block_size=self.bd_size,
                        small_block_size=self.small_block_size,
                        max_tokens=self.max_new_tokens,
                        mask_id=self.mask_id,
                        min_len=min_len,
                        seq_len=torch.tensor(seq_len, device=self.device),
                        use_block_cache=self.use_block_cache,
                        threshold=self.threshold,
                    )
                else:
                    generated_ids, nfe = self.model.mdm_sample(
                        batched_input_ids,
                        tokenizer=self.tokenizer,
                        block_size=self.bd_size,
                        small_block_size=self.small_block_size,
                        max_tokens=self.max_new_tokens,
                        mask_id=self.mask_id,
                        min_len=min_len,
                        seq_len=torch.tensor(seq_len, device=self.device),
                        use_block_cache=self.use_block_cache,
                        threshold=self.threshold,
                    )
            
            # 提取新生成的token，匹配chatbot的解码逻辑，并保持原始索引顺序
            for batch_pos, (orig_idx, req) in enumerate(batch):
                generated_answer = self.tokenizer.decode(
                    generated_ids[batch_pos][seq_len[batch_pos]:], 
                    skip_special_tokens=True
                )
            
                # 统计token数量
                if self.show_speed:
                    # num_tokens += len(self.tokenizer.encode(generated_answer, add_special_tokens=False))
                    num_tokens += (generated_ids[batch_pos][seq_len[batch_pos]:] != self.mask_id).sum()
                    num_nfe += nfe
                
                # 将结果放到正确的原始索引位置
                output[orig_idx] = generated_answer

                print('=' * 20)
                print('question: ', req.args[0])
                print('answer: ', generated_answer)
                print(f'tokens per nfe: {num_tokens / num_nfe}')
                print('=' * 20, end='\n\n')
            
        end_time = time.time()
        if self.show_speed:
            print(f"Total number of tokens generated: {num_tokens}")
            print(f"Total time taken: {end_time - start_time} seconds")
            print(f"Tokens per second: {num_tokens / (end_time - start_time)}")
            print(f"Total number of nfe: {num_nfe}")
            print(f"Avg tokens per nfe: {num_tokens / num_nfe}")
            
        return output


if __name__ == "__main__":
    cli_evaluate()
    
