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

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

from .tensor import Tensor
from .tensor import TensorOrScalar

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

if TYPE_CHECKING:
    import torch  # for static analyzers
    from .extensions import NormsMethods  # noqa: F401
else:
    # lazy import in PyTorchTensor
    torch = None


# stricter TensorType to get additional type information from the raw method
TensorType = TypeVar("TensorType", bound="PyTorchTensor")


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


class PyTorchTensor(BaseTensor):
    __slots__ = ()

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

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

    @property
    def raw(self) -> "torch.Tensor":
        return cast(torch.Tensor, super().raw)

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

    def numpy(self: TensorType) -> Any:
        a = self.raw.detach().cpu().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) -> Union[int, float, bool]:
        return self.raw.item()

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

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

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

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

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

    def arctanh(self: TensorType) -> TensorType:
        """
        improve once this issue has been fixed:
        https://github.com/pytorch/pytorch/issues/10324
        """
        return type(self)(0.5 * (torch.log1p(self.raw) - torch.log1p(-self.raw)))

    def sum(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        if axis is None and not keepdims:
            return type(self)(self.raw.sum())
        if axis is None:
            axis = tuple(range(self.ndim))
        return type(self)(self.raw.sum(dim=axis, keepdim=keepdims))

    def prod(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        if axis is None and not keepdims:
            return type(self)(self.raw.prod())
        if axis is None:
            axis = tuple(range(self.ndim))
        elif not isinstance(axis, Iterable):
            axis = (axis,)
        x = self.raw
        for i in sorted(axis, reverse=True):
            x = x.prod(i, keepdim=keepdims)
        return type(self)(x)

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

    def min(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        """
        simplify once this issue has been fixed:
        https://github.com/pytorch/pytorch/issues/28213
        """
        if axis is None and not keepdims:
            return type(self)(self.raw.min())
        if axis is None:
            axis = tuple(range(self.ndim))
        elif not isinstance(axis, Iterable):
            axis = (axis,)
        x = self.raw
        for i in sorted(axis, reverse=True):
            x, _ = x.min(i, keepdim=keepdims)
        return type(self)(x)

    def max(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        """
        simplify once this issue has been fixed:
        https://github.com/pytorch/pytorch/issues/28213
        """
        if axis is None and not keepdims:
            return type(self)(self.raw.max())
        if axis is None:
            axis = tuple(range(self.ndim))
        elif not isinstance(axis, Iterable):
            axis = (axis,)
        x = self.raw
        for i in sorted(axis, reverse=True):
            x, _ = x.max(i, keepdim=keepdims)
        return type(self)(x)

    def minimum(self: TensorType, other: TensorOrScalar) -> TensorType:
        if isinstance(other, Tensor):
            other_ = other.raw
        elif isinstance(other, int) or isinstance(other, float):
            other_ = torch.full_like(self.raw, other)
        else:
            raise TypeError(
                "expected x to be a Tensor, int or float"
            )  # pragma: no cover
        return type(self)(torch.min(self.raw, other_))

    def maximum(self: TensorType, other: TensorOrScalar) -> TensorType:
        if isinstance(other, Tensor):
            other_ = other.raw
        elif isinstance(other, int) or isinstance(other, float):
            other_ = torch.full_like(self.raw, other)
        else:
            raise TypeError(
                "expected x to be a Tensor, int or float"
            )  # pragma: no cover
        return type(self)(torch.max(self.raw, other_))

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

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

    def argsort(self: TensorType, axis: int = -1) -> TensorType:
        return type(self)(self.raw.argsort(dim=axis))

    def sort(self: TensorType, axis: int = -1) -> TensorType:
        return type(self)(self.raw.sort(dim=axis).values)  # type: ignore

    def uniform(
        self: TensorType, shape: ShapeOrScalar, low: float = 0.0, high: float = 1.0
    ) -> TensorType:
        return type(self)(
            torch.rand(shape, dtype=self.raw.dtype, device=self.raw.device)
            * (high - low)
            + low
        )

    def normal(
        self: TensorType, shape: ShapeOrScalar, mean: float = 0.0, stddev: float = 1.0
    ) -> TensorType:
        return type(self)(
            torch.randn(shape, dtype=self.raw.dtype, device=self.raw.device) * stddev
            + mean
        )

    def ones(self: TensorType, shape: ShapeOrScalar) -> TensorType:
        return type(self)(
            torch.ones(shape, dtype=self.raw.dtype, device=self.raw.device)
        )

    def zeros(self: TensorType, shape: ShapeOrScalar) -> TensorType:
        return type(self)(
            torch.zeros(shape, dtype=self.raw.dtype, device=self.raw.device)
        )

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

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

    def full_like(self: TensorType, fill_value: float) -> TensorType:
        return type(self)(torch.full_like(self.raw, fill_value))

    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")
        x = torch.zeros_like(self.raw)
        rows = np.arange(x.shape[0])
        x[rows, indices.raw] = value
        return type(self)(x)

    def from_numpy(self: TensorType, a: Any) -> TensorType:
        return type(self)(torch.as_tensor(a, device=self.raw.device))

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

    def _stack(
        self: TensorType, tensors: Iterable[TensorType], axis: int = 0
    ) -> TensorType:
        # stacks only "tensors", but not "self"
        tensors_ = unwrap_(*tensors)
        return type(self)(torch.stack(tensors_, dim=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)(self.raw.permute(*axes))

    def all(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        assert_bool(self)
        if axis is None and not keepdims:
            return type(self)(self.raw.all())
        if axis is None:
            axis = tuple(range(self.ndim))
        elif not isinstance(axis, Iterable):
            axis = (axis,)
        x = self.raw
        for i in sorted(axis, reverse=True):
            x = x.all(i, keepdim=keepdims)
        return type(self)(x)

    def any(
        self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
    ) -> TensorType:
        assert_bool(self)
        if axis is None and not keepdims:
            return type(self)(self.raw.any())
        if axis is None:
            axis = tuple(range(self.ndim))
        elif not isinstance(axis, Iterable):
            axis = (axis,)
        x = self.raw
        for i in sorted(axis, reverse=True):
            x = x.any(i, keepdim=keepdims)
        return type(self)(x)

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

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

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

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

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

    def log2(self: TensorType) -> TensorType:
        return type(self)(torch.log2(self.raw))

    def log10(self: TensorType) -> TensorType:
        return type(self)(torch.log10(self.raw))

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

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

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

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

    def squeeze(self: TensorType, axis: Optional[AxisAxes] = None) -> TensorType:
        if axis is None:
            return type(self)(self.raw.squeeze())
        if not isinstance(axis, Iterable):
            axis = (axis,)
        x = self.raw
        for i in sorted(axis, reverse=True):
            if x.shape[i] != 1:
                raise ValueError(
                    "cannot select an axis to squeeze out which has size not equal to one"
                )
            x = x.squeeze(dim=i)
        return type(self)(x)

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

    def full(self: TensorType, shape: ShapeOrScalar, value: float) -> TensorType:
        if not isinstance(shape, Iterable):
            shape = (shape,)
        return type(self)(
            torch.full(shape, value, dtype=self.raw.dtype, device=self.raw.device)
        )

    def index_update(
        self: TensorType, indices: Any, values: TensorOrScalar
    ) -> TensorType:
        indices, values_ = unwrap_(indices, values)
        if isinstance(indices, tuple):
            indices = unwrap_(*indices)
        x = self.raw.clone()
        x[indices] = values_
        return type(self)(x)

    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)(
            torch.arange(start=start, end=stop, step=step, device=self.raw.device)
        )

    def cumsum(self: TensorType, axis: Optional[int] = None) -> TensorType:
        if axis is None:
            return type(self)(self.raw.reshape(-1).cumsum(dim=0))
        return type(self)(self.raw.cumsum(dim=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)(self.raw.flip(dims=axis))

    def meshgrid(
        self: TensorType, *tensors: TensorType, indexing: str = "xy"
    ) -> Tuple[TensorType, ...]:
        tensors = unwrap_(*tensors)
        if indexing == "ij" or len(tensors) == 0:
            outputs = torch.meshgrid(self.raw, *tensors)  # type: ignore
        elif indexing == "xy":
            outputs = torch.meshgrid(tensors[0], self.raw, *tensors[1:])  # type: ignore
        else:
            raise ValueError(  # pragma: no cover
                f"Valid values for indexing are 'xy' and 'ij', got {indexing}"
            )
        results = [type(self)(out) for out in outputs]
        if indexing == "xy" and len(results) >= 2:
            results[0], results[1] = results[1], results[0]
        return tuple(results)

    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
            paddings = paddings[k:]
        paddings_ = list(x for p in reversed(paddings) for x in p)
        return type(self)(
            torch.nn.functional.pad(self.raw, paddings_, mode=mode, value=value)
        )

    def isnan(self: TensorType) -> TensorType:
        return type(self)(torch.isnan(self.raw))

    def isinf(self: TensorType) -> TensorType:
        return type(self)(torch.isinf(self.raw))  # type: ignore

    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)(
            torch.nn.functional.cross_entropy(self.raw, labels.raw, reduction="none")
        )

    @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:
            x = type(self)(x.raw.clone().requires_grad_())
            if has_aux:
                loss, aux = f(x, *args, **kwargs)
            else:
                loss = f(x, *args, **kwargs)
            loss = loss.raw
            loss.backward()
            assert x.raw.grad is not None
            grad = type(self)(x.raw.grad)
            assert grad.shape == x.shape
            loss = loss.detach()
            loss = type(self)(loss)
            if has_aux:
                if isinstance(aux, PyTorchTensor):
                    aux = PyTorchTensor(aux.raw.detach())
                elif isinstance(aux, tuple):
                    aux = tuple(
                        PyTorchTensor(t.raw.detach())
                        if isinstance(t, PyTorchTensor)
                        else t
                        for t in aux
                    )
                return loss, aux, grad
            else:
                return loss, grad

        return value_and_grad

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

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

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

    def where(self: TensorType, x: TensorOrScalar, y: TensorOrScalar) -> TensorType:
        if isinstance(x, Tensor):
            x_ = x.raw
        elif isinstance(x, int) or isinstance(x, float):
            x_ = torch.full_like(self.raw, x, dtype=torch.float32)
        else:
            raise TypeError(
                "expected x to be a Tensor, int or float"
            )  # pragma: no cover
        if isinstance(y, Tensor):
            y_ = y.raw
        elif isinstance(y, int) or isinstance(y, float):
            y_ = torch.full_like(self.raw, y, dtype=torch.float32)
        return type(self)(torch.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)(torch.matmul(self.raw, other.raw))

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

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

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

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

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

    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)
        elif isinstance(index, Tensor):
            index = index.raw
        return type(self)(self.raw[index])

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

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