# 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,
        save_dir=None,
        max_new_tokens=1024,
        temperature=0.0,
        do_sample=False,
        small_block_size=16,
        **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逻辑
        self.model.mdm_sample = types.MethodType(modeling_dfm.DFM_QwenForCausalLM.mdm_sample_cache_small_block, 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(model_path, trust_remote_code=True)
        
        # 添加缺失的属性
        self.is_instruct = is_instruct
        self.show_speed = show_speed
        self.save_dir = save_dir
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.do_sample = do_sample
        self.small_block_size = small_block_size
        self.model_path = model_path

    @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):
        chat_history = [{"role": "system", "content": "detailed thinking off"}] + chat_history
        return self.tokenizer.apply_chat_template(chat_history, add_generation_prompt=True, tokenize=False, enable_thinking=False)

    def loglikelihood(self, requests):
        raise NotImplementedError
    
    def loglikelihood_rolling(self, requests):
        raise NotImplementedError
    
    
    def generate_until(self, requests):
        output = []
        num_tokens = 0
        processed_count = 0
        if self.save_dir is not None:
            os.makedirs(self.save_dir, exist_ok=True)
            rank = self.rank
            save_path = os.path.join(self.save_dir, f'rank_{rank}.jsonl')
            print(f"save_path: {save_path}")
            if os.path.exists(save_path):
                print(f"load from {save_path}")
                with open(save_path, 'r', encoding='utf-8') as f:
                    output = [json.loads(line) for line in f]
                    processed_count = len(output)
                print(f"processed_count: {processed_count}")
        
        start_time = time.time()
        
        for i, req in enumerate(tqdm(requests, desc="Generating...")):
            if i < processed_count:
                continue
                
            question = req.args[0]
            model_inputs = self.tokenizer([question], return_tensors="pt").to(self.device)
            
            # 使用mdm_sample方法，匹配chatbot逻辑
            with torch.no_grad():
                if self.accelerator is not None:
                    generated_ids = self.accelerator.unwrap_model(self.model).mdm_sample(
                        model_inputs["input_ids"],
                        tokenizer=self.tokenizer,
                        block_size=128,
                        small_block_size=self.small_block_size,
                        max_tokens=self.max_new_tokens,
                        mask_id=self.tokenizer.encode("|<MASK>|")[0]
                    )
                else:
                    generated_ids = self.model.mdm_sample(
                        model_inputs["input_ids"],
                        tokenizer=self.tokenizer,
                        block_size=self.model.config.bd_size,
                        small_block_size=self.small_block_size,
                        max_tokens=self.max_new_tokens,
                        mask_id=self.tokenizer.encode("|<MASK>|")[0]
                    )
            
            # 提取新生成的token，匹配chatbot的解码逻辑
            generated_answer = self.tokenizer.decode(
                generated_ids[0][model_inputs["input_ids"].shape[1]:], 
                skip_special_tokens=True
            )
            
            # 统计token数量
            if self.show_speed:
                num_tokens += len(self.tokenizer.encode(generated_answer, add_special_tokens=False))
            
            output.append(generated_answer)

            if self.save_dir is not None:
                # 增量保存新生成的答案
                with open(save_path, 'a', encoding='utf-8') as f:
                    f.write(json.dumps(generated_answer, ensure_ascii=False) + '\n')

            print('=' * 20)
            print('question: ', question)
            print('answer: ', generated_answer)
            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)}")
            
        return output


if __name__ == "__main__":
    cli_evaluate()
    
