from enum import Enum
from typing import Callable, Dict, List, Optional, Union

import numpy as np

# @manual=//deeplearning/trt/python:py_tensorrt
import tensorrt as trt
import torch
from functorch import make_fx
from functorch.experimental import functionalize
from torch_tensorrt._utils import sanitized_torch_version
from torch_tensorrt.fx.passes.lower_basic_pass import (
    replace_op_with_indices,
    run_const_fold,
)

from packaging import version

from .types import Shape, TRTDataType


class Frameworks(Enum):
    NUMPY = "numpy"
    TORCH = "torch"
    TRT = "trt"


DataTypeEquivalence: Dict[
    TRTDataType, Dict[Frameworks, Union[TRTDataType, np.dtype, torch.dtype]]
] = {
    trt.int8: {
        Frameworks.NUMPY: np.int8,
        Frameworks.TORCH: torch.int8,
        Frameworks.TRT: trt.int8,
    },
    trt.int32: {
        Frameworks.NUMPY: np.int32,
        Frameworks.TORCH: torch.int32,
        Frameworks.TRT: trt.int32,
    },
    trt.int64: {
        Frameworks.NUMPY: np.int64,
        Frameworks.TORCH: torch.int64,
        Frameworks.TRT: trt.int64,
    },
    trt.float16: {
        Frameworks.NUMPY: np.float16,
        Frameworks.TORCH: torch.float16,
        Frameworks.TRT: trt.float16,
    },
    trt.float32: {
        Frameworks.NUMPY: np.float32,
        Frameworks.TORCH: torch.float32,
        Frameworks.TRT: trt.float32,
    },
    trt.bool: {
        Frameworks.NUMPY: bool,
        Frameworks.TORCH: torch.bool,
        Frameworks.TRT: trt.bool,
    },
}

if trt.__version__ >= "7.0":
    DataTypeEquivalence[trt.bool] = {
        Frameworks.NUMPY: np.bool_,
        Frameworks.TORCH: torch.bool,
        Frameworks.TRT: trt.bool,
    }


class LowerPrecision(Enum):
    FP32 = "fp32"
    FP16 = "fp16"
    INT8 = "int8"
    BF16 = "bf16"

    @staticmethod
    def from_str(label: str) -> Optional["LowerPrecision"]:
        if label in ("fp32", "float32", "float", "torch.float32"):
            return LowerPrecision.FP32
        elif label in ("fp16", "float16", "half", "torch.half", "torch.float16"):
            return LowerPrecision.FP16
        elif label in ("int8"):
            return LowerPrecision.INT8
        elif label in ("bf16", "bfloat16", "torch.bfloat16"):
            return LowerPrecision.BF16
        else:
            return None


def unified_dtype_converter(
    dtype: Union[TRTDataType, torch.dtype, np.dtype], to: Frameworks
) -> Union[np.dtype, torch.dtype, TRTDataType]:
    """
    Convert TensorRT, Numpy, or Torch data types to any other of those data types.

    Args:
        dtype (TRTDataType, torch.dtype, np.dtype): A TensorRT, Numpy, or Torch data type.
        to (Frameworks): The framework to convert the data type to.

    Returns:
        The equivalent data type in the requested framework.
    """
    assert to in Frameworks, f"Expected valid Framework for translation, got {to}"
    trt_major_version = int(trt.__version__.split(".")[0])
    if dtype in (np.int8, torch.int8, trt.int8):
        return DataTypeEquivalence[trt.int8][to]
    elif trt_major_version >= 7 and dtype in (np.bool_, torch.bool, trt.bool):
        return DataTypeEquivalence[trt.bool][to]
    elif dtype in (np.int32, torch.int32, trt.int32):
        return DataTypeEquivalence[trt.int32][to]
    elif dtype in (np.int64, torch.int64, trt.int64):
        return DataTypeEquivalence[trt.int64][to]
    elif dtype in (np.float16, torch.float16, trt.float16):
        return DataTypeEquivalence[trt.float16][to]
    elif dtype in (np.float32, torch.float32, trt.float32):
        return DataTypeEquivalence[trt.float32][to]
    else:
        raise TypeError("%s is not a supported dtype" % dtype)


def get_dynamic_dims(shape: Shape) -> List[int]:
    """
    This function finds the dynamic dimensions in the given
    shape. A dimension is dynamic if it's -1.

    Args:
        shape (Shape): A sequence of integer that represents
            the shape of a tensor.

    Returns:
        A list of integers contains all the dynamic dimensions
        in the given shape
    """
    dynamic_dims = []

    for i, s in enumerate(shape):
        if s == -1:
            dynamic_dims.append(i)

    return dynamic_dims


def proxytensor_trace(mod, inputs):
    mod.eval()

    def f(*inp):
        return mod(*inp)

    mod = make_fx(functionalize(f))(*inputs)

    # Remove const operation. For ex, nn.Linear has transpose operation on weight
    mod.graph.eliminate_dead_code()
    mod = run_const_fold(mod)
    mod = replace_op_with_indices(mod)
    return mod


def req_torch_version(min_torch_version: str = "2.dev"):
    """
    Create a decorator which verifies the Torch version installed
    against a specified version range

    Args:
        min_torch_version (str): The minimum required Torch version
        for the decorated function to work properly

    Returns:
        A decorator which raises a descriptive error message if
        an unsupported Torch version is used
    """

    def nested_decorator(f: Callable):
        def function_wrapper(*args, **kwargs):
            # Parse minimum and current Torch versions
            min_version = version.parse(min_torch_version)
            current_version = version.parse(sanitized_torch_version())

            if current_version < min_version:
                raise AssertionError(
                    f"Expected Torch version {min_torch_version} or greater, "
                    + f"when calling {f}. Detected version {torch.__version__}"
                )
            else:
                return f(*args, **kwargs)

        return function_wrapper

    return nested_decorator