from __future__ import annotations

from typing import Tuple

import torch
from torch import nn


def build_and_run_rnn_cell(a: int, b: int, c: int) -> torch.Tensor:
    """Single-step RNNCell."""
    cell = nn.RNNCell(input_size=c, hidden_size=a, nonlinearity="tanh")
    x1 = torch.randn(b, c)
    h0 = torch.zeros(b, a)
    h1 = cell(x1, h0)
    print("[RNNCell] h1:", tuple(h1.shape))
    return h1


def build_and_run_lstm_cell(a: int, b: int, c: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """Single-step LSTMCell."""
    cell = nn.LSTMCell(input_size=c, hidden_size=a)
    x1 = torch.randn(b, c)
    h0 = torch.zeros(b, a)
    c0 = torch.zeros(b, a)
    h1, c1 = cell(x1, (h0, c0))
    print("[LSTMCell] h1:", tuple(h1.shape), "c1:", tuple(c1.shape))
    return h1, c1


def MultiRNNCell_dynamic_call(a: int, b: int, c: int, d: int, e: int, cell_type: str = "rnn") -> torch.Tensor:
    """Multi-layer RNN/LSTM (batch_first)."""
    if cell_type.lower() == "lstm":
        rnn = nn.LSTM(input_size=e, hidden_size=b, num_layers=a, batch_first=True)
    else:
        rnn = nn.RNN(input_size=e, hidden_size=b, num_layers=a, nonlinearity="tanh", batch_first=True)
    inputs = torch.randn(c, d, e)
    if isinstance(rnn, nn.LSTM):
        h0 = torch.zeros(a, c, b); c0 = torch.zeros(a, c, b)
        outputs, (hn, cn) = rnn(inputs, (h0, c0))
        print("[Multi-LSTM] outputs:", tuple(outputs.shape))
    else:
        h0 = torch.zeros(a, c, b)
        outputs, hn = rnn(inputs, h0)
        print("[Multi-RNN] outputs:", tuple(outputs.shape))
    return outputs


if __name__ == "__main__":
    _ = build_and_run_rnn_cell(a=5, b=4, c=3)
    _ = build_and_run_lstm_cell(a=5, b=4, c=3)
    _ = MultiRNNCell_dynamic_call(a=2, b=8, c=4, d=10, e=6, cell_type="rnn")