from typing import Literal

import lucid
import lucid.nn as nn
import lucid.nn.functional as F

from lucid._tensor import Tensor
from lucid.nn.utils.rnn import PackedSequence
from lucid.types import Numeric, _DeviceType

from .activation import Tanh, ReLU


__all__ = ["RNNCell", "LSTMCell", "GRUCell", "RNNBase", "RNN", "LSTM", "GRU"]


def _get_activation(nonlinearity: str) -> type[nn.Module]:
    if nonlinearity == "tanh":
        return Tanh
    elif nonlinearity == "relu":
        return ReLU
    else:
        raise ValueError(
            f"Invalid nonlinearity '{nonlinearity}'. "
            "Supported nonlinearities are 'tanh' and 'relu'."
        )


class RNNCell(nn.Module):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        bias: bool = True,
        nonlinearity: Literal["tanh", "relu"] = "tanh",
    ) -> None:
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.bias = bias
        self.nonlinearity = _get_activation(nonlinearity)()

        sqrt_k = 1.0 / (hidden_size**0.5)
        self.weight_ih = nn.Parameter(
            lucid.random.uniform(-sqrt_k, sqrt_k, (self.hidden_size, self.input_size))
        )
        self.weight_hh = nn.Parameter(
            lucid.random.uniform(-sqrt_k, sqrt_k, (self.hidden_size, self.hidden_size))
        )

        if self.bias:
            self.bias_ih = nn.Parameter(
                lucid.random.uniform(-sqrt_k, sqrt_k, self.hidden_size)
            )
            self.bias_hh = nn.Parameter(
                lucid.random.uniform(-sqrt_k, sqrt_k, self.hidden_size)
            )
        else:
            self.bias_ih = None
            self.bias_hh = None

    def forward(self, input_: Tensor, hx: Tensor | None = None) -> Tensor:
        if input_.ndim not in (1, 2):
            raise ValueError(
                "RNNCell expected input with 1 or 2 dimensions, "
                f"got {input_.ndim} dimensions"
            )

        is_batched = input_.ndim == 2
        if not is_batched:
            input_ = input_.unsqueeze(axis=0)
        batch_size = input_.shape[0]

        if hx is None:
            hx = lucid.zeros(
                batch_size, self.hidden_size, dtype=input_.dtype, device=input_.device
            )
        else:
            if hx.ndim not in (1, 2):
                raise ValueError(
                    "RNNCell expected hidden state with 1 or 2 dimensions, "
                    f"got {hx.ndim} dimensions"
                )
            if hx.ndim == 1:
                hx = hx.unsqueeze(axis=0)

            if hx.shape[0] != batch_size or hx.shape[1] != self.hidden_size:
                raise ValueError(
                    "RNNCell expected hidden state with shape "
                    f"({batch_size}, {self.hidden_size}), got {hx.shape}"
                )

        hy = F.linear(input_, self.weight_ih, self.bias_ih)
        hy += F.linear(hx, self.weight_hh, self.bias_hh)
        ret = self.nonlinearity(hy)

        if not is_batched:
            ret = ret.squeeze(axis=0)
        return ret


class LSTMCell(nn.Module):
    def __init__(
        self, input_size: int, hidden_size: int, bias: bool = True, **kwargs
    ) -> None:
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.bias = bias

        sqrt_k = 1.0 / (hidden_size**0.5)
        self.weight_ih = nn.Parameter(
            lucid.random.uniform(
                -sqrt_k, sqrt_k, (4 * self.hidden_size, self.input_size)
            )
        )
        self.weight_hh = nn.Parameter(
            lucid.random.uniform(
                -sqrt_k, sqrt_k, (4 * self.hidden_size, self.hidden_size)
            )
        )

        if self.bias:
            self.bias_ih = nn.Parameter(
                lucid.random.uniform(-sqrt_k, sqrt_k, 4 * self.hidden_size)
            )
            self.bias_hh = nn.Parameter(
                lucid.random.uniform(-sqrt_k, sqrt_k, 4 * self.hidden_size)
            )
        else:
            self.bias_ih = None
            self.bias_hh = None

    def forward(
        self, input_: Tensor, hx: tuple[Tensor, Tensor] | None = None
    ) -> tuple[Tensor, Tensor]:
        if input_.ndim not in (1, 2):
            raise ValueError(
                "LSTMCell expected input with 1 or 2 dimensions, "
                f"got {input_.ndim} dimensions"
            )

        is_batched = input_.ndim == 2
        if not is_batched:
            input_ = input_.unsqueeze(axis=0)
        batch_size = input_.shape[0]

        if hx is None:
            h_t = lucid.zeros(
                batch_size, self.hidden_size, dtype=input_.dtype, device=input_.device
            )
            c_t = lucid.zeros(
                batch_size, self.hidden_size, dtype=input_.dtype, device=input_.device
            )
        else:
            h_t, c_t = hx
            if h_t.ndim not in (1, 2) or c_t.ndim not in (1, 2):
                raise ValueError(
                    "LSTMCell expected hidden state and cell state with 1 or 2 dimensions"
                )

            if h_t.ndim == 1:
                h_t = h_t.unsqueeze(axis=0)
            if c_t.ndim == 1:
                c_t = c_t.unsqueeze(axis=0)

            if h_t.shape[0] != batch_size or h_t.shape[1] != self.hidden_size:
                raise ValueError(
                    "LSTMCell expected hidden state with shape "
                    f"({batch_size}, {self.hidden_size}), got {h_t.shape}"
                )
            if c_t.shape[0] != batch_size or c_t.shape[1] != self.hidden_size:
                raise ValueError(
                    "LSTMCell expected cell state with shape "
                    f"({batch_size}, {self.hidden_size}), got {c_t.shape}"
                )

        gates = F.linear(input_, self.weight_ih, self.bias_ih)
        gates += F.linear(h_t, self.weight_hh, self.bias_hh)

        i_t, f_t, g_t, o_t = lucid.split(gates, 4, axis=1)
        i_t = F.sigmoid(i_t)
        f_t = F.sigmoid(f_t)
        g_t = F.tanh(g_t)
        o_t = F.sigmoid(o_t)

        c_t = f_t * c_t + i_t * g_t
        h_t = o_t * F.tanh(c_t)

        if not is_batched:
            h_t = h_t.squeeze(axis=0)
            c_t = c_t.squeeze(axis=0)
        return h_t, c_t


class GRUCell(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, bias: bool = True) -> None:
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.bias = bias

        sqrt_k = 1.0 / (hidden_size**0.5)
        self.weight_ih = nn.Parameter(
            lucid.random.uniform(
                -sqrt_k, sqrt_k, (3 * self.hidden_size, self.input_size)
            )
        )
        self.weight_hh = nn.Parameter(
            lucid.random.uniform(
                -sqrt_k, sqrt_k, (3 * self.hidden_size, self.hidden_size)
            )
        )

        if self.bias:
            self.bias_ih = nn.Parameter(
                lucid.random.uniform(-sqrt_k, sqrt_k, 3 * self.hidden_size)
            )
            self.bias_hh = nn.Parameter(
                lucid.random.uniform(-sqrt_k, sqrt_k, 3 * self.hidden_size)
            )
        else:
            self.bias_ih = None
            self.bias_hh = None

    def forward(self, input_: Tensor, hx: Tensor | None = None) -> Tensor:
        if input_.ndim not in (1, 2):
            raise ValueError(
                "GRUCell expected input with 1 or 2 dimensions, "
                f"got {input_.ndim} dimensions"
            )

        is_batched = input_.ndim == 2
        if not is_batched:
            input_ = input_.unsqueeze(axis=0)
        batch_size = input_.shape[0]

        if hx is None:
            hx = lucid.zeros(
                batch_size, self.hidden_size, dtype=input_.dtype, device=input_.device
            )
        else:
            if hx.ndim not in (1, 2):
                raise ValueError(
                    "GRUCell expected hidden state with 1 or 2 dimensions, "
                    f"got {hx.ndim} dimensions"
                )

            if hx.ndim == 1:
                hx = hx.unsqueeze(axis=0)
            if hx.shape[0] != batch_size or hx.shape[1] != self.hidden_size:
                raise ValueError(
                    "GRUCell expected hidden state with shape "
                    f"({batch_size}, {self.hidden_size}), got {hx.shape}"
                )

        input_gates = F.linear(input_, self.weight_ih, self.bias_ih)
        hidden_gates = F.linear(hx, self.weight_hh, self.bias_hh)

        i_r, i_z, i_n = lucid.split(input_gates, 3, axis=1)
        h_r, h_z, h_n = lucid.split(hidden_gates, 3, axis=1)

        r_t = F.sigmoid(i_r + h_r)
        z_t = F.sigmoid(i_z + h_z)
        n_t = F.tanh(i_n + r_t * h_n)

        h_t = (1 - z_t) * n_t + z_t * hx

        if not is_batched:
            h_t = h_t.squeeze(axis=0)
        return h_t


class RNNBase(nn.Module):
    def __init__(
        self,
        mode: Literal["RNN_TANH", "RNN_RELU", "LSTM", "GRU"],
        input_size: int,
        hidden_size: int,
        num_layers: int = 1,
        bias: bool = True,
        batch_first: bool = False,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.is_lstm = False
        nonlinearity = "tanh"

        if mode == "RNN_TANH":
            pass
        elif mode == "RNN_RELU":
            nonlinearity = "relu"
        elif mode == "LSTM":
            self.is_lstm = True
        elif mode == "GRU":
            pass
        else:
            raise ValueError(
                f"Invalid mode '{mode}'. Supported modes are 'RNN_TANH', "
                "'RNN_RELU', 'LSTM', or 'GRU'."
            )

        self.mode = mode
        self.nonlinearity = nonlinearity

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bias = bias
        self.batch_first = batch_first
        self.dropout = float(dropout)

        for layer in range(num_layers):
            layer_input_size = input_size if layer == 0 else hidden_size
            sqrt_k = 1.0 / (hidden_size**0.5)

            if mode in ("RNN_TANH", "RNN_RELU"):
                w_ih = nn.Parameter(
                    lucid.random.uniform(
                        -sqrt_k, sqrt_k, (hidden_size, layer_input_size)
                    )
                )
                w_hh = nn.Parameter(
                    lucid.random.uniform(-sqrt_k, sqrt_k, (hidden_size, hidden_size))
                )
            elif mode == "LSTM":
                w_ih = nn.Parameter(
                    lucid.random.uniform(
                        -sqrt_k, sqrt_k, (4 * hidden_size, layer_input_size)
                    )
                )
                w_hh = nn.Parameter(
                    lucid.random.uniform(
                        -sqrt_k, sqrt_k, (4 * hidden_size, hidden_size)
                    )
                )
            else:
                w_ih = nn.Parameter(
                    lucid.random.uniform(
                        -sqrt_k, sqrt_k, (3 * hidden_size, layer_input_size)
                    )
                )
                w_hh = nn.Parameter(
                    lucid.random.uniform(
                        -sqrt_k, sqrt_k, (3 * hidden_size, hidden_size)
                    )
                )

            self.register_parameter(f"weight_ih_l{layer}", w_ih)
            self.register_parameter(f"weight_hh_l{layer}", w_hh)

            if bias:
                b_ih = nn.Parameter(
                    lucid.random.uniform(-sqrt_k, sqrt_k, w_ih.shape[0])
                )
                b_hh = nn.Parameter(
                    lucid.random.uniform(-sqrt_k, sqrt_k, w_hh.shape[0])
                )
                self.register_parameter(f"bias_ih_l{layer}", b_ih)
                self.register_parameter(f"bias_hh_l{layer}", b_hh)
            else:
                self.register_parameter(f"bias_ih_l{layer}", None)
                self.register_parameter(f"bias_hh_l{layer}", None)

    def _rnn_cell(
        self,
        input_: Tensor,
        hx: Tensor,
        w_ih: Tensor,
        w_hh: Tensor,
        b_ih: Tensor | None,
        b_hh: Tensor | None,
    ) -> Tensor:
        hy = F.linear(input_, w_ih, b_ih)
        hy += F.linear(hx, w_hh, b_hh)

        if self.mode == "RNN_TANH":
            return F.tanh(hy)

        return F.relu(hy)

    def _lstm_cell(
        self,
        input_: Tensor,
        hx: Tensor,
        cx: Tensor,
        w_ih: Tensor,
        w_hh: Tensor,
        b_ih: Tensor | None,
        b_hh: Tensor | None,
    ) -> tuple[Tensor, Tensor]:
        gates = F.linear(input_, w_ih, b_ih)
        gates += F.linear(hx, w_hh, b_hh)

        i_t, f_t, g_t, o_t = lucid.split(gates, 4, axis=1)
        i_t = F.sigmoid(i_t)
        f_t = F.sigmoid(f_t)
        g_t = F.tanh(g_t)
        o_t = F.sigmoid(o_t)

        c_t = f_t * cx + i_t * g_t
        h_t = o_t * F.tanh(c_t)

        return h_t, c_t

    def _gru_cell(
        self,
        input_: Tensor,
        hx: Tensor,
        w_ih: Tensor,
        w_hh: Tensor,
        b_ih: Tensor | None,
        b_hh: Tensor | None,
    ) -> Tensor:
        input_gates = F.linear(input_, w_ih, b_ih)
        hidden_gates = F.linear(hx, w_hh, b_hh)

        i_r, i_z, i_n = lucid.split(input_gates, 3, axis=1)
        h_r, h_z, h_n = lucid.split(hidden_gates, 3, axis=1)

        r_t = F.sigmoid(i_r + h_r)
        z_t = F.sigmoid(i_z + h_z)
        n_t = F.tanh(i_n + r_t * h_n)
        h_t = (1 - z_t) * n_t + z_t * hx

        return h_t

    def _init_hidden(
        self, batch_size: int, dtype: Numeric, device: _DeviceType
    ) -> Tensor | tuple[Tensor, Tensor]:
        if self.is_lstm:
            h0 = lucid.zeros(
                self.num_layers,
                batch_size,
                self.hidden_size,
                dtype=dtype,
                device=device,
            )
            c0 = lucid.zeros(
                self.num_layers,
                batch_size,
                self.hidden_size,
                dtype=dtype,
                device=device,
            )
            return h0, c0
        return lucid.zeros(
            self.num_layers, batch_size, self.hidden_size, dtype=dtype, device=device
        )

    def forward(
        self,
        input_: Tensor | PackedSequence,
        hx: Tensor | tuple[Tensor, Tensor] | None = None,
    ) -> (
        tuple[Tensor | PackedSequence, Tensor]
        | tuple[Tensor | PackedSequence, tuple[Tensor, Tensor]]
    ):
        is_packed = isinstance(input_, PackedSequence)
        if is_packed:
            data = input_.data
            batch_sizes = input_.batch_sizes
            if data.ndim != 2:
                raise ValueError(
                    "RNNBase expected packed data with 2 dimensions, "
                    f"got {data.ndim} dimensions"
                )
            if batch_sizes.ndim != 1 or batch_sizes.shape[0] == 0:
                raise ValueError(
                    "PackedSequence batch_sizes must be a non-empty 1D tensor"
                )

            batch_size = int(batch_sizes[0].item())
            feat = data.shape[1]
            if feat != self.input_size:
                raise ValueError(
                    f"RNNBase expected input with feature size {self.input_size}, got {feat}"
                )
        else:
            if input_.ndim != 3:
                raise ValueError(
                    f"RNNBase expected input with 3 dimensions, got {input_.ndim} dimensions"
                )

            if self.batch_first:
                input_ = input_.swapaxes(0, 1)

            seq_len, batch_size, feat = input_.shape
            if feat != self.input_size:
                raise ValueError(
                    f"RNNBase expected input with feature size {self.input_size}, got {feat}"
                )

        if self.is_lstm:
            if hx is None:
                hx = self._init_hidden(batch_size, input_.dtype, input_.device)
            if not (
                isinstance(hx, (tuple, list))
                and len(hx) == 2
                and isinstance(hx[0], Tensor)
                and isinstance(hx[1], Tensor)
            ):
                raise ValueError("LSTM expects hx as a tuple of (h_0, c_0)")

            h0, c0 = hx
            if h0.ndim == 2:
                h0 = h0.unsqueeze(axis=0)
            if c0.ndim == 2:
                c0 = c0.unsqueeze(axis=0)

            if h0.ndim != 3 or c0.ndim != 3:
                raise ValueError("LSTM expects h_0 and c_0 with 3 dimensions")
            if h0.shape[0] != self.num_layers or c0.shape[0] != self.num_layers:
                raise ValueError("Incorrect number of layers in h_0 or c_0")
            if h0.shape[1] != batch_size or c0.shape[1] != batch_size:
                raise ValueError("Incorrect batch size in h_0 or c_0")
            if h0.shape[2] != self.hidden_size or c0.shape[2] != self.hidden_size:
                raise ValueError("Incorrect hidden size in h_0 or c_0")

            hx_h, hx_c = h0, c0

        else:
            if hx is None:
                hx = self._init_hidden(batch_size, input_.dtype, input_.device)
            if hx.ndim == 2:
                hx = hx.unsqueeze(axis=0)
            if hx.ndim != 3:
                raise ValueError(
                    f"RNNBase expected hidden state with 3 dimensions, got {hx.ndim} dimensions"
                )

            if hx.shape[0] != self.num_layers or hx.shape[1] != batch_size:
                raise ValueError("hx has incorrect shape")
            if hx.shape[2] != self.hidden_size:
                raise ValueError("Incorrect hidden size in hx")

        layer_input = data if is_packed else input_
        h_n_list: list[Tensor] = []
        c_n_list: list[Tensor] | None = [] if self.is_lstm else None

        for layer_idx in range(self.num_layers):
            w_ih = getattr(self, f"weight_ih_l{layer_idx}")
            w_hh = getattr(self, f"weight_hh_l{layer_idx}")
            b_ih = getattr(self, f"bias_ih_l{layer_idx}")
            b_hh = getattr(self, f"bias_hh_l{layer_idx}")

            if self.is_lstm:
                h_t = hx_h[layer_idx]
                c_t = hx_c[layer_idx]
            else:
                h_t = hx[layer_idx]

            outputs = []
            if is_packed:
                final_h: list[Tensor] = []
                final_c: list[Tensor] | None = [] if self.is_lstm else None
                offset = 0

                prev_bs: int | None = None
                max_len = int(batch_sizes.shape[0])
                for t in range(max_len):
                    bs = int(batch_sizes[t].item())
                    if bs == 0:
                        break

                    if prev_bs is None:
                        prev_bs = bs
                    if bs > prev_bs:
                        raise ValueError(
                            "PackedSequence batch_sizes must be non-increasing"
                        )

                    if bs < prev_bs:
                        final_h.append(h_t[bs:prev_bs])
                        if self.is_lstm and final_c is not None:
                            final_c.append(c_t[bs:prev_bs])

                        h_t = h_t[:bs]
                        if self.is_lstm:
                            c_t = c_t[:bs]

                    step_input = layer_input[offset : offset + bs]
                    offset += bs

                    if self.is_lstm:
                        h_t, c_t = self._lstm_cell(
                            step_input, h_t, c_t, w_ih, w_hh, b_ih, b_hh
                        )
                    elif self.mode == "GRU":
                        h_t = self._gru_cell(step_input, h_t, w_ih, w_hh, b_ih, b_hh)
                    else:
                        h_t = self._rnn_cell(step_input, h_t, w_ih, w_hh, b_ih, b_hh)

                    outputs.append(h_t)
                    prev_bs = bs

                final_h.append(h_t)
                if self.is_lstm and final_c is not None:
                    final_c.append(c_t)

                h_n_list.append(
                    lucid.concatenate(tuple(reversed(final_h)), axis=0).unsqueeze(
                        axis=0
                    )
                )
                if self.is_lstm and final_c is not None and c_n_list is not None:
                    c_n_list.append(
                        lucid.concatenate(tuple(reversed(final_c)), axis=0).unsqueeze(
                            axis=0
                        )
                    )

                layer_output = (
                    lucid.concatenate(tuple(outputs), axis=0)
                    if outputs
                    else layer_input[:0]
                )

            else:
                for t in range(seq_len):
                    step_input = layer_input[t]
                    if self.is_lstm:
                        h_t, c_t = self._lstm_cell(
                            step_input, h_t, c_t, w_ih, w_hh, b_ih, b_hh
                        )
                        outputs.append(h_t.unsqueeze(axis=0))
                    elif self.mode == "GRU":
                        h_t = self._gru_cell(step_input, h_t, w_ih, w_hh, b_ih, b_hh)
                        outputs.append(h_t.unsqueeze(axis=0))
                    else:
                        h_t = self._rnn_cell(step_input, h_t, w_ih, w_hh, b_ih, b_hh)
                        outputs.append(h_t.unsqueeze(axis=0))

                layer_output = lucid.concatenate(tuple(outputs), axis=0)

            if self.training and self.dropout > 0.0 and layer_idx < self.num_layers - 1:
                layer_output = F.dropout(layer_output, p=self.dropout)

            if not is_packed:
                h_n_list.append(h_t.unsqueeze(axis=0))
                if self.is_lstm and c_n_list is not None:
                    c_n_list.append(c_t.unsqueeze(axis=0))
            layer_input = layer_output

        if is_packed:
            output = PackedSequence(
                data=layer_input,
                batch_sizes=batch_sizes,
                sorted_indices=input_.sorted_indices,
                unsorted_indices=input_.unsorted_indices,
            )
        else:
            output = layer_input

        h_n = lucid.concatenate(tuple(h_n_list), axis=0)
        if self.is_lstm and c_n_list is not None:
            c_n = lucid.concatenate(tuple(c_n_list), axis=0)

        if is_packed:
            if input_.unsorted_indices is not None:
                h_n = h_n[:, input_.unsorted_indices]
                if self.is_lstm and c_n_list is not None:
                    c_n = c_n[:, input_.unsorted_indices]
        else:
            if self.batch_first:
                output = output.swapaxes(0, 1)

        if self.is_lstm and c_n_list is not None:
            return output, (h_n, c_n)
        return output, h_n


class RNN(RNNBase):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int = 1,
        nonlinearity: Literal["tanh", "relu"] = "tanh",
        bias: bool = True,
        batch_first: bool = False,
        dropout: float = 0.0,
    ) -> None:
        if nonlinearity == "tanh":
            mode = "RNN_TANH"
        elif nonlinearity == "relu":
            mode = "RNN_RELU"
        else:
            raise ValueError(
                f"Invalid nonlinearity '{nonlinearity}'. "
                "Supported nonlinearities are 'tanh' and 'relu'."
            )

        super().__init__(
            mode=mode,
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            bias=bias,
            batch_first=batch_first,
            dropout=dropout,
        )


class LSTM(RNNBase):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int = 1,
        bias: bool = True,
        batch_first: bool = False,
        dropout: float = 0.0,
    ) -> None:
        mode = "LSTM"
        super().__init__(
            mode=mode,
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            bias=bias,
            batch_first=batch_first,
            dropout=dropout,
        )


class GRU(RNNBase):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int = 1,
        bias: bool = True,
        batch_first: bool = False,
        dropout: float = 0.0,
    ) -> None:
        mode = "GRU"
        super().__init__(
            mode=mode,
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            bias=bias,
            batch_first=batch_first,
            dropout=dropout,
        )
