from typing import Union
from energizer.tensor import Tensor
import numpy as np
import numpy.lib.stride_tricks as np_stride_tricks
import mlx.core as mx

from energizer.function import Function
import energizer.derivatives as dv


def max(tensor: Tensor, floor: Union[int, float] = 0) -> Tensor:
    if tensor.device == "gpu":
        return Tensor(
            mx.maximum(floor, tensor.data),
            requires_grad=tensor.requires_grad,
            grad_fn=Function(dv.max_backward, [tensor, floor]),
            device=tensor.device,
        )
    return Tensor(
        np.maximum(floor, tensor.data),
        requires_grad=tensor.requires_grad,
        grad_fn=Function(dv.max_backward, [tensor, floor]),
        device=tensor.device,
    )


def as_strided(
    tensor: Tensor, shape: tuple, strides: tuple, storage_offset: int = 0
) -> "Tensor":
    array = tensor.data

    itemsize = array.itemsize
    byte_strides = tuple(s * itemsize for s in strides)

    if storage_offset > 0:
        flat_array = array.flatten()
        if storage_offset >= len(flat_array):
            raise ValueError(
                f"storage_offset {storage_offset} is out of bounds for array of size {len(flat_array)}"
            )
        offset_array = flat_array[storage_offset:]
        array = offset_array

    if tensor.device == "gpu":
        strided_data = mx.as_strided(
            array, shape=shape, strides=byte_strides, storage_offset=storage_offset
        )
    else:
        strided_data = np_stride_tricks.as_strided(
            array, shape=shape, strides=byte_strides
        )

    return Tensor(
        strided_data,
        requires_grad=tensor.requires_grad,
        grad_fn=Function(
            dv.as_strided_backward, [tensor, shape, strides, storage_offset]
        ),
        device=tensor.device,
    )


def trace(tensor: Tensor) -> Tensor:
    if tensor.device == "gpu":
        result = mx.trace(tensor.data)
    else:
        result = np.trace(tensor.data)
    return Tensor(
        result,
        requires_grad=tensor.requires_grad,
        grad_fn=Function(dv.trace_backward, [tensor]),
        device=tensor.device,
    )


def tanh(tensor: Tensor) -> Tensor:
    if tensor.device == "gpu":
        return Tensor(
            mx.tanh(tensor.data),
            requires_grad=tensor.requires_grad,
            grad_fn=Function(dv.tanh_backward, [tensor]),
            device=tensor.device,
        )
    return Tensor(
        np.tanh(tensor.data),
        requires_grad=tensor.requires_grad,
        grad_fn=Function(dv.tanh_backward, [tensor]),
        device=tensor.device,
    )


def softmax(tensor: Tensor) -> Tensor:
    if tensor.device == "gpu":
        return Tensor(
            mx.softmax(tensor.data),
            requires_grad=tensor.requires_grad,
            grad_fn=Function(dv.softmax_backward, [tensor]),
            device=tensor.device,
        )
    return Tensor(
        np.softmax(tensor.data),
        requires_grad=tensor.requires_grad,
        grad_fn=Function(dv.softmax_backward, [tensor]),
        device=tensor.device,
    )
