# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

from typing import List, Optional, Sequence, Tuple, Callable
import torch
import torch_npu
import pytest
import functools
import re
import numpy as np
import os
import glob
import csv
import shutil
import atexit

_float_dtypes = [
    'float32', 'float16', 'bfloat16'
]
_int_dtypes = [
    'int32', 'int64', 'int16', 'int8'
]
_uint_dtypes = [
    'uint8', 'uint16', 'uint32', 'uint64'
]
_all_dtypes_no_bool = _float_dtypes + _int_dtypes
_all_dtypes = _all_dtypes_no_bool + ['bool']
_32bit_dtypes = ('int32')
_16bit_dtypes = ['float16', 'bfloat16', 'int16']


def libdevice_numel(shape: Sequence[int]) -> int:
    n = 1
    for d in shape:
        n *= int(d)
    return n


def generate_numpy(shape, dtype, low=None, high=None):
    if dtype in _int_dtypes + _uint_dtypes:
        iinfo = np.iinfo(getattr(np, dtype))
        low = iinfo.min if low is None else max(low, iinfo.min)
        high = iinfo.max if high is None else min(high, iinfo.max)
        dty = getattr(np, dtype)
        return np.random.randint(low, high, shape, dtype=dty)
    elif dtype == 'float16' or dtype == 'float32':
        return np.random.normal(0, 1, shape).astype(dtype)
    elif dtype == 'bfloat16':
        return (np.random.normal(0, 1, shape).astype('float32').view('uint32') & np.uint32(0xffff0000)).view('float32')
    elif dtype == 'bool':
        return np.random.randint(low=0, high=2, size=shape).astype(bool)
    else:
        raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype))


def generate_tensor(shape, dtype):
    if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16':
        return torch.randn(size=shape, dtype=eval('torch.' + dtype))
    elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16':
        return torch.randint(low=0, high=2000, size=shape, dtype=eval('torch.' + dtype))
    elif dtype == 'int8':
        return torch.randint(low=0, high=127, size=shape, dtype=eval('torch.' + dtype))
    elif dtype == 'bool':
        return torch.randint(low=0, high=2, size=shape).bool()
    elif dtype == 'uint8':
        return torch.randint(low=0, high=255, size=shape, dtype=torch.uint8)
    else:
        raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype))


def get_triton_sig_typename(dtype):
    if dtype == 'float32':
        tyname = "*fp32"
    elif dtype == 'int32':
        tyname = "*i32"
    elif dtype == 'int64':
        tyname = "*i64"
    elif dtype == 'float16':
        tyname = "*fp16"
    elif dtype == 'int16':
        tyname = "*i16"
    elif dtype == 'int8':
        tyname = "*i8"
    elif dtype == 'bool':
        tyname = "*i1"
    else:
        raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype))
    return tyname

# Relative error: abs(x_ref - x_cal) / abs(x_ref)
# Absolute error: abs(x_ref - x_cal)

# calculation type operators require different error range
# It is a stricter verification and not satisfied now, save it here
def validate_cal(dtype, y_cal, y_ref):
    if dtype == 'float16':
        if torch.mean(y_ref) < 0.001:
            assert torch.abs(y_cal - y_ref) < 0.001, "|y_cal - y_ref| < 0.001 is required !"
        else:
            diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.001
            # all true
            assert diff.all(), "Relative error is less than 0.001 !"
    if dtype == 'float32':
        if torch.mean(y_ref) < 0.0001:
            assert torch.abs(y_cal - y_ref) < 0.0001, "|y_cal - y_ref| < 0.0001 is required !"
        else:
            diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.0001
            assert diff.all(), "Relative error is less than 0.001 !"
    elif dtype == 'bfloat16':
        diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.001
        assert diff.all(), "Relative error is less than 0.001 !"
    elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8':
        assert torch.equal(y_cal, y_ref)
    elif dtype == 'uint8' or dtype == 'uint16' or dtype == 'uint32' or dtype == 'uint64':
        assert torch.equal(y_cal, y_ref)
    elif dtype == 'bool':
        assert torch.equal(y_cal, y_ref)
    else:
        raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype))

# moving and comparison ops require no precision error
def validate_cmp(dtype, y_cal, y_ref, overflow_mode: Optional[str] = None):
    y_cal=y_cal.npu()
    y_ref=y_ref.npu()
    if overflow_mode == "saturate":
        if dtype in ('float16'):
            min_value = -torch.finfo(dtype).min
            max_value = torch.finfo(dtype).max
        elif dtype in ['int32', 'int16', 'int8']:
            min_value = torch.iinfo(dtype).min
            max_value = torch.iinfo(dtype).max
        elif dtype == 'bool':
            min_value = 0
            max_value = 1
        else:
            raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype))
        y_ref = torch.clamp(y_ref, min=min_value, max=max_value)
    if dtype == 'float16':
        torch.testing.assert_close(y_ref, y_cal,  rtol=1e-03, atol=1e-03, equal_nan=True)
    elif dtype == 'bfloat16':
        torch.testing.assert_close(y_ref.to(torch.float32), y_cal.to(torch.float32),  rtol=1e-03, atol=1e-03, equal_nan=True)
    elif dtype == 'float32':
        torch.testing.assert_close(y_ref, y_cal,  rtol=1e-04, atol=1e-04, equal_nan=True)
    elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8':
        assert torch.equal(y_cal, y_ref)
    elif dtype == 'uint8' or dtype == 'uint16' or dtype == 'uint32' or dtype == 'uint64':
        assert torch.equal(y_cal, y_ref)
    elif dtype == 'bool':
        assert torch.equal(y_cal, y_ref)
    else:
        raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype))

def validate_cmp_with_expection(dtype, y_cal, y_ref, expect):
    if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16':
        if expect:
            assert torch.allclose(y_ref, y_cal,  rtol=1e-03, atol=1e-03, equal_nan=True)
        else:
            assert not torch.allclose(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True)
    elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8' \
        or dtype == 'uint8' or dtype == 'uint16' or dtype == 'uint32' or dtype == 'uint64':
        if expect:
            assert torch.equal(y_cal, y_ref)
        else:
            assert not torch.equal(y_cal, y_ref)
    else:
        raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype))

# Use the following pytest fixture to run one test case by only single worker.
# Refer to https://pytest-xdist.readthedocs.io/en/stable/how-to.html#making-session-scoped-fixtures-execute-only-once
@pytest.fixture(scope="function")
def pytest_runonce(worker_id, request, cache):
    if (cache.get(request.node.nodeid, "none")) == "none":
        cache.set(request.node.nodeid, worker_id)
    else:
        file_name = f"pytest_{worker_id}.txt"
        with open(file_name, 'a') as file:
            file.write(f"{request.node.nodeid} is already processed by {worker_id}")
        return True
    yield True
    cache.set(request.node.nodeid, "none")

def raises_with_match(expected_exception, match_pattern):
    def decorator(test_func):
        @functools.wraps(test_func)
        def wrapper(*args, **kwargs):
            with pytest.raises(expected_exception, match=match_pattern):
                return test_func(*args, **kwargs)
        return wrapper
    return decorator

def capture_output(expected_output):
    def decorator(test_func):
        @functools.wraps(test_func)
        def wrapper(*args, **kwargs):
            capsys = kwargs.pop('capsys', None)
            if capsys is None:
                try:
                    capsys = pytest.fixture(capsys)()
                except:
                    raise RuntimeError("This decorator requires pytest's capsys fixture")
            test_func(capsys, *args, **kwargs)
            captured = capsys.readouterr()
            # pybind11::scoped_ostream_redirect captures std::cout with \x00 inserted
            # for now, no idea how to eliminate \x00 from C++ side.
            cleaned = re.sub(r"\x00", "", captured.out)
            assert expected_output in cleaned
        return wrapper
    return decorator

def safe_max(tensor):
    """安全的max计算，处理NaN/Inf情况"""
    # 检查NaN并替换
    if torch.isnan(tensor).any():
        print("警告：输入张量包含NaN值")
        tensor = torch.nan_to_num(tensor, nan=0.0)
    
    # 检查Inf并替换
    if torch.isinf(tensor).any():
        print("警告：输入张量包含Inf值")
        tensor = torch.nan_to_num(tensor, posinf=1e10, neginf=-1e10)
    
    # 处理空张量
    if tensor.numel() == 0:
        print("警告：输入张量为空，返回0")
        return torch.tensor(0.0)
    
    return torch.max(tensor)

def print_max_error(inputs, actual, expected, rel_tol=1e-4):
    """
    增强版张量比较函数，支持多输入张量
    
    参数:
        inputs (Tensor或tuple[Tensor]): 模型输入张量（单个或多个）
        actual (Tensor): 实际输出张量
        expected (Tensor): 预期输出张量
        rel_tol (float): 相对误差容限阈值（默认1e-5）
    """
    # 1. 设备转移（NPU -> CPU）
    if actual.device.type == 'npu' or expected.device.type == 'npu':
        actual = actual.cpu()
        expected = expected.cpu()
        print("已将张量转移到CPU处理")
    
    # 2. 数据类型转换
    actual = actual.float()  # 确保为浮点数类型
    expected = expected.float()
    
    actual_origin = actual
    expected_origin = expected

    # 3. 处理特殊值
    actual = torch.nan_to_num(actual, nan=0.0, posinf=1e10, neginf=-1e10)
    expected = torch.nan_to_num(expected, nan=0.0, posinf=1e10, neginf=-1e10)
    
    # 4. 安全计算绝对误差
    abs_error = torch.abs(actual - expected)
    
    # 5. 计算相对误差（带保护机制）
    denominator = torch.max(torch.abs(expected), torch.tensor(1e-6))
    rel_error = abs_error / denominator
    
    # 6. 使用安全max函数获取最大误差
    max_abs_error = safe_max(abs_error)
    max_rel_error = safe_max(rel_error)
    
    # 7. 检查计算结果
    if torch.isnan(max_abs_error) or torch.isinf(max_abs_error):
        print("错误：无法计算有效绝对最大误差")
        print(f"绝对误差NaN数量: {torch.isnan(abs_error).sum().item()}")
        return
    
    if torch.isnan(max_rel_error) or torch.isinf(max_rel_error):
        print("错误：无法计算有效相对最大误差")
        print(f"相对误差NaN数量: {torch.isnan(rel_error).sum().item()}")
        return
    
    # 8. 查找最大绝对误差位置
    try:
        abs_indices = (abs_error == max_abs_error).nonzero(as_tuple=True)
        rel_indices = (rel_error == max_rel_error).nonzero(as_tuple=True)
    except Exception as e:
        print(f"查找位置时出错: {e}")
        return
    
    # 9. 打印绝对误差报告
    print("\n===== 绝对误差报告 =====")
    print(f"最大绝对误差值: {max_abs_error.item()}")
    if len(abs_indices[0]) > 0:
        print(f"出现位置数: {len(abs_indices[0])}")
        for i in range(min(3, len(abs_indices[0]))):  # 最多显示3个位置
            pos = tuple(idx[i].item() for idx in abs_indices)
            print(f"位置 {pos}:")
            
            # 处理多输入情况
            if isinstance(inputs, (tuple, list)):
                for j, input_tensor in enumerate(inputs):
                    print(f"  输入{j}值 = {input_tensor[pos].item() if input_tensor is not None else 'N/A'}")
            else:
                print(f"  输入值 = {inputs[pos].item() if inputs is not None else 'N/A'}")
                
            print(f"  实际值 = {actual_origin[pos].item()}")
            print(f"  预期值 = {expected_origin[pos].item()}")
            print(f"  绝对误差 = {abs_error[pos].item()}")
            print(f"  相对误差 = {rel_error[pos].item()}")
    else:
        print("未找到最大绝对误差位置")
    
    # 10. 打印相对误差报告
    print("\n===== 相对误差报告 =====")
    print(f"最大相对误差值: {max_rel_error.item():.6e} (容限: {rel_tol})")
    if len(rel_indices[0]) > 0:
        print(f"出现位置数: {len(rel_indices[0])}")
        for i in range(min(3, len(rel_indices[0]))):  # 最多显示3个位置
            pos = tuple(idx[i].item() for idx in rel_indices)
            print(f"位置 {pos}:")
            
            # 处理多输入情况
            if isinstance(inputs, (tuple, list)):
                for j, input_tensor in enumerate(inputs):
                    print(f"  输入{j}值 = {input_tensor[pos].item() if input_tensor is not None else 'N/A':.6f}")
            else:
                print(f"  输入值 = {inputs[pos].item() if inputs is not None else 'N/A':.6f}")
                
            print(f"  实际值 = {actual_origin[pos].item():.6f}")
            print(f"  预期值 = {expected_origin[pos].item():.6f}")
            print(f"  绝对误差 = {abs_error[pos].item():.6e}")
            print(f"  相对误差 = {rel_error[pos].item():.6e}")
            
            if rel_error[pos].item() > rel_tol:
                print("  超过相对误差容限")
    else:
        print("未找到最大相对误差位置")

def get_output_file(module_name: str) -> str:
    return f"{module_name}.txt"

def get_temp_file(module_name: str) -> str:
    return f".{module_name}_temp_results.txt"

def extract_triton_performance(result_dir: str = "./result_dir"):
    try:
        triton_folders = glob.glob(os.path.join(result_dir, "*_pt"))
        if not triton_folders:
            print("未找到triton profiler输出文件夹")
            return None
        triton_folders.sort(key=os.path.getmtime, reverse=True)
        latest_triton_folder = triton_folders[0]
        print(f"\n找到最新的profiler文件夹: {latest_triton_folder}")
        csv_path = os.path.join(latest_triton_folder, "ASCEND_PROFILER_OUTPUT", "op_statistic.csv")
        if not os.path.exists(csv_path):
            print(f"未找到op_statistic.csv文件: {csv_path}")
            return None
        avg_time = None
        with open(csv_path, 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            header = next(reader)
            try:
                avg_time_col = header.index("Avg Time(us)")
            except ValueError:
                print("CSV文件中未找到'Avg Time(us)'列")
                return None
            for row in reader:
                if len(row) > 1 and "triton" in row[1].lower():
                    try:
                        avg_time = float(row[avg_time_col])
                        print(f" Triton kernel平均执行时间: {avg_time:.4f} us")
                        break
                    except (ValueError, IndexError):
                        continue
        if avg_time is None:
            print("未找到包含'triton'的kernel性能数据")
        shutil.rmtree(latest_triton_folder)
        print(f"已清理profiler文件夹: {latest_triton_folder}\n")
        return avg_time
    except Exception as e:
        print(f"提取性能数据时发生错误: {str(e)}")
        return None

_registered_modules = set()

def append_temp_result(module_name: str, shape: tuple, element_count: int, avg_time: float):
    temp_file = get_temp_file(module_name)
    with open(temp_file, 'a', encoding='utf-8') as f:
        shape_str = ','.join(str(s) for s in shape)
        f.write(f"{shape_str},{element_count},{avg_time}\n")
    if module_name not in _registered_modules:
        _registered_modules.add(module_name)
        atexit.register(generate_final_report, module_name)

def generate_final_report(module_name: str):
    temp_file = get_temp_file(module_name)
    output_file = get_output_file(module_name)
    if not os.path.exists(temp_file):
        print("没有性能数据可写入")
        return
    results = []
    with open(temp_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                parts = line.split(',')
                shape = tuple(int(p) for p in parts[:-2])
                element_count = int(parts[-2])
                avg_time = float(parts[-1])
                results.append((shape, element_count, avg_time))

    best_by_shape = {}
    for shape, element_count, avg_time in results:
        current = best_by_shape.get(shape)
        if current is None or avg_time < current[1]:
            best_by_shape[shape] = (element_count, avg_time)

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("=" * 80 + "\n")
        f.write("【所有配置平均耗时汇总】\n")
        for shape, element_count, avg_time in results:
            shape_str = str(shape)
            element_str = f"{element_count:12,d}"
            time_str = f"{avg_time:8.3f} us"
            line = f"形状 {shape_str} | 元素数={element_str} | 平均耗时={time_str}\n"
            f.write(line)
        f.write("=" * 80 + "\n")
        f.write("【各Shape最短耗时汇总】\n")
        for shape, (element_count, avg_time) in best_by_shape.items():
            shape_str = str(shape)
            element_str = f"{element_count:12,d}"
            time_str = f"{avg_time:8.3f} us"
            line = f"形状 {shape_str} | 元素数={element_str} | 最短耗时={time_str}\n"
            f.write(line)
        f.write("=" * 80 + "\n")
    os.remove(temp_file)
    print(f"\n 性能报告已保存到: {output_file}")

_base_default_config = ((16384, 4096), 4096, 16384, 2048)

_base_full_configs = [
    ((256, 4096), 64, 16384, 1024),
    ((256, 4096), 64, 16384, 2048),
    ((256, 4096), 64, 16384, 4096),
    ((256, 4096), 128, 8192, 1024),
    ((256, 4096), 128, 8192, 2048),
    ((256, 4096), 128, 8192, 4096),
    ((256, 4096), 256, 4096, 1024),
    ((256, 4096), 256, 4096, 2048),
    ((256, 4096), 256, 4096, 4096),
    ((256, 4096), 512, 2048, 512),
    ((256, 4096), 512, 2048, 1024),
    ((256, 4096), 512, 2048, 2048),
    ((2048, 4096), 64, 131072, 1024),
    ((2048, 4096), 64, 131072, 2048),
    ((2048, 4096), 128, 65536, 1024),
    ((2048, 4096), 128, 65536, 2048),
    ((2048, 4096), 256, 32768, 1024),
    ((2048, 4096), 256, 32768, 2048),
    ((2048, 4096), 512, 16384, 1024),
    ((2048, 4096), 512, 16384, 2048),
    ((2048, 4096), 1024, 8192, 1024),
    ((2048, 4096), 1024, 8192, 2048),
    ((2048, 4096), 2048, 4096, 1024),
    ((2048, 4096), 2048, 4096, 2048),
    ((16384, 4096), 64, 1048576, 1024),
    ((16384, 4096), 64, 1048576, 2048),
    ((16384, 4096), 128, 524288, 1024),
    ((16384, 4096), 128, 524288, 2048),
    ((16384, 4096), 256, 262144, 1024),
    ((16384, 4096), 256, 262144, 2048),
    ((16384, 4096), 512, 131072, 1024),
    ((16384, 4096), 512, 131072, 2048),
    ((16384, 4096), 1024, 65536, 1024),
    ((16384, 4096), 1024, 65536, 2048),
    ((16384, 4096), 2048, 32768, 1024),
    ((16384, 4096), 2048, 32768, 2048),
    ((16384, 4096), 4096, 16384, 1024),
    ((16384, 4096), 4096, 16384, 2048),
]

def make_default_param_list(dtypes):
    return [[dtype, *_base_default_config] for dtype in dtypes]

def make_full_param_list(dtypes):
    result = []
    for dtype in dtypes:
        for cfg in _base_full_configs:
            result.append([dtype, *cfg])
    return result

def run_with_profiler(kernel_fn: Callable, shape: tuple, op_name: str):
    experimental_config = torch_npu.profiler._ExperimentalConfig(
        aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization,
        profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False
    )

    with torch_npu.profiler.profile(
        activities=[
            torch_npu.profiler.ProfilerActivity.NPU],
        with_stack=False,
        record_shapes=False,
        profile_memory=False,
        schedule=torch_npu.profiler.schedule(wait=1,
                                            warmup=1,
                                            active=30,
                                            repeat=1,
                                            skip_first=1),
        experimental_config=experimental_config,
        on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result_dir")
    ) as prof:
        for i in range(10):
            kernel_fn()
            torch.npu.synchronize()
            prof.step()

    avg_time = extract_triton_performance()
    if avg_time is not None:
        element_count = libdevice_numel(shape)
        append_temp_result(op_name, shape, element_count, avg_time)