from energizer.function import Function
from energizer.neural_network import Module, Parameter
from energizer.tensor import Tensor
import numpy as np
import mlx.core as mx
import energizer.functionnal as F

class ConvNd(Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
        padding: int = 0,
        bias: bool = True,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        if self.device == "gpu":
            std = mx.sqrt(2.0 / (in_channels * kernel_size))
            self.weight = Parameter(
                mx.random.normal(shape=(out_channels, in_channels, kernel_size)),
                requires_grad=True,
                device=self.device,
            )
        else:
            std = np.sqrt(2.0 / (in_channels * kernel_size))
            self.weight = Parameter(
                np.random.randn(out_channels, in_channels, kernel_size) * std,
                requires_grad=True,
                device=self.device,
            )
        if bias:
            if self.device == "gpu":
                self.bias = Parameter(
                    mx.zeros(out_channels), requires_grad=True, device=self.device
                )
            else:
                self.bias = Parameter(
                    np.zeros(out_channels), requires_grad=True, device=self.device
                )
        else:
            self.bias = None
        self._parameters["weight"] = self.weight
        if self.bias is not None:
            self._parameters["bias"] = self.bias

    def forward(self) -> Tensor:
        pass


class Conv1d(ConvNd):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
        padding: int = 0,
        bias: bool = True,
    ):
        super().__init__(in_channels, out_channels, kernel_size, stride, padding, bias)

    def forward(self, input: Tensor) -> Tensor:
        self._input = input
        self._input_padded = (
            self._pad1d(input, self.padding) if self.padding > 0 else input
        )

        batch_size, in_channels, length = self._input_padded.data.shape
        out_length = ((length - self.kernel_size) // self.stride) + 1

        x_2d = self._input_padded.data.reshape(-1, length)

        windows = []
        for i in range(x_2d.shape[0]):
            row_windows = F.as_strided(
                Tensor(x_2d[i]),
                shape=(out_length, self.kernel_size),
                strides=(self.stride, 1),
            ).data
            windows.append(row_windows)

        windows_stacked = np.stack(windows, axis=0)

        windows_reshaped = (
            windows_stacked.reshape(
                batch_size, in_channels, out_length, self.kernel_size
            )
            .transpose(0, 2, 1, 3)
            .reshape(batch_size, out_length, -1)
        )

        weight_reshaped = self.weight.data.reshape(self.out_channels, -1)

        result = np.einsum("bij,oj->bio", windows_reshaped, weight_reshaped)
        result = result.transpose(0, 2, 1)

        if self.bias is not None:
            result += self.bias.data.reshape(1, -1, 1)

        return Tensor(
            result,
            requires_grad=input.requires_grad,
        )

    def _pad1d(self, input: Tensor, padding: int) -> Tensor:
        batch_size, channels, length = input.data.shape
        padded_length = length + 2 * padding
        padded_data = np.zeros((batch_size, channels, padded_length))
        padded_data[:, :, padding : padding + length] = input.data
        return Tensor(padded_data, requires_grad=input.requires_grad)

    def _unfold1d(
        self, input: Tensor, kernel_size: int, stride: int, out_length: int
    ) -> Tensor:
        batch_size, channels, length = input.data.shape
        unfolded_data = np.zeros((batch_size, out_length, channels, kernel_size))
        for b in range(batch_size):
            for c in range(channels):
                channel_data = input.data[b, c, :]

                window = F.as_strided(
                    Tensor(channel_data),
                    shape=(out_length, kernel_size),
                    strides=(stride, 1),
                )

                unfolded_data[b, :, c, :] = window.data

        return Tensor(unfolded_data, requires_grad=input.requires_grad)

    def to(self, device: str):
        self.weight = self.weight.to(device)
        if self.bias is not None:
            self.bias = self.bias.to(device)
        return self


class Conv2d(Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: tuple,
        stride: tuple = (1, 1),
        padding: tuple = (0, 0),
        bias: bool = True,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = (
            kernel_size
            if isinstance(kernel_size, tuple)
            else (kernel_size, kernel_size)
        )
        self.stride = stride if isinstance(stride, tuple) else (stride, stride)
        self.padding = padding if isinstance(padding, tuple) else (padding, padding)

        if self.device == "gpu":
            std = mx.sqrt(
                2.0 / (in_channels * self.kernel_size[0] * self.kernel_size[1])
            )
            self.weight = Parameter(
                mx.random.normal(
                    shape=(
                        out_channels,
                        in_channels,
                        self.kernel_size[0],
                        self.kernel_size[1],
                    )
                ),
                requires_grad=True,
                device=self.device,
            )
        else:
            std = np.sqrt(
                2.0 / (in_channels * self.kernel_size[0] * self.kernel_size[1])
            )
            self.weight = Parameter(
                np.random.randn(
                    out_channels, in_channels, self.kernel_size[0], self.kernel_size[1]
                )
                * std,
                requires_grad=True,
                device=self.device,
            )
        if bias:
            if self.device == "gpu":
                self.bias = Parameter(
                    mx.zeros(out_channels), requires_grad=True, device=self.device
                )
            else:
                self.bias = Parameter(
                    np.zeros(out_channels), requires_grad=True, device=self.device
                )
        else:
            self.bias = None

    def forward(self, input: Tensor) -> Tensor:
        if self.padding != (0, 0):
            input = self._pad2d(input, self.padding)

        batch_size, in_channels, height, width = input.data.shape
        kh, kw = self.kernel_size
        sh, sw = self.stride

        out_height = ((height - kh) // sh) + 1
        out_width = ((width - kw) // sw) + 1

        windows = self._im2col(input, kh, kw, sh, sw, out_height, out_width)

        if input.device == "gpu":
            weight_reshaped = input.mlx().__class__(
                self.weight.data.reshape(self.out_channels, -1)
            )

            result = mx.zeros((batch_size, out_height * out_width, self.out_channels))
            for b in range(batch_size):
                result[b] = windows[b] @ weight_reshaped.T
            result = result.transpose(0, 2, 1).reshape(
                batch_size, self.out_channels, out_height, out_width
            )
            if self.bias is not None:
                result += self.bias.data.reshape(1, -1, 1, 1)
        else:
            weight_reshaped = self.weight.data.reshape(self.out_channels, -1)

            result = np.zeros((batch_size, out_height * out_width, self.out_channels))
            for b in range(batch_size):
                result[b] = windows[b] @ weight_reshaped.T

            result = result.transpose(0, 2, 1).reshape(
                batch_size, self.out_channels, out_height, out_width
            )

            if self.bias is not None:
                result += self.bias.data.reshape(1, -1, 1, 1)

        return Tensor(result, requires_grad=input.requires_grad, device=input.device)

    def _pad2d(self, x: Tensor, padding: tuple) -> Tensor:
        ph, pw = padding
        batch_size, channels, height, width = x.data.shape
        if ph == 0 and pw == 0:
            return x
        padded_data = np.zeros((batch_size, channels, height + 2 * ph, width + 2 * pw))
        padded_data[:, :, ph : ph + height, pw : pw + width] = x.data
        return Tensor(padded_data, requires_grad=x.requires_grad, device=x.device)

    def _im2col(
        self,
        x: Tensor,
        kh: int,
        kw: int,
        sh: int,
        sw: int,
        out_height: int,
        out_width: int,
    ):
        batch_size, channels, height, width = x.data.shape
        if x.device == "gpu":
            unfolded_data = mx.zeros(
                (batch_size, out_height * out_width, channels * kh * kw)
            )
        else:
            unfolded_data = np.zeros(
                (batch_size, out_height * out_width, channels * kh * kw)
            )

        for b in range(batch_size):
            col_index = 0
            for i in range(out_height):
                for j in range(out_width):
                    h_start = i * sh
                    w_start = j * sw
                    window = x.data[
                        b, :, h_start : h_start + kh, w_start : w_start + kw
                    ]
                    if x.device == "gpu":
                        window_flat = window.reshape(-1)
                    else:
                        window_flat = window.ravel()
                    unfolded_data[b, col_index] = window_flat
                    col_index += 1

        return unfolded_data
