from __future__ import annotations
from numpy import ndarray
from typing import Optional, Self, Any
from nova._interfaces._base_tensor import TensorBase
from nova._typing import Dim, Size, Dtype, Hook
from nova.utils.hooks import HooksHandle
from nova.autograd.function import Function
from nova.autograd.engine import Context

class Tensor(TensorBase):

    _data_internal: ndarray
    _dtype_internal: Dtype
    requires_grad: bool
    grad_fn: Optional[Function]
    grad: Optional[ndarray]
    copy: bool
    _is_leaf: bool
    _backward_hooks: list[Hook]
    _retain_grad: bool
    _inputs: list[Tensor]
    _ctx: Context
    # Constructor

    def __init__(
        self,
        data: Any,
        dtype: Optional[Dtype] = None,
        requires_grad: bool = False,
        grad_fn: Optional[Function] = None,
        copy: bool = True,
    ) -> None: ...

    # State for serialization

    def __getstate__(self): ...
    def __setstate__(self, state: dict): ...

    # Autograd Operations

    def backward(
        self, gradient: Optional[Tensor | ndarray] = None, retain_graph: bool = False
    ) -> None: ...
    def zero_grad(self, set_to_none: bool = True) -> None: ...
    def requires_grad_(self, requires_grad: bool = True) -> Self: ...
    def detach(self) -> Tensor: ...
    def retain_grad(self) -> Self: ...
    def register_hook(self, hook: Hook) -> HooksHandle: ...

    # Properties

    @property
    def data(self) -> ndarray: ...
    @property
    def T(self) -> Tensor: ...
    @property
    def shape(self) -> Size: ...
    @property
    def dtype(self) -> Dtype: ...
    @property
    def strides(self) -> Size: ...
    @property
    def itemsize(self) -> int: ...
    @property
    def ndim(self) -> int: ...
    @property
    def nbytes(self) -> int: ...
    @property
    def is_leaf(self) -> bool: ...
    @property
    def is_cuda(self) -> bool: ...
    @property
    def device(self) -> str: ...
    def size(self, dim: Optional[Dim] = None) -> Dim: ...
    def dim(self) -> int: ...
    def numel(self) -> int: ...

    # Comparison Operations

    def __lt__(self, other: Tensor | float | int) -> Tensor: ...
    def __le__(self, other: Tensor | float | int) -> Tensor: ...
    def __gt__(self, other: Tensor | float | int) -> Tensor: ...
    def __ge__(self, other: Tensor | float | int) -> Tensor: ...
    def __eq__(self, other: Any) -> Tensor: ...
    def __ne__(self, other: Any) -> Tensor: ...
    def eq(self, other: Tensor) -> Tensor: ...
    def ne(self, other: Tensor) -> Tensor: ...
    def lt(self, other: Tensor) -> Tensor: ...
    def le(self, other: Tensor) -> Tensor: ...
    def gt(self, other: Tensor) -> Tensor: ...
    def ge(self, other: Tensor) -> Tensor: ...
    def maximum(self, other: Tensor | float | int) -> Tensor: ...
    def minimum(self, other: Tensor | float | int) -> Tensor: ...

    # Basic Arithmetic Operations

    def __add__(self, other: Tensor | float | int) -> Tensor: ...
    def __radd__(self, other: Tensor | float | int) -> Tensor: ...
    def __iadd__(self, other: Tensor | float | int) -> Self: ...
    def add(self, other: Tensor | float | int) -> Tensor: ...
    def add_(self, other: Tensor | float | int) -> Self: ...
    def __sub__(self, other: Tensor | float | int) -> Tensor: ...
    def __rsub__(self, other: Tensor | float | int) -> Tensor: ...
    def __isub__(self, other: Tensor | float | int) -> Self: ...
    def sub(self, other: Tensor | float | int) -> Tensor: ...
    def sub_(self, other: Tensor | float | int) -> Self: ...
    def __mul__(self, other: Tensor | float | int) -> Tensor: ...
    def __rmul__(self, other: Tensor | float | int) -> Tensor: ...
    def __imul__(self, other: Tensor | float | int) -> Self: ...
    def mul(self, other: Tensor | float | int) -> Tensor: ...
    def mul_(self, other: Tensor | float | int) -> Self: ...
    def __truediv__(self, other: Tensor | float | int) -> Tensor: ...
    def __rtruediv__(self, other: Tensor | float | int) -> Tensor: ...
    def __itruediv__(self, other: Tensor | float | int) -> Self: ...
    def div(self, other: Tensor | float | int) -> Tensor: ...
    def div_(self, other: Tensor | float | int) -> Self: ...
    def __floordiv__(self, other: Tensor | float | int) -> Tensor: ...
    def __rfloordiv__(self, other: Tensor | float | int) -> Tensor: ...
    def __ifloordiv__(self, other: Tensor | float | int) -> Self: ...
    def floor_divide(self, other: Tensor | float | int) -> Tensor: ...
    def floor_divide_(self, other: Tensor | int | float) -> Self: ...
    def __mod__(self, other: Tensor | float | int) -> Tensor: ...
    def __rmod__(self, other: Tensor | float | int) -> Tensor: ...
    def __imod__(self, other: Tensor | float | int) -> Self: ...
    def mod(self, other: Tensor | float | int) -> Tensor: ...
    def mod_(self, other: Tensor | float | int) -> Self: ...
    def __pow__(self, other: Tensor | float | int) -> Tensor: ...
    def __rpow__(self, other: Tensor | float | int) -> Tensor: ...
    def __ipow__(self, other: Tensor | float | int) -> Self: ...
    def pow_(self, other: Tensor | float | int) -> Self: ...
    def pow(self, other: Tensor | float | int) -> Tensor: ...

    # Unary Operations

    def __abs__(self) -> Tensor: ...
    def abs(self) -> Tensor: ...
    def abs_(self) -> Self: ...
    def __neg__(self) -> Tensor: ...
    def exp(self) -> Tensor: ...
    def exp_(self) -> Self: ...
    def log(self) -> Tensor: ...
    def log_(self) -> Tensor: ...
    def sqrt(self) -> Tensor: ...
    def sqrt_(self) -> Self: ...
    def floor(self) -> Tensor: ...
    def floor_(self) -> Self: ...
    def sign(self) -> Tensor: ...
    def sign_(self) -> Tensor: ...
    def ceil(self) -> Tensor: ...
    def ceil_(self) -> Tensor: ...

    # Activation Functions

    def relu(self, input: Tensor | ndarray) -> Tensor: ...
    def relu_(self) -> Self: ...
    def leaky_relu(
        self, input: Tensor | ndarray, alpha: Tensor | float = 0.01
    ) -> Tensor: ...
    def leaky_relu_(self, alpha: Tensor | float = 0.01) -> Tensor: ...
    def prelu(
        self, input: Tensor | ndarray, weight: float | ndarray = 0.25
    ) -> Tensor: ...
    def prelu_(self, weight: float | ndarray = 0.25) -> Self: ...
    def gelu(self, input: Tensor | ndarray) -> Tensor: ...
    def gelu_(self) -> Self: ...
    def sigmoid(self, input: Tensor | ndarray) -> Tensor: ...
    def sigmoid_(self) -> Self: ...

    # Linear Algebra Operations

    def __matmul__(self, other: Tensor) -> Tensor: ...
    def __rmatmul__(self, other: Tensor) -> Tensor: ...
    def __imatmul__(self, other: Tensor) -> Tensor: ...
    def dot(self, other: Tensor) -> Tensor: ...
    def inv(self) -> Tensor: ...
    def det(self) -> Tensor: ...
    def diag(self, diagonal: int = 0) -> Tensor: ...
    def trace(self) -> Tensor: ...
    def norm(
        self, ord: int | str = 2, dim: Optional[Dim] = None, keepdims: bool = False
    ) -> Tensor: ...

    # Shape Manipulation

    def reshape(self, *size: Dim) -> Tensor: ...
    def flatten(self, start_dim: int = 0, end_dim: int = -1) -> Self: ...
    def view(self, *size: Size) -> Tensor: ...
    def permute(self, *dims: Optional[Dim]) -> Tensor: ...
    def squeeze(self, dim: Optional[Dim] = None) -> Tensor: ...
    def unsqueeze(self, dim: Dim) -> Tensor: ...
    def tile(self, repeats: int | tuple[int, ...]) -> Tensor: ...
    def repeat(self, repeats: int, dim: Dim = 0) -> Tensor: ...
    def extend(self, *size: Size) -> Tensor: ...
    def pad(
        self, pad_width: tuple[tuple[int, ...], ...], mode: str = "constant"
    ) -> Tensor: ...

    # Indexing Operations

    def __getitem__(self, idx: Tensor | ndarray | slice | int) -> Tensor: ...
    def __setitem__(
        self, key: Tensor | int | slice | Any, value: Tensor | int | float | Any
    ) -> Self: ...

    # Reduction Operations

    def sum(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...
    def mean(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...
    def var(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...
    def std(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...
    def max(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...
    def min(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...
    def all(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...
    def any(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...

    # Trigonometric methods
    def sin(self) -> Tensor: ...
    def sin_(self) -> Tensor: ...
    def cos(self) -> Tensor: ...
    def cos_(self) -> Tensor: ...
    def tan(self) -> Tensor: ...
    def tan_(self) -> Tensor: ...
    def sinh(self) -> Tensor: ...
    def sinh_(self) -> Tensor: ...
    def cosh(self) -> Tensor: ...
    def cosh_(self) -> Tensor: ...
    def tanh(self) -> Tensor: ...
    def tanh_(self) -> Tensor: ...
    def cot(self) -> Tensor: ...
    def cot_(self) -> Tensor: ...
    def sec(self) -> Tensor: ...
    def sec_(self) -> Tensor: ...
    def csc(self) -> Tensor: ...
    def csc_(self) -> Tensor: ...
    def atan2(self, other: Tensor | float | int) -> Tensor: ...

    # Inverse Trigonometric methods
    def arcsin(self) -> Tensor: ...
    def arcsin_(self) -> Tensor: ...
    def arccos(self) -> Tensor: ...
    def arccos_(self) -> Tensor: ...
    def arctan(self) -> Tensor: ...
    def arctan_(self) -> Tensor: ...
    def asinh(self) -> Tensor: ...
    def asinh_(self) -> Tensor: ...
    def acosh(self) -> Tensor: ...
    def acosh_(self) -> Tensor: ...
    def atanh(self) -> Tensor: ...
    def atanh_(self) -> Tensor: ...
    def arccot(self) -> Tensor: ...
    def arccot_(self) -> Tensor: ...
    def arcsec(self) -> Tensor: ...
    def arcsec_(self) -> Tensor: ...
    def arccsc(self) -> Tensor: ...
    def arccsc_(self) -> Tensor: ...

    # Argument Operations

    def argmax(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...
    def argmin(self, dim: Optional[Dim] = None, keepdims: bool = False) -> Tensor: ...
    def argsort(self, dim: Optional[Dim] = None, kind=None, order=None) -> Tensor: ...
    def argwhere(self) -> Tensor: ...

    # Utility Operations

    def split(self, sections: int | list[int], dim: int = 0) -> list[Tensor]: ...
    def clamp(
        self, min: Optional[float] = None, max: Optional[float] = None
    ) -> Tensor: ...
    def clamp_(
        self, min: Optional[float] = None, max: Optional[float] = None
    ) -> Self: ...
    def allclose(
        self, other, rtol=1e-5, atol=1e-8, equal_nan: bool = False
    ) -> bool: ...
    def clone(self) -> Tensor: ...
    def item(self) -> float | int: ...

    # Type Casting and Conversion

    def type(self, dtype: Optional[Dtype] = None) -> Tensor | str: ...
    def is_contiguous(self) -> bool: ...
    def contiguous(self) -> Tensor: ...
    def to(self, dtype: Dtype) -> None: ...
    def float(self) -> None: ...
    def double(self) -> None: ...
    def int(self) -> None: ...
    def long(self) -> None: ...
    def bool(self) -> None: ...
    def numpy(self) -> ndarray: ...

    # In-place Initialization Operations

    def normal_(self, mean: float = 0, std: float = 1) -> None: ...
    def uniform_(self, low: float = -1, high: float = 1) -> None: ...
    def random_(self) -> None: ...
    def zero_(self) -> None: ...
    def ones_(self) -> None: ...
    def fill_(self, value: Any) -> None: ...
    def copy_(self, src: Tensor) -> Self: ...

    # Low-level Memory Operations

    def as_strided(self, size: Size, strides: Size) -> Tensor: ...
