"""
kiri/nn/recurrent.py

Embedding, RNN, LSTM — for NLP and sequence assignments.
"""

import numpy as np
from kiri.backend.detect import BACKEND
from kiri.nn.layers import Module
from kiri.tensor import _CPUTensor

if BACKEND == "mlx":
    import mlx.core as mx
    import mlx.nn as mlx_nn


# ─── Embedding ────────────────────────────────────────────────────────────────

class Embedding(Module):
    """
    Lookup table: maps integer token IDs → dense vectors.

    Parameters
    ----------
    num_embeddings : vocabulary size
    embedding_dim  : vector dimension

    Input  : (N, seq_len) integer array
    Output : (N, seq_len, embedding_dim) float array

    Example
    -------
    emb = kiri.nn.Embedding(10000, 128)
    tokens = np.array([[3, 17, 42, 0]])   # (1, 4) token IDs
    out = emb(tokens)                      # (1, 4, 128)
    """

    def __init__(self, num_embeddings, embedding_dim, padding_idx=None):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim  = embedding_dim
        self.padding_idx    = padding_idx

        scale = np.sqrt(2.0 / embedding_dim)
        if BACKEND == "mlx":
            self.weight = mx.random.normal((num_embeddings, embedding_dim)) * scale
        else:
            self.weight = _CPUTensor(
                (np.random.randn(num_embeddings, embedding_dim) * scale).astype(np.float32)
            )
            if padding_idx is not None:
                self.weight.data[padding_idx] = 0.0

    def forward(self, idx):
        """idx: np.ndarray of shape (N, seq_len) with integer token IDs."""
        if BACKEND == "mlx":
            return self.weight[mx.array(idx)]
        else:
            # idx can be numpy array or _CPUTensor
            if isinstance(idx, _CPUTensor):
                idx_data = idx.data.astype(int)
            else:
                idx_data = np.asarray(idx, dtype=int)

            out_data = self.weight.data[idx_data]   # (N, seq_len, embedding_dim)
            out = _CPUTensor(out_data, (self.weight,), "embedding")

            def _backward():
                if isinstance(idx, _CPUTensor):
                    i = idx.data.astype(int).ravel()
                else:
                    i = idx_data.ravel()
                np.add.at(self.weight.grad, i, out.grad.reshape(-1, self.embedding_dim))
                if self.padding_idx is not None:
                    self.weight.grad[self.padding_idx] = 0.0

            out._backward = _backward
            return out

    def parameters(self):
        return [self.weight]

    def __repr__(self):
        return f"Embedding({self.num_embeddings}, {self.embedding_dim})"


# ─── RNN ──────────────────────────────────────────────────────────────────────

class RNN(Module):
    """
    Elman RNN (single layer).

    h_t = tanh(x_t @ W_ih^T + b_ih + h_{t-1} @ W_hh^T + b_hh)

    Input  : (N, seq_len, input_size)
    Output : (N, seq_len, hidden_size), h_n: (N, hidden_size)
    """

    def __init__(self, input_size, hidden_size, bias=True):
        super().__init__()
        self.input_size  = input_size
        self.hidden_size = hidden_size

        k = np.sqrt(1.0 / hidden_size)
        if BACKEND == "mlx":
            self.W_ih = mx.random.uniform(-k, k, (hidden_size, input_size))
            self.W_hh = mx.random.uniform(-k, k, (hidden_size, hidden_size))
            self.b_ih = mx.zeros((hidden_size,)) if bias else None
            self.b_hh = mx.zeros((hidden_size,)) if bias else None
        else:
            self.W_ih = _CPUTensor(np.random.uniform(-k, k, (hidden_size, input_size)).astype(np.float32))
            self.W_hh = _CPUTensor(np.random.uniform(-k, k, (hidden_size, hidden_size)).astype(np.float32))
            self.b_ih = _CPUTensor(np.zeros(hidden_size, dtype=np.float32)) if bias else None
            self.b_hh = _CPUTensor(np.zeros(hidden_size, dtype=np.float32)) if bias else None

    def forward(self, x, h0=None):
        """
        x  : (N, seq_len, input_size)
        h0 : (N, hidden_size) initial hidden state, zeros if None
        """
        if BACKEND == "mlx":
            return self._forward_mlx(x, h0)
        return self._forward_cpu(x, h0)

    def _forward_cpu(self, x, h0):
        if isinstance(x, _CPUTensor):
            x_data = x.data
        else:
            x_data = np.asarray(x, dtype=np.float32)

        N, T, _ = x_data.shape
        h = np.zeros((N, self.hidden_size), dtype=np.float32) if h0 is None else (
            h0.data if isinstance(h0, _CPUTensor) else np.asarray(h0, dtype=np.float32)
        )

        outputs = []
        hs      = [h]

        for t in range(T):
            xt = _CPUTensor(x_data[:, t, :], requires_grad=False)
            h_t = _CPUTensor(h, requires_grad=False)
            pre = xt.matmul(self.W_ih.T) + h_t.matmul(self.W_hh.T)
            if self.b_ih is not None:
                pre = pre + self.b_ih + self.b_hh
            h = np.tanh(pre.data)
            outputs.append(h[:, np.newaxis, :])
            hs.append(h)

        all_out = np.concatenate(outputs, axis=1)
        return _CPUTensor(all_out, requires_grad=False), _CPUTensor(h, requires_grad=False)

    def _forward_mlx(self, x, h0):
        if isinstance(x, np.ndarray):
            x = mx.array(x)
        N, T, _ = x.shape
        h = mx.zeros((N, self.hidden_size)) if h0 is None else h0

        outputs = []
        for t in range(T):
            xt  = x[:, t, :]
            pre = xt @ self.W_ih.T + h @ self.W_hh.T
            if self.b_ih is not None:
                pre = pre + self.b_ih + self.b_hh
            h = mx.tanh(pre)
            outputs.append(h[:, None, :])

        return mx.concatenate(outputs, axis=1), h

    def parameters(self):
        p = [self.W_ih, self.W_hh]
        if self.b_ih is not None:
            p += [self.b_ih, self.b_hh]
        return p

    def __repr__(self):
        return f"RNN(input_size={self.input_size}, hidden_size={self.hidden_size})"


# ─── LSTM ─────────────────────────────────────────────────────────────────────

class LSTM(Module):
    """
    Single-layer LSTM.

    Input  : (N, seq_len, input_size)
    Output : (N, seq_len, hidden_size), (h_n, c_n)
    """

    def __init__(self, input_size, hidden_size, bias=True):
        super().__init__()
        self.input_size  = input_size
        self.hidden_size = hidden_size

        k = np.sqrt(1.0 / hidden_size)
        # 4 gates: input, forget, cell, output — packed together
        def W(shape): return np.random.uniform(-k, k, shape).astype(np.float32)

        if BACKEND == "mlx":
            self.W_ih = mx.array(W((4*hidden_size, input_size)))
            self.W_hh = mx.array(W((4*hidden_size, hidden_size)))
            self.b_ih = mx.zeros((4*hidden_size,)) if bias else None
            self.b_hh = mx.zeros((4*hidden_size,)) if bias else None
        else:
            self.W_ih = _CPUTensor(W((4*hidden_size, input_size)))
            self.W_hh = _CPUTensor(W((4*hidden_size, hidden_size)))
            self.b_ih = _CPUTensor(np.zeros(4*hidden_size, dtype=np.float32)) if bias else None
            self.b_hh = _CPUTensor(np.zeros(4*hidden_size, dtype=np.float32)) if bias else None

    def forward(self, x, hc0=None):
        if BACKEND == "mlx":
            return self._forward_mlx(x, hc0)
        return self._forward_cpu(x, hc0)

    def _forward_cpu(self, x, hc0):
        if isinstance(x, _CPUTensor):
            x_data = x.data
        else:
            x_data = np.asarray(x, dtype=np.float32)

        N, T, _ = x_data.shape
        H = self.hidden_size

        if hc0 is None:
            h = np.zeros((N, H), dtype=np.float32)
            c = np.zeros((N, H), dtype=np.float32)
        else:
            h, c = hc0
            h = h.data if isinstance(h, _CPUTensor) else np.asarray(h, dtype=np.float32)
            c = c.data if isinstance(c, _CPUTensor) else np.asarray(c, dtype=np.float32)

        outputs = []
        Wih = self.W_ih.data
        Whh = self.W_hh.data
        bih = self.b_ih.data if self.b_ih else 0
        bhh = self.b_hh.data if self.b_hh else 0

        for t in range(T):
            gates = x_data[:, t, :] @ Wih.T + h @ Whh.T + bih + bhh
            i_g = _sigmoid(gates[:, :H])
            f_g = _sigmoid(gates[:, H:2*H])
            g_g = np.tanh(gates[:, 2*H:3*H])
            o_g = _sigmoid(gates[:, 3*H:])
            c   = f_g * c + i_g * g_g
            h   = o_g * np.tanh(c)
            outputs.append(h[:, np.newaxis, :])

        all_out = np.concatenate(outputs, axis=1)
        return (
            _CPUTensor(all_out, requires_grad=False),
            (_CPUTensor(h, requires_grad=False), _CPUTensor(c, requires_grad=False))
        )

    def _forward_mlx(self, x, hc0):
        if isinstance(x, np.ndarray):
            x = mx.array(x)
        N, T, _ = x.shape
        H = self.hidden_size

        h = mx.zeros((N, H)) if hc0 is None else hc0[0]
        c = mx.zeros((N, H)) if hc0 is None else hc0[1]

        outputs = []
        for t in range(T):
            gates = x[:, t, :] @ self.W_ih.T + h @ self.W_hh.T
            if self.b_ih is not None:
                gates = gates + self.b_ih + self.b_hh
            i_g = mx.sigmoid(gates[:, :H])
            f_g = mx.sigmoid(gates[:, H:2*H])
            g_g = mx.tanh(gates[:, 2*H:3*H])
            o_g = mx.sigmoid(gates[:, 3*H:])
            c   = f_g * c + i_g * g_g
            h   = o_g * mx.tanh(c)
            outputs.append(h[:, None, :])

        return mx.concatenate(outputs, axis=1), (h, c)

    def parameters(self):
        p = [self.W_ih, self.W_hh]
        if self.b_ih is not None:
            p += [self.b_ih, self.b_hh]
        return p

    def __repr__(self):
        return f"LSTM(input_size={self.input_size}, hidden_size={self.hidden_size})"


def _sigmoid(x):
    return 1 / (1 + np.exp(-x.clip(-500, 500)))
