from itertools import repeat

import equinox as eqx
import jax
from beartype.typing import Any, Callable, Sequence
from equinox.nn import StatefulLayer
from jaxtyping import Array, Float, PRNGKeyArray

from jaxonlayers.layers.abstract import AbstractNorm, AbstractNormStateful


def make_ntuple(x: Any, n: int) -> tuple[Any, ...]:
    if isinstance(x, collections.abc.Iterable):  # pyright: ignore
        return tuple(x)
    return tuple(repeat(x, n))


class ConvNormActivation(StatefulLayer):
    conv: eqx.nn.Conv
    norm: AbstractNorm | AbstractNormStateful | None
    activation: eqx.nn.Lambda | None

    def __init__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        kernel_size: int | Sequence[int] = 3,
        stride: int | Sequence[int] = 1,
        padding: str | int | Sequence[int] | Sequence[tuple[int, int]] | None = None,
        groups: int = 1,
        activation_layer: Callable[..., Array] | None = jax.nn.relu,
        dilation: int | Sequence[int] = 1,
        use_bias: bool | None = None,
        *,
        norm_layer: Callable[..., AbstractNorm | AbstractNormStateful] | None,
        key: PRNGKeyArray,
        dtype: Any,
    ):
        if padding is None:
            if isinstance(kernel_size, int) and isinstance(dilation, int):
                padding = (kernel_size - 1) // 2 * dilation
            else:
                _conv_dim = (
                    len(kernel_size)
                    if isinstance(kernel_size, Sequence)
                    else len(dilation)  # pyright: ignore
                )
                kernel_size = make_ntuple(kernel_size, _conv_dim)
                dilation = make_ntuple(dilation, _conv_dim)
                padding = tuple(
                    (kernel_size[i] - 1) // 2 * dilation[i] for i in range(_conv_dim)
                )
        if use_bias is None:
            use_bias = norm_layer is None

        key, subkey = jax.random.split(key)

        self.conv = eqx.nn.Conv(
            num_spatial_dims=num_spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            use_bias=use_bias,
            dtype=dtype,
            key=subkey,
        )

        self.norm = None
        if norm_layer is not None:
            self.norm = norm_layer()

        if activation_layer is not None:
            self.activation = eqx.nn.Lambda(activation_layer)
        else:
            self.activation = None

    def __call__(
        self,
        x: Float[Array, "c *num_spatial_dims"],
        state: eqx.nn.State,
        key: PRNGKeyArray | None = None,
    ) -> tuple[Float[Array, "c_out *num_spatial_dims_out"], eqx.nn.State]:
        x = self.conv(x)

        if self.norm:
            if isinstance(self.norm, StatefulLayer) and self.norm.is_stateful():
                assert state is not None
                x, state = self.norm(x, state)  # pyright: ignore
            else:
                x = self.norm(x)  # pyright: ignore

        if self.activation:
            x = self.activation(x)

        return x, state
