from typing import Callable, Iterator, Optional, Self, SupportsIndex, Any, overload
from types import NoneType
from collections import deque

import numpy as np
import weakref

import lucid
from lucid import types
from lucid.types import (
    _ArrayOrScalar,
    _NumPyArray,
    _MLXArray,
    _Scalar,
    _DeviceType,
    _BuiltinNumeric,
    Numeric,
)

from lucid._tensor.base import _TensorBase, _TensorInplace
from lucid._backend.core import BackwardOperation, Operation, noop
from lucid._backend.metal import mx, parse_mlx_indexing, check_metal_availability


_HookType = Callable[["Tensor", _NumPyArray | _MLXArray], None]

_dtype_map = {int: types.Int64, float: types.Float64, complex: types.Complex64}


class Tensor(_TensorBase, _TensorInplace):
    def __init__(
        self,
        data: _ArrayOrScalar | _MLXArray,
        requires_grad: bool = False,
        keep_grad: bool = False,
        dtype: _BuiltinNumeric | Numeric | None = None,
        device: _DeviceType = "cpu",
    ) -> None:
        self._is_free = False
        self._is_bool_tensor = False

        if dtype is bool:
            self._is_bool_tensor = True
            dtype = None
        if dtype in {int, float, complex}:
            dtype = _dtype_map[dtype]

        dtype: Numeric | None
        assert isinstance(dtype, (Numeric, NoneType)), type(dtype)

        if dtype is not None and dtype.is_bit_free:
            raise TypeError(
                f"Implicit dtypes is not supported for Tensor instantiation."
            )
        if device not in {"cpu", "gpu"}:
            raise lucid.UnknownDeviceError(device)

        self.data: _NumPyArray | _MLXArray
        if not isinstance(data, (_NumPyArray, _MLXArray)):
            self.data = np.array(data, dtype=dtype.cpu if dtype is not None else dtype)

            if isinstance(data, (_Scalar, list, tuple)):
                self._is_free = True

        elif isinstance(data, _NumPyArray):
            self.data = data
            if dtype is not None and data.dtype != dtype.cpu:
                self.data = data.astype(dtype.cpu)

        elif isinstance(data, _MLXArray):
            check_metal_availability()
            device = "gpu"
            self.data = data
            if dtype is not None and data.dtype != dtype.gpu:
                self.data = data.astype(dtype.gpu)

        if device == "gpu" and isinstance(self.data, _NumPyArray):
            check_metal_availability()
            self.data = mx.array(self.data)

        if "bool" in str(self.data.dtype):
            self._is_bool_tensor = True
        else:
            if self._is_bool_tensor:
                self.data = self.data.astype(bool if device == "cpu" else mx.bool_)

        self._op: Operation | None = None
        self._backward_op: BackwardOperation = noop
        self._prev: list[Tensor] = []
        self._backward_hooks: list[_HookType] = []
        self._version: int = 0

        self.grad: Optional[_NumPyArray | _MLXArray] = None
        self.device = device
        self.requires_grad = requires_grad and lucid.grad_enabled()
        self.keep_grad = keep_grad

        self.dtype: Numeric | bool
        if self._is_bool_tensor:
            self.dtype = bool
        else:
            self.dtype = (
                dtype if dtype is not None else types.to_numeric_type(self.data.dtype)
            )

    @property
    def is_leaf(self) -> bool:
        return self.requires_grad and len(self._prev) == 0

    @property
    def is_free(self) -> bool:
        return self._is_free

    @classmethod
    def is_all_free(cls, *tensors: Self) -> bool:
        return all(t.is_free for t in tensors)

    @classmethod
    def copy_data(cls, data: _NumPyArray | _MLXArray) -> _NumPyArray | _MLXArray:
        if isinstance(data, _NumPyArray):
            return data.copy()
        elif isinstance(data, _MLXArray):
            return mx.array(data)
        else:
            raise TypeError(f"Unexpected type: '{type(data).__name__}'")

    @classmethod
    def copy_grad(cls, grad: _NumPyArray | _MLXArray) -> _NumPyArray | _MLXArray:
        return cls.copy_data(data=grad)

    def new_tensor(self) -> Tensor:
        return Tensor(
            self.data, self.requires_grad, self.keep_grad, self.dtype, self.device
        )

    def free(self) -> Self:
        self._is_free = True
        return self

    def eval(self) -> Self:
        if self.is_gpu():
            mx.eval(self.data)
            stopped = mx.stop_gradient(self.data)
            if stopped is not None:
                self.data = stopped
        return self

    def backward(self, retain_grad: bool = False, retain_graph: bool = False) -> None:
        import lucid.autograd as autograd

        autograd.backward(self, retain_grad=retain_grad, retain_graph=retain_graph)

    def register_hook(self, hook: _HookType) -> Callable:
        self._backward_hooks.append(hook)
        return lambda: self._backward_hooks.remove(hook)

    @property
    def shape(self) -> tuple[int, ...]:
        return self.data.shape

    @property
    def ndim(self) -> int:
        return self.data.ndim

    @property
    def size(self) -> int:
        return self.data.size

    @property
    def flops(self) -> int:
        total = 0
        queue = deque([self])
        visited = set()

        while queue:
            cur = queue.popleft()
            visited.add(cur)
            if cur._op is not None:
                total += cur._op.flops

            for next in cur._prev:
                if next not in visited:
                    queue.appendleft(next)

        return total

    def item(self) -> _Scalar | bool:
        if self.size > 1:
            raise ValueError(
                "Tensor must be 0-dimensional(scalar) to pop its item. "
                "Use `tensor.data` to access the data directly.",
            )
        item = self.data[..., 0] if self.ndim > 0 else self.data
        if self.dtype is bool:
            return bool(item)

        if item % 1 == 0:
            return int(item)
        else:
            return float(item)

    def tolist(self) -> list:
        return self.data.tolist()

    def numpy(self) -> _NumPyArray:
        return np.array(self.data)

    def mlx(self) -> _MLXArray:
        return mx.array(self.data)

    def detach(self) -> Self:
        data_copy = Tensor.copy_data(self.data)
        return Tensor(
            data_copy, requires_grad=False, dtype=self.dtype, device=self.device
        )

    def zero(self) -> None:
        self._version += 1
        if self.is_cpu():
            self.data = np.zeros_like(self.data)
        else:
            self.data = mx.zeros_like(self.data)

    def zero_grad(self) -> None:
        self.grad = None

    def clear_node(self, clear_op: bool = True) -> None:
        self._prev = []
        self._backward_op = noop
        if clear_op:
            self._op = None

    def astype(self, dtype: type | Numeric) -> Self:
        new_dtype = dtype
        if isinstance(dtype, Numeric):
            self._is_bool_tensor = False
            if dtype.is_bit_free:
                new_dtype = dtype.auto_parse(self.data.dtype, device=self.device)
            else:
                new_dtype = dtype.parse(device=self.device)

        elif dtype is bool:
            self._is_bool_tensor = True
            if self.is_gpu():
                new_dtype = mx.bool_

        self.data = self.data.astype(new_dtype)
        if self._is_bool_tensor:
            self.dtype = bool
        else:
            self.dtype = types.to_numeric_type(self.data.dtype)

        self._version += 1
        return self

    @overload
    def to(self, device: _DeviceType) -> Self: ...

    @overload
    def to(self, dtype: type | Numeric) -> Self: ...

    def to(self, device_or_dtype: _DeviceType | type | Numeric) -> Self:
        if isinstance(device_or_dtype, str):
            device = device_or_dtype
            if self.device == device:
                return self

            if device == "cpu":
                self.data = np.array(self.data)
                if self.grad is not None:
                    self.grad = np.array(self.grad)

            elif device == "gpu":
                check_metal_availability()
                self.data = mx.array(self.data)
                if self.grad is not None:
                    self.grad = mx.array(self.grad)

            else:
                raise lucid.UnknownDeviceError(device)

            self.device = device
            return self
        else:
            dtype = device_or_dtype
            return self.astype(dtype)

    def is_cpu(self) -> bool:
        return self.device == "cpu"

    def is_gpu(self) -> bool:
        return self.device == "gpu"

    def __getitem__(self, idx: SupportsIndex | Self) -> Self:
        new_idx = idx
        if isinstance(idx, Tensor):
            new_idx = idx.data
            if self.is_gpu() and idx.is_free:
                new_idx = mx.array(idx.data)

        if isinstance(idx, tuple):
            new_idx = tuple()
            for id in idx:
                if isinstance(id, Tensor):
                    if self.is_gpu() and id.is_free:
                        id = mx.array(id.data)
                    else:
                        id = id.data
                new_idx += (id,)

        if self.is_gpu():
            new_idx = parse_mlx_indexing(new_idx)
        else:
            if isinstance(idx, Tensor) and idx.is_gpu():
                raise lucid.DeviceMismatchError(to="gpu", from_="cpu")

        sliced_data = self.data[new_idx]
        new_tensor = Tensor(
            sliced_data, self.requires_grad, self.keep_grad, self.dtype, self.device
        )

        def _grad_func() -> None:
            if self.grad is None:
                self.grad = (
                    np.zeros_like(self.data)
                    if self.is_cpu()
                    else mx.zeros_like(self.data)
                )
            if new_tensor.grad is None:
                return
            new_grad = lucid._match_grad_shape(
                self.data[new_idx], new_tensor.grad, device=self.device
            )
            lucid._set_tensor_grad(self, new_grad, at=new_idx)

        if self.requires_grad and lucid.grad_enabled():
            new_tensor._prev = [self]
            new_tensor._backward_op = BackwardOperation(
                forward_op_ref=None,
                grad_func=None,
                tensor_refs=(weakref.ref(self),),
                device=self.device,
                versions=(self._version,),
                custom_closure=_grad_func,
            )

        if self.is_free:
            new_tensor.free()

        return new_tensor

    def __setitem__(
        self, idx: SupportsIndex | Self, value: Self | _ArrayOrScalar
    ) -> None:
        if self.requires_grad:
            raise RuntimeError(
                "Cannot perform in-place item setting on a "
                + "Tensor that requires gradients. "
            )

        new_idx = idx
        if isinstance(idx, Tensor):
            new_idx = idx.data
        elif isinstance(idx, tuple):
            new_idx = tuple()
            for id in idx:
                if isinstance(id, Tensor):
                    id = id.data
                new_idx += (id,)

        if self.is_gpu():
            new_idx = parse_mlx_indexing(new_idx)

        data: _ArrayOrScalar
        if isinstance(value, Tensor):
            if self.device != value.device:
                if value.is_free:
                    data = value.to(self.device).data
                else:
                    raise lucid.DeviceMismatchError(to=self.device, from_=value.device)
            data = value.data
        else:
            data = value
        self.data[new_idx] = data

    def __len__(self) -> int:
        if self.ndim == 0:
            return self.size
        else:
            return self.shape[0]

    def __iter__(self) -> Iterator[Self]:
        for i in range(self.shape[0]):
            yield self[i]

    def __repr__(self) -> str:
        return f"Tensor({self.data}, grad={self.grad}, device={self.device})"

    def __str__(self) -> str:
        return str(self.data)

    def __hash__(self) -> int:
        return hash(id(self))

    def __deepcopy__(self, *args: Any) -> Self:
        cls = self.__class__
        copied_data = Tensor.copy_data(self.data)

        if cls is Tensor:
            new = Tensor(
                copied_data, self.requires_grad, self.keep_grad, self.dtype, self.device
            )
        else:
            base = Tensor(copied_data, dtype=self.dtype, device=self.device)
            new = cls(base)

        if self.grad is not None and (
            self.keep_grad or getattr(new, "keep_grad", False)
        ):
            new.grad = Tensor.copy_grad(self.grad)
        else:
            new.grad = None

        new._op = None
        new._backward_op = noop
        new._prev = []
        new._backward_hooks = []

        new._is_free = self._is_free
        new._is_bool_tensor = self._is_bool_tensor

        return new

    def __bool__(self) -> bool:
        if self.data.size != 1:
            raise RuntimeError(
                "Boolean value of Tensor with more than one value is ambiguous. "
                "Use `tensor.any()` or `tensor.all()` instead."
            )
        return bool(self.data.item())

    def any(self, axis: int | None = None, keepdims: bool = False) -> bool | Self:
        if self.is_cpu():
            result = np.any(self.data, axis=axis, keepdims=keepdims)
            return bool(result) if axis is None else Tensor(result, device="cpu")
        else:
            mx.eval(self.data)
            result = mx.any(self.data, axis=axis, keepdims=keepdims)
            return bool(result.item()) if axis is None else Tensor(result, device="gpu")

    def all(self, axis=None, keepdims=False) -> bool | Self:
        if self.is_cpu():
            result = np.all(self.data, axis=axis, keepdims=keepdims)
            return bool(result) if axis is None else Tensor(result, device="cpu")
        else:
            mx.eval(self.data)
            result = mx.all(self.data, axis=axis, keepdims=keepdims)
            return bool(result.item()) if axis is None else Tensor(result, device="gpu")

    def char(self) -> Self:
        return self.astype(lucid.Char)

    def short(self) -> Self:
        return self.astype(lucid.Short)

    def int(self) -> Self:
        return self.astype(lucid.Int)

    def long(self) -> Self:
        return self.astype(lucid.Long)

    def half(self) -> Self:
        return self.astype(lucid.Half)

    def float(self) -> Self:
        return self.astype(lucid.Float)

    def double(self) -> Self:
        return self.astype(lucid.Double)

    def bool(self) -> Self:
        return self.astype(bool)
