#   Copyright (c) 2022 PaddlePaddle Authors. 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 implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function

import inspect
import os
import fcntl

import paddle
import paddle.fluid.core as core

type_dict_paddle_to_str = {
    paddle.bool: 'bool',
    paddle.uint8: 'uint8',
    paddle.int8: 'int8',
    paddle.int16: 'int16',
    paddle.int32: 'int32',
    paddle.int64: 'int64',
    paddle.float16: 'float16',
    paddle.float32: 'float32',
    paddle.float64: 'float64',
    paddle.complex128: 'complex128',
    paddle.complex64: 'complex64',
}

type_dict_str_to_paddle = {
    'int32': paddle.int32,
    'int64': paddle.int64,
    'float32': paddle.float32,
    'float16': paddle.float16,
    'bool': paddle.bool,
    'uint8': paddle.uint8,
    'int8': paddle.int8,
    'complex128': paddle.complex128,
    'complex64': paddle.complex64,
    'int16': paddle.int16,
}

xpu_test_op_white_list = []
xpu_test_type_white_list = []
xpu_test_op_type_white_list = []
xpu_test_device_op_white_list = []
xpu_test_device_op_type_white_list = []


class XPUOpTestWrapper(object):
    def create_classes(self):
        base_class = None
        classes = []
        return base_class, classes


def get_op_white_list():
    op_white_list = xpu_test_op_white_list
    if os.getenv('XPU_TEST_OP_WHITE_LIST') is not None:
        op_white_list.extend(
            os.getenv('XPU_TEST_OP_WHITE_LIST').strip().split(','))
    return list(set(op_white_list))


def get_type_white_list():
    type_white_list = xpu_test_type_white_list
    if os.getenv('XPU_TEST_TYPE_WHITE_LIST') is not None:
        type_white_list.extend(
            os.getenv('XPU_TEST_TYPE_WHITE_LIST').strip().split(','))
    return list(set(type_white_list))


def get_op_type_white_list():
    op_type_white_list = xpu_test_op_type_white_list
    if os.getenv('XPU_TEST_OP_TYPE_WHITE_LIST') is not None:
        op_type_white_list.extend(
            os.getenv('XPU_TEST_OP_TYPE_WHITE_LIST').strip().split(','))
    return list(set(op_type_white_list))


def get_device_op_white_list():
    device_op_white_list = xpu_test_device_op_white_list
    if os.getenv('XPU_TEST_DEVICE_OP_WHITE_LIST') is not None:
        device_op_white_list.extend(
            os.getenv('XPU_TEST_DEVICE_OP_WHITE_LIST').strip().split(','))
    return list(set(device_op_white_list))


def get_device_op_type_white_list():
    device_op_type_white_list = xpu_test_device_op_type_white_list
    if os.getenv('XPU_TEST_DEVICE_OP_TYPE_WHITE_LIST') is not None:
        device_op_type_white_list.extend(
            os.getenv('XPU_TEST_DEVICE_OP_TYPE_WHITE_LIST').strip().split(','))
    return list(set(device_op_type_white_list))


def make_xpu_op_list(xpu_version):
    ops = []
    raw_op_list = core.get_xpu_device_op_list(xpu_version)
    version_str = "xpu2" if xpu_version == core.XPUVersion.XPU2 else "xpu1"
    op_white_list = get_op_white_list()
    type_white_list = get_type_white_list()
    op_type_white_list = get_op_type_white_list()
    device_op_white_list = get_device_op_white_list()
    device_op_type_white_list = get_device_op_type_white_list()
    print('op_white_list:', op_white_list)
    print('type_white_list:', type_white_list)
    print('op_type_white_list:', op_type_white_list)
    print('device_op_white_list:', device_op_white_list)
    print('device_op_type_white_list:', device_op_type_white_list)

    for op_name, type_list in raw_op_list.items():
        device_op_name = version_str + '_' + op_name
        if op_name in op_white_list or device_op_name in device_op_white_list:
            continue
        for op_type in type_list:
            if op_type in type_white_list or op_type not in type_dict_paddle_to_str.keys(
            ):
                continue

            device_op_type_name = device_op_name + '_' + type_dict_paddle_to_str[
                op_type]
            if device_op_type_name in device_op_type_white_list:
                continue

            op_type_name = op_name + '_' + type_dict_paddle_to_str[op_type]
            if op_type_name in op_type_white_list:
                continue

            ops.append(op_type_name)
    return ops


def get_xpu_op_support_types(op_name, dev_id=0):
    xpu_version = core.get_xpu_device_version(dev_id)
    support_type_list = core.get_xpu_device_op_support_types(op_name,
                                                             xpu_version)
    support_type_str_list = [
        type_dict_paddle_to_str[x] for x in support_type_list
    ]
    return support_type_str_list


def record_op_test(op_name, test_type):
    dirname = os.getenv('XPU_OP_LIST_DIR')
    filename = 'xpu_op_test'
    if dirname is not None:
        filename = os.path.join(dirname, filename)
    with open(filename, 'a') as f:
        fcntl.flock(f, fcntl.LOCK_EX)
        f.write(op_name + '_' + test_type + '\n')


def is_empty_grad_op_type(xpu_version, op, test_type):
    xpu_op_list = core.get_xpu_device_op_list(xpu_version)
    grad_op = op + '_grad'
    if grad_op not in xpu_op_list.keys():
        return True

    grad_op_types = xpu_op_list[op]
    paddle_test_type = type_dict_str_to_paddle[test_type]
    if paddle_test_type not in grad_op_types:
        return True

    return False


def create_test_class(func_globals,
                      test_class,
                      test_type,
                      test_grad=True,
                      ignore_deivce_version=[],
                      test_deivce_version=[]):
    xpu_version = core.get_xpu_device_version(0)
    if xpu_version in ignore_deivce_version:
        return

    if len(test_deivce_version) != 0 and xpu_version not in test_deivce_version:
        return

    test_class_obj = test_class()
    register_classes = inspect.getmembers(test_class_obj, inspect.isclass)
    op_name = test_class_obj.op_name
    no_grad = is_empty_grad_op_type(xpu_version, op_name, test_type)

    for test_class in register_classes:
        if test_class[0] == '__class__':
            continue
        class_obj = test_class[1]
        cls_name = "{0}_{1}".format(test_class[0], str(test_type))
        func_globals[cls_name] = type(cls_name, (class_obj, ),
                                      {'in_type': test_type})

    if hasattr(test_class_obj, 'use_dynamic_create_class'
               ) and test_class_obj.use_dynamic_create_class:
        base_class, dynamic_classes = test_class_obj.dynamic_create_class()
        for dy_class in dynamic_classes:
            cls_name = "{0}_{1}".format(dy_class[0], str(test_type))
            attr_dict = dy_class[1]
            attr_dict['in_type'] = test_type
            func_globals[cls_name] = type(cls_name, (base_class, ), attr_dict)

    record_op_test(op_name, test_type)
    if not no_grad:
        record_op_test(op_name + '_grad', test_type)


def get_test_cover_info():
    xpu_version = core.get_xpu_device_version(0)
    version_str = "xpu2" if xpu_version == core.XPUVersion.XPU2 else "xpu1"
    xpu_op_list = make_xpu_op_list(xpu_version)
    xpu_op_covered = []

    dirname = os.getenv('XPU_OP_LIST_DIR')
    filename = 'xpu_op_test'
    if dirname is not None:
        filename = os.path.join(dirname, filename)
    if os.path.exists(filename) and os.path.isfile(filename):
        with open(filename) as f:
            for line in f:
                test_op_name = line.strip()
                if test_op_name in xpu_op_list:
                    xpu_op_covered.append(test_op_name)
    diff_list = list(set(xpu_op_list).difference(set(xpu_op_covered)))
    total_len = len(set(xpu_op_list))
    covered_len = len(set(xpu_op_covered))
    print('{} test: {}/{}'.format(version_str, covered_len, total_len))
    if (len(diff_list) != 0):
        print("These ops need to be tested on {0}! ops:{1}".format(
            version_str, ','.join(diff_list)))


if __name__ == '__main__':
    get_test_cover_info()
