# coding=utf-8
# Copyright (c) 2024, HUAWEI CORPORATION.  All rights reserved.
# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
#
# 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 or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
批量测试推理脚本 - 支持中英文测试样例
用于测试0416和0422模型的各个checkpoint
"""

import sys
import time
import logging
import argparse

import torch
from torch import distributed as dist

logging.basicConfig(format="")
logging.getLogger().setLevel(logging.INFO)

# 中文测试样例
CHINESE_PROMPTS = [
    "你好",
    "介绍一下北京",
    "为什么天空是蓝色的",
    "你知道杭州么",
    # "如果今天下雨你需要",
    # "请用中文解释什么是人工智能。",
    # "中国的首都是",
    # "解释什么是机器学习",
    # "今天天气如何"
    # "请写一首关于春天的短诗。",
    # "翻译成英文：今天天气很好。",
    # "请用中文写一个简单的Python函数，计算两个数的和。",
    # "请解释牛顿第一定律。",
    # "请用中文描述一下北京的天气。",
    # "1+1等于多少？",
    # "以下是关于中国历史的单项选择题，请直接给出正确答案的选项。秦朝统一中国是在哪一年？A. 公元前221年B. 公元前206年  C. 公元前256年 D. 公元前230年 答案：",
    # "国务院今日召开新闻发布会，宣布新一轮经济刺激政策将",
    # "人类大脑包含约860亿个神经元，这些神经元之间通过",
    # "她打开信封，看到第一行字，双手开始颤抖，因为",
    # "所有的哺乳动物都是温血动物。鲸鱼是哺乳动物。因此，",
    # "法国的首都是巴黎。德国的首都是柏林。日本的首都是",
    # '你有什么特长？',
    '为什么天空是蓝色的',
    '请用Python写一个计算斐波那契数列的函数',
    '解释一下"光合作用"的基本过程',
    '如果明天下雨，我应该如何出门',
    '比较一下猫和狗作为宠物的优缺点',
    '解释什么是机器学习',
    '推荐一些中国的美食'
]

# 英文测试样例
temp = "The following is a question. " * 10
ENGLISH_PROMPTS = [
    # temp + "Hello, how are you today?",
    # temp + "What is the capital of the United States?",
    # temp + "Please write a short poem about the ocean.",
    # temp + "Translate to Chinese: Artificial Intelligence is changing the world.",
    # temp + "Please explain the concept of machine learning.",
    # temp + "If you could travel anywhere, where would you go?",
    # temp + "Write a simple Python function to check if a number is even.",
    # temp + "What is the difference between CPU and GPU?",
    # temp + "Please describe what a neural network is.",
    # temp + "What is 2 times 3?",
    "The president announced today that the new economic policy would",
    "The human brain contains approximately 86 billion neurons. These neurons",
    "She opened the letter and read the first line. Her hands began to tremble because",
    "All mammals are warm-blooded. Whales are mammals. Therefore,",
    "The capital of France is Paris. The capital of Germany is Berlin. The capital of Japan is"
]


def print_flush(prev_str, curr_str):
    difference = ''.join([char2 for char1, char2 in zip(prev_str, curr_str) if char1 != char2])

    if len(prev_str) < len(curr_str):
        difference += curr_str[len(prev_str):]

    sys.stdout.write(difference)


def task_factory(args, model):
    """测试任务工厂"""
    task_map = {
        "greedy": task_greedy_search,
        "do_sample": task_do_sample,
        "batch_test": task_batch_test,
        "chat": task_chat,
    }

    total_tasks = args.task

    if total_tasks is None:
        total_tasks = ["greedy"]

    for task in total_tasks:
        if task not in task_map.keys():
            raise ValueError(f"Task name incorrect: {task}")

        task_map.get(task)(args, model)


def task_greedy_search(args, model):
    """Greedy Search - 单条测试"""
    instruction = "你好，请介绍一下你自己。"

    t = time.time()
    output = model.generate(
        [instruction],
        do_sample=False,
        max_new_tokens=args.max_new_tokens,
        stream=False
    )

    if dist.get_rank() == 0:
        logging.info("\n=============== Greedy Search ================")
        logging.info("\n输入:\n%s\n\n模型输出:\n%s", instruction, output)
        logging.info("==============================================")
        logging.info("\n耗时: %ss", round(time.time() - t, 2))

    dist.barrier()


def task_do_sample(args, model):
    """Do Sample - 单条测试"""
    # instruction = "你好，请介绍一下你自己。"
    instruction = "The president announced today that the new economic policy would"

    t = time.time()
    output = model.generate(
        [instruction],
        do_sample=True,
        top_k=args.top_k if args.top_k else 50,
        top_p=args.top_p if args.top_p else 0.95,
        temperature=args.temperature if args.temperature else 0.7,
        max_new_tokens=args.max_new_tokens,
        stream=False
    )

    if dist.get_rank() == 0:
        logging.info("\n=============== Do Sample ================")
        logging.info("\n输入:\n%s\n\n模型输出:\n%s", instruction, output)
        logging.info("==========================================")
        logging.info("\n耗时: %ss", round(time.time() - t, 2))

    dist.barrier()


def task_batch_test(args, model):
    """批量测试 - 测试所有中英文样例"""
    all_prompts = CHINESE_PROMPTS + ENGLISH_PROMPTS 
    
    logging.info("\n" + "=" * 60)
    logging.info("批量测试开始 - 共 %d 条测试样例", len(all_prompts))
    logging.info("=" * 60)
    
    results = []
    
    for i, instruction in enumerate(all_prompts):
        lang = "中文" if i < len(CHINESE_PROMPTS) else "英文"
        
        t = time.time()
        # output = model.generate(
        #     [instruction],
        #     do_sample=False,
        #     max_new_tokens=args.max_new_tokens,
        #     stream=False,
        # )
        output = model.generate(
            [instruction],
            do_sample=True,
            top_k=args.top_k if args.top_k else 50,
            top_p=args.top_p if args.top_p else 0.95,
            temperature=args.temperature if args.temperature else 0.7,
            max_new_tokens=args.max_new_tokens,
            stream=False
        )

        elapsed = round(time.time() - t, 2)
        
        if dist.get_rank() == 0:
            logging.info("\n--- 测试 %d [%s] ---", i + 1, lang)
            logging.info("输入: %s", instruction)
            logging.info("输出: %s", output[0] if isinstance(output, list) else output)
            logging.info("耗时: %ss", elapsed)
            
            results.append({
                "id": i + 1,
                "language": lang,
                "prompt": instruction,
                "output": output[0] if isinstance(output, list) else output,
                "elapsed": elapsed
            })
        
        dist.barrier()
    
    if dist.get_rank() == 0:
        logging.info("\n" + "=" * 60)
        logging.info("批量测试完成")
        logging.info("=" * 60)
        
        # 打印结果汇总
        logging.info("\n=== 测试结果汇总 ===")
        for r in results:
            logging.info("[%d] %s: %s -> %s (耗时%s)", 
                        r["id"], r["language"], r["prompt"][:20], 
                        r["output"][:50] if r["output"] else "无输出", r["elapsed"])


def task_chat(args, model):
    """交互式对话模式"""
    histories_no_template = []
    histories_template = []
    instruction = None
    prompt = ""
    input_template = "\n\nYou >> "
    command_clear = ["clear"]

    while True:
        terminate_runs = torch.zeros(1, dtype=torch.int64, device=torch.cuda.current_device())

        if dist.get_rank() == 0:
            if not histories_no_template and not histories_template:
                logging.info("===========================================================")
                logging.info("1. 输入 q, quit, exit 退出")
                logging.info("2. 输入 clear, new 开始新对话")
                logging.info("===========================================================")

            prompt = input(input_template)
            prompt = prompt.encode('utf-8', errors='ignore').decode('utf-8')
            if prompt.strip() in ["q", "exit", "quit"]:
                terminate_runs += 1

            if prompt.strip() in ["clear", "new"]:
                subprocess.call(command_clear)
                histories_no_template = []
                histories_template = []
                continue

            if not prompt.strip():
                continue

        dist.all_reduce(terminate_runs)
        dist.barrier()
        if terminate_runs > 0:
            break

        if not instruction.strip():
            continue

        responses = model.generate(
            instruction,
            do_sample=True,
            top_k=args.top_k if args.top_k else 50,
            top_p=args.top_p if args.top_p else 0.95,
            temperature=args.temperature if args.temperature else 0.7,
            max_new_tokens=args.max_new_tokens,
            stream=True,
            broadcast=True
        )

        # 处理响应
        prev = ""
        curr = ""
        output = ""
        for response in responses:
            curr = response[0]
            print_flush(prev, curr)
            prev = curr
            output = curr

        histories_template.append({"role": "assistant", "content": output})
        histories_no_template.append((prompt, output))
        if len(histories_no_template) > 3:
            histories_no_template.pop()


def main():
    """主函数 - 需要在MindSpeed环境中调用"""
    from mindspeed_llm import megatron_adaptor
    from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_local_spec
    from megatron.core.transformer.spec_utils import import_module
    from megatron.training import get_args, print_rank_0
    from megatron.legacy.model import GPTModel
    from megatron.training.initialize import initialize_megatron
    from megatron.training.arguments import core_transformer_config_from_args
    from mindspeed_llm.tasks.inference.module import GPTModelInfer, MegatronModuleForCausalLM

    def model_provider(pre_process=True, post_process=True):
        args = get_args()
        use_te = args.transformer_impl == "transformer_engine"

        print_rank_0('building GPT model ...')
        
        if args.yaml_cfg is not None:
            from megatron.training.yaml_arguments import core_transformer_config_from_yaml
            config = core_transformer_config_from_yaml(args, "language_model")
        else:
            config = core_transformer_config_from_args(args)

        if args.use_mcore_models:
            if args.spec is not None:
                transformer_layer_spec = import_module(args.spec)
            else:
                if use_te:
                    transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(args.num_experts, args.moe_grouped_gemm)
                else:
                    transformer_layer_spec = get_gpt_layer_local_spec(args.num_experts, args.moe_grouped_gemm)

            model = GPTModelInfer(
                config=config,
                transformer_layer_spec=transformer_layer_spec,
                vocab_size=args.padded_vocab_size,
                max_sequence_length=args.max_position_embeddings,
                pre_process=pre_process,
                post_process=post_process,
                fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
                parallel_output=True if args.sequence_parallel else False,
                share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
                position_embedding_type=args.position_embedding_type,
                rotary_percent=args.rotary_percent,
                seq_len_interpolation_factor=args.rotary_seq_len_interpolation_factor
            )
        else:
            model = GPTModel(
                config,
                parallel_output=True if args.sequence_parallel else False,
                pre_process=pre_process,
                post_process=post_process
            )

        return model

    initialize_megatron(args_defaults={'no_load_rng': True, 'no_load_optim': True})
    args = get_args()

    model = MegatronModuleForCausalLM.from_pretrained(
        model_provider=model_provider,
        pretrained_model_name_or_path=args.load
    )

    task_factory(args, model)


if __name__ == "__main__":
    main()