# -------------------------------------------------------------------------
#  This file is part of the MindStudio project.
# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
#
# MindStudio is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
# -------------------------------------------------------------------------

from atk.configs.case_config import CaseConfig
from atk.common.utils import cal_tensor_numel


class CaseConfigReport:
    def __init__(self, case: CaseConfig):
        self.case = case

    @property
    def name(self):
        return self.case.name

    @property
    def id(self):
        return self.case.id

    @property
    def api(self):
        return self.case.api

    @property
    def api_type(self):
        return self.case.api_type

    @property
    def version(self):
        return self.case.version

    @property
    def shape(self):
        return self._get_by_attr("shape")

    @property
    def dtype(self):
        return self._get_by_attr("dtype")

    @property
    def range(self):
        return self._get_by_attr("range_values")

    @property
    def type(self):
        return self._get_by_attr("type")

    @property
    def numel(self):
        numel_ret = {}

        def get_numel_by_input_type(input_type_args):
            case_input_type = getattr(self.case, input_type_args)
            cases_shapes = []
            if not case_input_type:
                numel_ret[input_type] = None
                return
            if not isinstance(case_input_type, list):
                case_input_type = [case_input_type]
            for input_case in case_input_type:
                if isinstance(input_case, list):
                    if input_case[0].type == "tensors":
                        cases_shapes.extend(input_case)
                elif input_case.type == "tensor":
                    cases_shapes.append(input_case)
            number_elements = 0
            for case in cases_shapes:
                if case.range_values in ["null", ["null"]]:
                    continue
                number_elements += cal_tensor_numel(case.shape)

            numel_ret[input_type_args] = number_elements

        for input_type in ["method_inputs", "tensor_input", "inputs"]:
            get_numel_by_input_type(input_type)
        return numel_ret

    @staticmethod
    def get_titles():
        titles = [
            "id",
            "name",
            "version",
            "api",
            "api_type",
            "type",
            "shape",
            "dtype",
            "range",
            "numel",
        ]
        return titles

    def _get_by_attr(self, name):
        attr_ret = {}

        def get_attr_by_input_type(input_type_args):
            input_ret = []
            case_input_type = getattr(self.case, input_type_args)
            if not case_input_type:
                attr_ret[input_type_args] = None
                return
            if not isinstance(case_input_type, list):
                case_input_type = [case_input_type]
            for input_case in case_input_type:
                if isinstance(input_case, list):
                    values = [getattr(case, name) for case in input_case]
                else:
                    if name == "shape":
                        values = (
                            [0]
                            if not input_case.shape
                               or "tensor" not in input_case.type
                            else input_case.shape
                        )
                    else:
                        values = [getattr(input_case, name)]
                input_ret.append(values)
            attr_ret[input_type_args] = input_ret

        for input_type in ["method_inputs", "tensor_input", "inputs"]:
            get_attr_by_input_type(input_type)
        return attr_ret