from typing import (
    Tuple,
    cast,
    Union,
    Any,
    TypeVar,
    TYPE_CHECKING,
    Iterable,
    Optional,
    overload,
    Callable,
)
from typing_extensions import Literal
import numpy as np
from importlib import import_module
import functools

from ..types import Axes, AxisAxes, Shape, ShapeOrScalar

from .. import index

from .tensor import Tensor
from .tensor import TensorOrScalar
from .tensor import TensorType

from .base import BaseTensor
from .base import unwrap_
from .base import unwrap1

if TYPE_CHECKING:
    import tensorflow as tf  # for static analyzers
    from .extensions import NormsMethods  # noqa: F401
else:
    # lazy import in TensorFlowTensor
    tf = None

FuncType = Callable[..., Any]
F = TypeVar("F", bound=FuncType)


def samedevice(f: F) -> F:
    @functools.wraps(f)
    def wrapper(self: "TensorFlowTensor", *args: Any, **kwargs: Any) -> Any:
        with tf.device(self.raw.device):
            return f(self, *args, **kwargs)

    return cast(F, wrapper)


def common_dtype(f: F) -> F:
    @functools.wraps(f)
    def wrapper(self: "TensorFlowTensor", *args: Any, **kwargs: Any) -> Any:
        dtypes = {self.dtype} | {arg.dtype for arg in args if isinstance(arg, Tensor)}
        if len(dtypes) == 1:
            # all dtypes are the same, nothing more to do
            return f(self, *args, **kwargs)
        numpy_dtypes = [np.dtype(dtype.name) for dtype in dtypes]
        common = np.find_common_type(numpy_dtypes, [])
        common = getattr(tf, common.name)
        if self.dtype != common:
            self = self.astype(common)
        args = tuple(
            arg.astype(common)
            if isinstance(arg, Tensor) and arg.dtype != common
            else arg
            for arg in args
        )
        return f(self, *args, **kwargs)

    return cast(F, wrapper)


def assert_bool(x: Any) -> None:
    if not isinstance(x, Tensor):
        return
    if x.dtype != tf.bool:
        raise ValueError(f"requires dtype bool, got {x.dtype}, consider t.bool().all()")


class TensorFlowTensor(BaseTensor):
    __slots__ = ()

    # more specific types for the extensions
    norms: "NormsMethods[TensorFlowTensor]"

    def __init__(self, raw: "tf.Tensor"):  # type: ignore
        global tf
        if tf is None:
            tf = import_module("tensorflow")
        super().__init__(raw)

    @property
    def raw(self) -> "tf.Tensor":  # type: ignore
        return super().raw

    def numpy(self: TensorType) -> Any:
        a = self.raw.numpy()
        if a.flags.writeable:
            # without the check, we would attempt to set it on array
            # scalars, and that would fail
            a.flags.writeable = False
        return a

    def item(self: TensorType) -> Union[int, float, bool]:
        return self.numpy().item()  # type: ignore

    @property
    def shape(self: TensorType) -> Shape:
        return tuple(self.raw.shape.as_list())

    def reshape(self: TensorType, shape: Union[Shape, int]) -> TensorType:
        if isinstance(shape, int):
            shape = (shape,)
        return type(self)(tf.reshape(self.raw, shape))

    def astype(self: TensorType, dtype: Any) -> TensorType:
        return type(self)(tf.cast(self.raw, dtype))

    def clip(self: TensorType, min_: float, max_: float) -> TensorType:
        return type(self)(tf.clip_by_value(self.raw, min_, max_))

    def square(self: TensorType) -> TensorType:
        return type(self)(tf.square(self.raw))

    def arctanh(self: TensorType) -> TensorType:
        return type(self)(tf.atanh(self.raw))

    def sum(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        if self.raw.dtype == tf.bool:
            return self.astype(tf.int64).sum(axis=axis, keepdims=keepdims)
        return type(self)(tf.reduce_sum(self.raw, axis=axis, keepdims=keepdims))

    def prod(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        if self.raw.dtype == tf.bool:
            return self.astype(tf.int64).prod(axis=axis, keepdims=keepdims)
        return type(self)(tf.reduce_prod(self.raw, axis=axis, keepdims=keepdims))

    def mean(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        if self.raw.dtype not in [tf.float16, tf.float32, tf.float64]:
            raise ValueError(
                f"Can only calculate the mean of floating types. Got {self.raw.dtype} instead."
            )
        return type(self)(tf.reduce_mean(self.raw, axis=axis, keepdims=keepdims))

    def min(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        return type(self)(tf.reduce_min(self.raw, axis=axis, keepdims=keepdims))

    def max(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        return type(self)(tf.reduce_max(self.raw, axis=axis, keepdims=keepdims))

    def minimum(self: TensorType, other: TensorOrScalar) -> TensorType:
        return type(self)(tf.minimum(self.raw, unwrap1(other)))

    def maximum(self: TensorType, other: TensorOrScalar) -> TensorType:
        return type(self)(tf.maximum(self.raw, unwrap1(other)))

    def argmin(self: TensorType, axis: Optional[int] = None) -> TensorType:
        return type(self)(tf.argmin(self.raw, axis=axis))

    def argmax(self: TensorType, axis: Optional[int] = None) -> TensorType:
        return type(self)(tf.argmax(self.raw, axis=axis))

    def argsort(self: TensorType, axis: Optional[int] = -1) -> TensorType:
        return type(self)(tf.argsort(self.raw, axis=axis))

    def sort(self: TensorType, axis: Optional[int] = -1) -> TensorType:
        return type(self)(tf.sort(self.raw, axis=axis))

    @samedevice
    def uniform(
        self: TensorType, shape: ShapeOrScalar, low: float = 0.0, high: float = 1.0
    ) -> TensorType:
        if not isinstance(shape, Iterable):
            shape = (shape,)
        return type(self)(
            tf.random.uniform(shape, minval=low, maxval=high, dtype=self.raw.dtype)
        )

    @samedevice
    def normal(
        self: TensorType, shape: ShapeOrScalar, mean: float = 0.0, stddev: float = 1.0
    ) -> TensorType:
        if not isinstance(shape, Iterable):
            shape = (shape,)
        return type(self)(
            tf.random.normal(shape, mean=mean, stddev=stddev, dtype=self.raw.dtype)
        )

    @samedevice
    def ones(self: TensorType, shape: ShapeOrScalar) -> TensorType:
        return type(self)(tf.ones(shape, dtype=self.raw.dtype))

    @samedevice
    def zeros(self: TensorType, shape: ShapeOrScalar) -> TensorType:
        return type(self)(tf.zeros(shape, dtype=self.raw.dtype))

    def ones_like(self: TensorType) -> TensorType:
        return type(self)(tf.ones_like(self.raw))

    def zeros_like(self: TensorType) -> TensorType:
        return type(self)(tf.zeros_like(self.raw))

    def full_like(self: TensorType, fill_value: float) -> TensorType:
        fill_value = tf.cast(fill_value, self.raw.dtype)
        return type(self)(tf.fill(self.raw.shape, fill_value))

    @samedevice
    def onehot_like(
        self: TensorType, indices: TensorType, *, value: float = 1
    ) -> TensorType:
        if self.ndim != 2:
            raise ValueError("onehot_like only supported for 2D tensors")
        if indices.ndim != 1:
            raise ValueError("onehot_like requires 1D indices")
        if len(indices) != len(self):
            raise ValueError("length of indices must match length of tensor")
        value = tf.cast(value, self.raw.dtype)
        return type(self)(
            tf.one_hot(
                indices.raw,
                depth=self.raw.shape[-1],
                on_value=value,
                dtype=self.raw.dtype,
            )
        )

    @samedevice
    def from_numpy(self: TensorType, a: Any) -> TensorType:
        return type(self)(tf.convert_to_tensor(a))

    def _concatenate(
        self: TensorType, tensors: Iterable[TensorType], axis: int = 0
    ) -> TensorType:
        # concatenates only "tensors", but not "self"
        tensors_ = unwrap_(*tensors)
        return type(self)(tf.concat(tensors_, axis=axis))

    def _stack(
        self: TensorType, tensors: Iterable[TensorType], axis: int = 0
    ) -> TensorType:
        # stacks only "tensors", but not "self"
        tensors_ = unwrap_(*tensors)
        return type(self)(tf.stack(tensors_, axis=axis))

    def transpose(self: TensorType, axes: Optional[Axes] = None) -> TensorType:
        if axes is None:
            axes = tuple(range(self.ndim - 1, -1, -1))
        return type(self)(tf.transpose(self.raw, perm=axes))

    def all(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        assert_bool(self)
        return type(self)(tf.reduce_all(self.raw, axis=axis, keepdims=keepdims))

    def any(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        assert_bool(self)
        return type(self)(tf.reduce_any(self.raw, axis=axis, keepdims=keepdims))

    def logical_and(self: TensorType, other: TensorOrScalar) -> TensorType:
        assert_bool(self)
        assert_bool(other)
        return type(self)(tf.logical_and(self.raw, unwrap1(other)))

    def logical_or(self: TensorType, other: TensorOrScalar) -> TensorType:
        assert_bool(self)
        assert_bool(other)
        return type(self)(tf.logical_or(self.raw, unwrap1(other)))

    def logical_not(self: TensorType) -> TensorType:
        assert_bool(self)
        return type(self)(tf.logical_not(self.raw))

    def exp(self: TensorType) -> TensorType:
        return type(self)(tf.exp(self.raw))

    def log(self: TensorType) -> TensorType:
        return type(self)(tf.math.log(self.raw))

    def log2(self: TensorType) -> TensorType:
        return type(self)(tf.math.log(self.raw) / tf.math.log(2.0))

    def log10(self: TensorType) -> TensorType:
        return type(self)(tf.math.log(self.raw) / tf.math.log(10.0))

    def log1p(self: TensorType) -> TensorType:
        return type(self)(tf.math.log1p(self.raw))

    def tile(self: TensorType, multiples: Axes) -> TensorType:
        multiples = unwrap1(multiples)
        if len(multiples) != self.ndim:
            raise ValueError("multiples requires one entry for each dimension")
        return type(self)(tf.tile(self.raw, multiples))

    def softmax(self: TensorType, axis: int = -1) -> TensorType:
        return type(self)(tf.nn.softmax(self.raw, axis=axis))

    def log_softmax(self: TensorType, axis: int = -1) -> TensorType:
        return type(self)(tf.nn.log_softmax(self.raw, axis=axis))

    def squeeze(self: TensorType, axis: Optional[AxisAxes] = None) -> TensorType:
        return type(self)(tf.squeeze(self.raw, axis=axis))

    def expand_dims(self: TensorType, axis: int) -> TensorType:
        return type(self)(tf.expand_dims(self.raw, axis=axis))

    @samedevice
    def full(self: TensorType, shape: ShapeOrScalar, value: float) -> TensorType:
        if not isinstance(shape, Iterable):
            shape = (shape,)
        return type(self)(tf.fill(shape, value))

    def index_update(
        self: TensorType, indices: Any, values: TensorOrScalar
    ) -> TensorType:
        indices, values_ = unwrap_(indices, values)
        del values
        if isinstance(indices, tuple):
            indices = unwrap_(*indices)

        x = self.raw
        if isinstance(indices, int):
            if isinstance(values_, int) or isinstance(values_, float):
                values_ = tf.fill(x.shape[-1:], values_)
            return type(self)(
                tf.tensor_scatter_nd_update(x, [[indices]], values_[None])
            )
        elif isinstance(indices, tuple) and any(
            isinstance(idx, slice) for idx in indices
        ):
            if (
                len(indices) == x.ndim == 2
                and indices[0] == index[:]
                and not isinstance(indices[1], slice)
            ):
                x = tf.transpose(x)
                if isinstance(values_, int) or isinstance(values_, float):
                    values_ = tf.fill(x.shape[-1:], values_)
                result = tf.tensor_scatter_nd_update(x, [[indices[-1]]], values_[None])
                return type(self)(tf.transpose(result))
            else:
                raise NotImplementedError  # pragma: no cover
        elif isinstance(indices, tuple):
            if all(idx.dtype in [tf.int32, tf.int64] for idx in indices):
                indices = [
                    tf.cast(idx, tf.int64) if idx.dtype == tf.int32 else idx
                    for idx in indices
                ]
            indices = tf.stack(indices, axis=-1)
            if isinstance(values_, int) or isinstance(values_, float):
                values_ = tf.fill((indices.shape[0],), values_)
            return type(self)(tf.tensor_scatter_nd_update(x, indices, values_))
        else:
            raise ValueError  # pragma: no cover

    @samedevice
    def arange(
        self: TensorType,
        start: int,
        stop: Optional[int] = None,
        step: Optional[int] = None,
    ) -> TensorType:
        if step is None:
            step = 1
        if stop is None:
            stop = start
            start = 0
        return type(self)(tf.range(start, stop, step))

    def cumsum(self: TensorType, axis: Optional[int] = None) -> TensorType:
        if axis is None:
            x = tf.reshape(self.raw, (-1,))
            return type(self)(tf.cumsum(x, axis=0))
        return type(self)(tf.cumsum(self.raw, axis=axis))

    def flip(self: TensorType, axis: Optional[AxisAxes] = None) -> TensorType:
        if axis is None:
            axis = tuple(range(self.ndim))
        if not isinstance(axis, Iterable):
            axis = (axis,)
        return type(self)(tf.reverse(self.raw, axis=axis))

    def meshgrid(
        self: TensorType, *tensors: TensorType, indexing: str = "xy"
    ) -> Tuple[TensorType, ...]:
        tensors = unwrap_(*tensors)
        outputs = tf.meshgrid(self.raw, *tensors, indexing=indexing)
        return tuple(type(self)(out) for out in outputs)

    def pad(
        self: TensorType,
        paddings: Tuple[Tuple[int, int], ...],
        mode: str = "constant",
        value: float = 0,
    ) -> TensorType:
        if len(paddings) != self.ndim:
            raise ValueError("pad requires a tuple for each dimension")
        for p in paddings:
            if len(p) != 2:
                raise ValueError("pad requires a tuple for each dimension")
        if not (mode == "constant" or mode == "reflect"):
            raise ValueError("pad requires mode 'constant' or 'reflect'")
        if mode == "reflect":
            # PyTorch's pad has limited support for 'reflect' padding
            if self.ndim != 3 and self.ndim != 4:
                raise NotImplementedError  # pragma: no cover
            k = self.ndim - 2
            if paddings[:k] != ((0, 0),) * k:
                raise NotImplementedError  # pragma: no cover
        return type(self)(tf.pad(self.raw, paddings, mode=mode, constant_values=value))

    def isnan(self: TensorType) -> TensorType:
        return type(self)(tf.math.is_nan(self.raw))

    def isinf(self: TensorType) -> TensorType:
        return type(self)(tf.math.is_inf(self.raw))

    def crossentropy(self: TensorType, labels: TensorType) -> TensorType:
        if self.ndim != 2:
            raise ValueError("crossentropy only supported for 2D logits tensors")
        if self.shape[:1] != labels.shape:
            raise ValueError("labels must be 1D and must match the length of logits")
        return type(self)(
            tf.nn.sparse_softmax_cross_entropy_with_logits(labels.raw, self.raw)
        )

    @overload
    def _value_and_grad_fn(
        self: TensorType, f: Callable[..., TensorType]
    ) -> Callable[..., Tuple[TensorType, TensorType]]:
        ...

    @overload  # noqa: F811 (waiting for pyflakes > 2.1.1)
    def _value_and_grad_fn(
        self: TensorType, f: Callable[..., TensorType], has_aux: Literal[False]
    ) -> Callable[..., Tuple[TensorType, TensorType]]:
        ...

    @overload  # noqa: F811 (waiting for pyflakes > 2.1.1)
    def _value_and_grad_fn(
        self: TensorType,
        f: Callable[..., Tuple[TensorType, Any]],
        has_aux: Literal[True],
    ) -> Callable[..., Tuple[TensorType, Any, TensorType]]:
        ...

    def _value_and_grad_fn(  # noqa: F811 (waiting for pyflakes > 2.1.1)
        self: TensorType, f: Callable, has_aux: bool = False
    ) -> Callable[..., Tuple]:
        def value_and_grad(x: TensorType, *args: Any, **kwargs: Any) -> Tuple:
            # using tf.identity to make x independent from possible other instances of x in args
            x_ = TensorFlowTensor(tf.identity(x.raw))
            del x
            with tf.GradientTape() as tape:
                tape.watch(x_.raw)
                if has_aux:
                    loss, aux = f(x_, *args, **kwargs)
                else:
                    loss = f(x_, *args, **kwargs)
            grad = tape.gradient(loss.raw, x_.raw)
            grad = TensorFlowTensor(grad)
            assert grad.shape == x_.shape
            if has_aux:
                return loss, aux, grad
            else:
                return loss, grad

        return value_and_grad

    def sign(self: TensorType) -> TensorType:
        return type(self)(tf.sign(self.raw))

    def sqrt(self: TensorType) -> TensorType:
        return type(self)(tf.sqrt(self.raw))

    def tanh(self: TensorType) -> TensorType:
        return type(self)(tf.tanh(self.raw))

    def float32(self: TensorType) -> TensorType:
        return self.astype(tf.float32)

    def where(self: TensorType, x: TensorOrScalar, y: TensorOrScalar) -> TensorType:
        x, y = unwrap_(x, y)
        return type(self)(tf.where(self.raw, x, y))

    def matmul(self: TensorType, other: TensorType) -> TensorType:
        if self.ndim != 2 or other.ndim != 2:
            raise ValueError(
                f"matmul requires both tensors to be 2D, got {self.ndim}D and {other.ndim}D"
            )
        return type(self)(tf.matmul(self.raw, other.raw))

    @common_dtype
    def __lt__(self: TensorType, other: TensorOrScalar) -> TensorType:
        return type(self)(self.raw.__lt__(unwrap1(other)))

    @common_dtype
    def __le__(self: TensorType, other: TensorOrScalar) -> TensorType:
        return type(self)(self.raw.__le__(unwrap1(other)))

    @common_dtype
    def __eq__(self: TensorType, other: TensorOrScalar) -> TensorType:  # type: ignore
        return type(self)(self.raw.__eq__(unwrap1(other)))

    @common_dtype
    def __ne__(self: TensorType, other: TensorOrScalar) -> TensorType:  # type: ignore
        return type(self)(self.raw.__ne__(unwrap1(other)))

    @common_dtype
    def __gt__(self: TensorType, other: TensorOrScalar) -> TensorType:
        return type(self)(self.raw.__gt__(unwrap1(other)))

    @common_dtype
    def __ge__(self: TensorType, other: TensorOrScalar) -> TensorType:
        return type(self)(self.raw.__ge__(unwrap1(other)))

    def __getitem__(self: TensorType, index: Any) -> TensorType:
        if isinstance(index, tuple):
            index = tuple(x.raw if isinstance(x, Tensor) else x for x in index)
            basic = all(
                x is None or x is Ellipsis or isinstance(x, int) or isinstance(x, slice)
                for x in index
            )
            if not basic:
                # workaround for missing support for this in TensorFlow
                index = [tf.convert_to_tensor(x) for x in index]
                shapes = [tuple(x.shape) for x in index]
                shape = tuple(max(x) for x in zip(*shapes))
                int64 = any(x.dtype == tf.int64 for x in index)
                for i in range(len(index)):
                    t = index[i]
                    if int64:
                        t = tf.cast(t, tf.int64)  # pragma: no cover
                    assert t.ndim == len(shape)
                    tiling = []
                    for b, k in zip(shape, t.shape):
                        if k == 1:
                            tiling.append(b)
                        elif k == b:
                            tiling.append(1)
                        else:
                            raise ValueError(  # pragma: no cover
                                f"{tuple(t.shape)} cannot be broadcasted to {shape}"
                            )
                    index[i] = tf.tile(t, tiling)
                index = tf.stack(index, axis=-1)
                return type(self)(tf.gather_nd(self.raw, index))
        elif (
            isinstance(index, range)
            or isinstance(index, list)
            or isinstance(index, np.ndarray)
        ):
            return type(self)(tf.gather(self.raw, index))
        elif isinstance(index, Tensor):
            if index.raw.dtype == tf.bool:
                return type(self)(self.raw.__getitem__(index.raw))
            else:
                return type(self)(tf.gather(self.raw, index.raw))
        return type(self)(self.raw.__getitem__(index))

    def take_along_axis(self: TensorType, index: TensorType, axis: int) -> TensorType:
        axis = batch_dims = axis % self.ndim
        if axis != self.ndim - 1:
            raise NotImplementedError(
                "take_along_axis is currently only supported for the last axis"
            )
        return type(self)(
            tf.gather(self.raw, index.raw, axis=axis, batch_dims=batch_dims)
        )

    def bool(self: TensorType) -> TensorType:
        return self.astype(tf.bool)
