# Copyright (c) 2023 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 annotations

import builtins
import copy
import inspect
import sys
import time
import types
import weakref
from collections import OrderedDict
from collections.abc import Callable
from contextlib import contextmanager
from dataclasses import is_dataclass
from functools import lru_cache
from typing import TYPE_CHECKING, Any, TypeVar
from weakref import WeakValueDictionary

import numpy as np

import paddle
from paddle.jit.dy2static.utils import (
    TransformOptions,
    dataclass_as_dict,
    dataclass_from_dict,
)
from paddle.utils import flatten, map_structure

from .envs import (
    ENV_SOT_LOG_LEVEL,
    ENV_SOT_SPECIALIZED_DIM_NUMBERS,
    ENV_STRICT_MODE,
)
from .paddle_api_config import (
    break_graph_functions,
    paddle_api_list,
    paddle_api_module_prefix,
)

if TYPE_CHECKING:
    from collections.abc import Callable

    from paddle._typing import NestedStructure

T = TypeVar("T")
T1 = TypeVar("T1")
T2 = TypeVar("T2")
T3 = TypeVar("T3")
ConstTypes = (int, float, str, bool, type(None), bytes)


class Singleton(type):
    _instances: dict[Any, Any] = {}

    def __call__(cls, *args: Any, **kwargs: Any):
        if cls not in cls._instances:
            cls._instances[cls] = super().__call__(*args, **kwargs)
        return cls._instances[cls]


class NameGenerator:
    def __init__(self, prefix):
        self.counter = 0
        self.prefix = prefix

    def next(self):
        name = self.prefix + str(self.counter)
        self.counter += 1
        return name

    def match_name(self, name: str) -> bool:
        return name.startswith(self.prefix)


class SymbolRegistry:
    def __init__(self):
        self.symbol_generator = NameGenerator(prefix="___t_")
        self.tmp_names_record = OrderedDict()
        self.declared_symbols: set[str] = set()
        self.symbol_table = {}

    def next_symbol(self) -> str:
        return self.symbol_generator.next()

    def request_symbol(self, expr: str) -> str:
        if expr in self.symbol_table:
            return self.symbol_table[expr]
        symbol = self.next_symbol()
        self.symbol_table[expr] = symbol
        return symbol

    def gen_expr(self, expr: str, gen_expr_fn):
        symbol = self.symbol_table[expr]
        if symbol in self.declared_symbols:
            return symbol
        self.declared_symbols.add(symbol)
        return f"({symbol} := ({gen_expr_fn()}))"


_symbol_registry = SymbolRegistry()


@contextmanager
def switch_symbol_registry():
    global _symbol_registry
    original_registry = _symbol_registry
    _symbol_registry = SymbolRegistry()
    yield
    _symbol_registry = original_registry


def current_symbol_registry():
    global _symbol_registry
    return _symbol_registry


class ResumeFnNameFactory(metaclass=Singleton):
    def __init__(self) -> None:
        self.gen = NameGenerator('resume_')

    def next(self):
        name = self.gen.next()
        return name


class SIRToCodeMap(metaclass=Singleton):
    def __init__(self):
        self._map = {}

    def register(self, sir, code):
        self._map[sir.name] = code

    def get(self, sir):
        return self._map.get(sir.name)


def log(level, *args):
    cur_level = ENV_SOT_LOG_LEVEL.get()
    if level <= cur_level:
        print(*args, end="", flush=True)


def log_do(level, fn):
    cur_level = ENV_SOT_LOG_LEVEL.get()
    if level <= cur_level:
        fn()


def log_format(level, str, *args):
    cur_level = ENV_SOT_LOG_LEVEL.get()
    if level <= cur_level:
        print(str.format(*args), end="", flush=True)


def log_enabled(level):
    return level <= ENV_SOT_LOG_LEVEL.get()


@lru_cache
def log_once(msg):
    print(msg, flush=True)


def no_eval_frame(func):
    def no_eval_frame_func(*args, **kwargs):
        old_cb = paddle.framework.core.set_eval_frame(None)
        try:
            retval = func(*args, **kwargs)
        except:
            raise
        finally:
            paddle.framework.core.set_eval_frame(old_cb)
        return retval

    return no_eval_frame_func


def is_comprehensive_name(name):
    return name in ["<listcomp>", "<dictcomp>", "<setcomp>", "<genexpr>"]


def is_paddle_api(func):
    if isinstance(func, paddle.nn.Layer):  # ignore all the classes
        return False
    if hasattr(func, "__self__"):  # ignore all the methods
        return False
    if inspect.isclass(
        func
    ):  # paddle.Tensor should not be wrapped, but how about other situations?
        return False
    return in_paddle_module(func) or func in paddle_api_list


def already_unified_in_dynamic_and_static_graph(fn):
    if is_paddle_api(fn):
        return True
    return not TransformOptions.check_fn_need_transform(
        fn, TransformOptions.ToStaticMode.SOT
    )


def need_capture_control_flow(fn):
    return TransformOptions.check_fn_need_capture_control_flow(fn)


def is_builtin_fn(fn):
    special_builtin_fns = [weakref.ref]
    if fn in special_builtin_fns:
        return True
    if isinstance(fn, types.BuiltinFunctionType):
        return True
    for member_name, member in inspect.getmembers(builtins):
        if member is fn and isinstance(member, type):
            return True
    return False


def in_paddle_module(func):
    if hasattr(func, "__module__"):
        module_str = func.__module__
        if module_str is None:
            return False
        log(5, "find paddle function with __module__: ", module_str, "\n")
        if hasattr(func, "__name__"):
            log(
                5, "                     with __name__  : ", func.__name__, "\n"
            )
        log(5, "                     with results   : ")
        for prefix in paddle_api_module_prefix:
            if module_str.startswith(prefix):
                log(5, " True\n")
                return True
    log(5, " False\n")
    return False


def is_break_graph_api(func):
    return func in break_graph_functions


def is_namedtuple_class(cls):
    if not inspect.isclass(cls):
        return False
    if not issubclass(cls, tuple):
        return False
    # The signature created by nametuple function
    namedtuple_attrs = {"_make", "_asdict", "_fields", "_replace"}
    cls_attrs = set(dir(cls))
    return namedtuple_attrs.issubset(cls_attrs)


def map_if(
    *structures: NestedStructure[T1],
    pred: Callable[[T1], bool],
    true_fn: Callable[[T1], T2],
    false_fn: Callable[[T1], T3],
) -> NestedStructure[T2 | T3]:
    def replace(*args):
        if pred(*args):
            return true_fn(*args)
        return false_fn(*args)

    return map_structure(replace, *structures)


def flatten_extend(structure):
    for item in flatten(structure):
        if isinstance(item, slice):
            yield item.start
            yield item.stop
            yield item.step
        else:
            yield item


def map_if_extend(structure, pred, true_fn, false_fn):
    """support extended structures like slice and SliceVariable"""

    def wrapped_pred(x):
        if isinstance(x, slice):
            return True
        if is_dataclass(x) and not isinstance(x, type):
            return True
        return pred(x)

    def wrapped_true_fn(x):
        if isinstance(x, (slice)):
            l = [x.start, x.stop, x.step]
            l = map_if_extend(l, pred, true_fn, false_fn)
            return slice(*l)

        if is_dataclass(x) and not isinstance(x, type):
            dt_dict = dataclass_as_dict(x)
            dt_dict = map_if_extend(dt_dict, pred, true_fn, false_fn)
            return dataclass_from_dict(type(x), dt_dict)

        return true_fn(x)

    return map_if(
        structure, pred=wrapped_pred, true_fn=wrapped_true_fn, false_fn=false_fn
    )


def count_if(*structures, pred):
    def is_true(*args):
        if pred(*args):
            return 1
        return 0

    return sum(flatten(map_structure(is_true, *structures)))


class Cache:
    def __init__(self, weak=False, copy=False):
        if not weak:
            self.cache = {}
        else:
            self.cache = WeakValueDictionary()
        self.hit_num = 0
        self.copy = copy

    def __call__(self, *args, **kwargs):
        cache_key = self.key_fn(*args, **kwargs)
        if not hashable(cache_key):
            return self.value_fn(*args, **kwargs)
        if cache_key in self.cache:
            log(5, "cache hit: ", cache_key, "\n")
            self.hit_num += 1
            cache_item = self.cache[cache_key]
            if self.copy:
                cache_item = copy.deepcopy(cache_item)
            return cache_item
        value = self.value_fn(*args, **kwargs)
        self.cache[cache_key] = value
        return value

    def clear(self):
        self.cache.clear()
        self.hit_num = 0

    def key_fn(self, *args, **kwargs):
        raise NotImplementedError

    def value_fn(self, *args, **kwargs):
        raise NotImplementedError


def execute_time(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        print("Execute time:", execution_time)
        return result

    return wrapper


def meta_str(shape, dtype, stop_gradient):
    return f"(shape: {shape}, dtype: {dtype}, stop_gradient: {stop_gradient})"


def is_strict_mode():
    return ENV_STRICT_MODE.get()


def list_find_index_by_id(li: list[Any], item: Any) -> int:
    return [id(it) for it in li].index(id(item))


def list_contain_by_id(li: list[Any], item: Any) -> int:
    return id(item) in [id(it) for it in li]


def get_unbound_method(obj, name):
    # TODO(dev): Consider the case of patching methods to instances
    return getattr(obj.__class__, name)


class SotUndefinedVar(metaclass=Singleton):
    pass


def hashable(obj):
    try:
        hash(obj)
        return True
    except TypeError as e:
        return False


def printable(obj):
    try:
        str(obj)
        return True
    except Exception as e:
        return False


class StepInfo:
    BACK_TRACE_STEPS = 20

    def __init__(self):
        self.step_count = -1

    def need_back_trace(self):
        return self.step_count < self.BACK_TRACE_STEPS


class StepInfoManager(metaclass=Singleton):
    def __init__(self):
        self.step_record = {}
        self.current_code = None
        self.current_step_info = None

    @contextmanager
    def step_guard(self, code):
        try:
            old_code = self.current_code
            old_info = self.current_step_info

            self.current_code = code
            if code not in self.step_record:
                self.step_record[code] = StepInfo()
            self.current_step_info = self.step_record[code]

            self.current_step_info.step_count += 1
            yield
        finally:
            self.current_code = old_code
            self.current_step_info = old_info

    @property
    def need_back_trace(self):
        return (
            self.current_step_info is not None
            and self.current_step_info.need_back_trace()
        )

    @property
    def current_step(self):
        return self.current_step_info.step_count

    def clear(self):
        self.step_record.clear()
        self.current_code = None
        self.current_step = -1


def get_api_fullname(api):
    api_name = api.__name__
    module_str = api.__module__
    while len(module_str) > 0:
        if module_str not in sys.modules:
            return api_name
        module = sys.modules[module_str]
        if hasattr(module, api_name):
            return module_str + "." + api_name
        module_str = module_str.rpartition(".")[0]
    return None


def get_numpy_ufuncs():
    ufuncs = [
        ufunc
        for _, ufunc in inspect.getmembers(
            np, lambda member: isinstance(member, np.ufunc)
        )
    ]
    unary_ufuncs = filter(lambda ufunc: ufunc.nin == 1, ufuncs)
    binary_ufuncs = filter(lambda ufunc: ufunc.nin == 2, ufuncs)
    return list(unary_ufuncs), list(binary_ufuncs)


def do_until_stop_iteration(fn: Callable[[], T]) -> list[T]:
    from paddle.jit.sot.utils.exceptions import SotCapturedStopIteration

    res = []
    while True:
        try:
            res.append(fn())
        except SotCapturedStopIteration:
            break
    return res


def update_list_inplace(
    original_list: list[T], new_contents: list[T]
) -> list[T]:
    original_list.clear()
    original_list.extend(new_contents)
    return original_list


def get_obj_stable_repr(obj) -> str:
    if hasattr(obj, '__qualname__'):
        return obj.__qualname__
    if hasattr(obj, '__name__'):
        return obj.__name__

    class_name = obj.__class__.__name__

    # If module is available and not __main__, include it
    if hasattr(obj, "__class__") and hasattr(obj.__class__, "__module__"):
        module = obj.__class__.__module__
        if module not in ("__main__", "builtins"):
            return f"{module}.{class_name}()"

    return f"{class_name}()"


def get_min_non_specialized_number() -> int:
    specialized_dim_numbers_raw_str = (
        ENV_SOT_SPECIALIZED_DIM_NUMBERS.get().lower()
    )
    assert specialized_dim_numbers_raw_str in [
        "no",
        "0",
        "01",
    ], f"Unsupported specialized_dim_numbers: {specialized_dim_numbers_raw_str}"
    to_min_non_specialized_number = {
        # specialized numbers, minimum non-specialized number
        "no": 0,
        "0": 1,
        "01": 2,
    }
    return to_min_non_specialized_number[specialized_dim_numbers_raw_str]