import torch

# Text: sequence of word indices, variable length
review = [42, 17, 891, 0, 4, 230, ...]  # "This movie was great..."
# Shape: (sequence_length,)

# Time series: sensor readings over time
stock_prices = torch.randn(252, 5)   # 252 trading days, 5 features
# Shape: (time_steps, features)

# Audio: samples in a waveform
audio = torch.randn(16000)           # 1 second at 16kHz
# Shape: (num_samples,)

# Video: frames over time
video = torch.randn(30, 3, 224, 224) # 30 frames, RGB, 224x224
# Shape: (frames, channels, height, width)

# When batching: sequences get an extra batch dimension
# Shape convention: (batch_size, sequence_length, features) — with batch_first=True
# OR: (sequence_length, batch_size, features) — PyTorch default

import torch
import torch.nn as nn
import torch.nn.functional as F

# Manual single-step RNN to build intuition
class ManualRNNCell(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.Wx = nn.Linear(input_size, hidden_size)   # weights for input
        self.Wh = nn.Linear(hidden_size, hidden_size, bias=False)  # weights for hidden state

    def forward(self, x_t, h_prev):
        """
        x_t:   (batch, input_size)
        h_prev: (batch, hidden_size)
        returns h_t: (batch, hidden_size)
        """
        return torch.tanh(self.Wx(x_t) + self.Wh(h_prev))


# Process a sequence of 10 steps
input_size, hidden_size, batch_size = 20, 64, 4
seq_len = 10

cell = ManualRNNCell(input_size, hidden_size)

# Initial hidden state: all zeros
h = torch.zeros(batch_size, hidden_size)

# Process sequence step by step
sequence = torch.randn(seq_len, batch_size, input_size)  # (T, B, features)
hidden_states = []

for t in range(seq_len):
    x_t = sequence[t]   # (batch, input_size)
    h = cell(x_t, h)    # update hidden state
    hidden_states.append(h)
    print(f"Step {t+1:2d}: h range [{h.min():.3f}, {h.max():.3f}]")

# Final hidden state h contains a summary of the entire sequence
print(f"\nFinal hidden state shape: {h.shape}")

import torch
import torch.nn as nn

# nn.RNN handles the loop for you
rnn = nn.RNN(
    input_size=20,       # size of each input vector
    hidden_size=64,      # size of hidden state
    num_layers=1,        # number of stacked RNN layers
    batch_first=True,    # input shape: (batch, seq_len, features)
    nonlinearity='tanh'  # 'tanh' or 'relu'
)

# Input: batch of 4 sequences, each 10 steps long, 20 features per step
x = torch.randn(4, 10, 20)             # (batch, seq_len, input_size)
h0 = torch.zeros(1, 4, 64)            # (num_layers, batch, hidden_size)

output, h_n = rnn(x, h0)

print(f"output shape: {output.shape}")  # (4, 10, 64) — hidden state at every step
print(f"h_n shape:    {h_n.shape}")     # (1, 4, 64)  — final hidden state

import torch
import torch.nn as nn

# Demonstrate gradient flow through many time steps
rnn = nn.RNN(input_size=10, hidden_size=32, batch_first=True)

# Sequence of length 100 (a short paragraph)
x = torch.randn(1, 100, 10, requires_grad=True)
h0 = torch.zeros(1, 1, 32)

output, h_n = rnn(x, h0)

# Loss depends only on the FINAL hidden state — must propagate back 100 steps
loss = h_n.sum()
loss.backward()

# Examine gradient magnitude at early vs late positions
grads = x.grad  # (1, 100, 10)
early_grad_norm  = grads[0, 0, :].norm().item()    # gradient at step 1
late_grad_norm   = grads[0, 99, :].norm().item()   # gradient at step 100

print(f"Gradient norm at step   1 (earliest): {early_grad_norm:.6f}")
print(f"Gradient norm at step 100 (latest):   {late_grad_norm:.6f}")
print(f"Ratio (late/early): {late_grad_norm / (early_grad_norm + 1e-10):.1f}x")

import numpy as np
import matplotlib.pyplot as plt

def tanh_derivative(x):
    return 1 - np.tanh(x) ** 2

x_vals = np.linspace(-3, 3, 200)
deriv  = tanh_derivative(x_vals)

plt.figure(figsize=(8, 3))
plt.plot(x_vals, np.tanh(x_vals), label='tanh(x)', color='steelblue')
plt.plot(x_vals, deriv,           label="tanh'(x) — at most 1.0", color='tomato')
plt.axhline(1.0, color='gray', linestyle='--', alpha=0.5)
plt.axhline(0.0, color='gray', linestyle='-', alpha=0.3)
plt.legend(); plt.xlabel('x'); plt.title("tanh and its derivative")
plt.tight_layout(); plt.show()

# Simulate gradient decay over 100 steps
# Assume each tanh derivative is a constant 0.8 (optimistic!)
T = 100
grad_at_step = [(0.8 ** (T - t)) for t in range(T)]

print(f"Gradient at step   1 (if deriv=0.8 per step): {0.8**99:.2e}")
print(f"Gradient at step  50:                          {0.8**50:.2e}")
print(f"Gradient at step  90:                          {0.8**10:.2e}")
print(f"Gradient at step 100 (final):                  {0.8**0:.2e}")
# Step 1: 2.04e-10 — effectively zero
# Step 100: 1.0    — gradient flows normally

import torch
import torch.nn as nn

class ManualLSTMCell(nn.Module):
    """
    Manual LSTM cell — shows every gate explicitly.
    Identical to nn.LSTMCell.
    """
    def __init__(self, input_size, hidden_size):
        super().__init__()
        # All four transformations are typically batched into one
        # large linear layer for efficiency, but written separately for clarity
        self.forget_gate = nn.Linear(input_size + hidden_size, hidden_size)
        self.input_gate   = nn.Linear(input_size + hidden_size, hidden_size)
        self.cell_gate    = nn.Linear(input_size + hidden_size, hidden_size)  # candidate
        self.output_gate  = nn.Linear(input_size + hidden_size, hidden_size)

    def forward(self, x_t, h_prev, c_prev):
        """
        x_t:   (batch, input_size)
        h_prev: (batch, hidden_size) — short-term memory
        c_prev: (batch, hidden_size) — long-term memory (cell state)
        """
        # Concatenate inputs and previous hidden state
        combined = torch.cat([h_prev, x_t], dim=1)  # (batch, input+hidden)

        # ── Forget gate: what to erase from cell state ──
        # sigmoid: output in (0,1); 0 = forget everything, 1 = remember everything
        f = torch.sigmoid(self.forget_gate(combined))   # (batch, hidden)

        # ── Input gate: what new info to add ──
        i = torch.sigmoid(self.input_gate(combined))    # (batch, hidden)

        # ── Candidate cell values (what the input gate selects from) ──
        g = torch.tanh(self.cell_gate(combined))        # (batch, hidden)

        # ── Cell state update (the gradient highway!) ──
        # Old cell * forget gate + new candidate * input gate
        c_t = f * c_prev + i * g                        # (batch, hidden)

        # ── Output gate: what to expose as hidden state ──
        o = torch.sigmoid(self.output_gate(combined))   # (batch, hidden)

        # ── New hidden state ──
        h_t = o * torch.tanh(c_t)                      # (batch, hidden)

        return h_t, c_t


# Test the manual LSTM cell
cell = ManualLSTMCell(input_size=20, hidden_size=64)
x_t  = torch.randn(4, 20)   # batch=4
h_0  = torch.zeros(4, 64)
c_0  = torch.zeros(4, 64)

h_1, c_1 = cell(x_t, h_0, c_0)
print(f"h_1 shape: {h_1.shape}")  # (4, 64)
print(f"c_1 shape: {c_1.shape}")  # (4, 64)
print(f"h_1 range: [{h_1.min():.3f}, {h_1.max():.3f}]")  # within [-1, 1] (tanh bounded)
print(f"c_1 range: [{c_1.min():.3f}, {c_1.max():.3f}]")  # can grow beyond [-1, 1]

import torch
import torch.nn as nn

# nn.LSTM handles the full sequence loop
lstm = nn.LSTM(
    input_size=20,       # size of each input vector
    hidden_size=64,      # size of hidden state AND cell state
    num_layers=2,        # stacked LSTM layers
    batch_first=True,    # input: (batch, seq_len, features)
    dropout=0.3,         # dropout between LSTM layers (not after last)
    bidirectional=False  # see Section 8 for bidirectional
)

# Input: batch of 4 sequences, 10 steps, 20 features each
x = torch.randn(4, 10, 20)

# Initial states (both h and c for LSTM!)
# Shape: (num_layers * num_directions, batch, hidden_size)
h0 = torch.zeros(2, 4, 64)  # num_layers=2
c0 = torch.zeros(2, 4, 64)

output, (h_n, c_n) = lstm(x, (h0, c0))

print(f"Input:    {x.shape}")
print(f"output:   {output.shape}")  # (4, 10, 64) — hidden state at every step
print(f"h_n:      {h_n.shape}")     # (2, 4, 64)  — final hidden state per layer
print(f"c_n:      {c_n.shape}")     # (2, 4, 64)  — final cell state per layer

# For many-to-one tasks (e.g., sentiment), use only the final output
final_hidden = output[:, -1, :]   # (4, 64) — last time step's hidden state
# OR equivalently: h_n[-1, :, :]  — last layer's final hidden state
print(f"Final hidden: {final_hidden.shape}")

import torch
import torch.nn as nn

# GRU — simpler than LSTM (no cell state, 2 gates instead of 3)
gru = nn.GRU(
    input_size=20,
    hidden_size=64,
    num_layers=2,
    batch_first=True,
    dropout=0.3,
    bidirectional=False
)

x = torch.randn(4, 10, 20)
h0 = torch.zeros(2, 4, 64)   # only h (no cell state c!)

output, h_n = gru(x, h0)      # GRU returns (output, h_n) — no c_n

print(f"GRU output: {output.shape}")  # (4, 10, 64)
print(f"GRU h_n:    {h_n.shape}")     # (2, 4, 64)

# Parameter count comparison
lstm_params = sum(p.numel() for p in nn.LSTM(20, 64, batch_first=True).parameters())
gru_params  = sum(p.numel() for p in nn.GRU(20, 64, batch_first=True).parameters())
print(f"\nLSTM parameters (same hidden_size): {lstm_params:,}")
print(f"GRU  parameters (same hidden_size): {gru_params:,}")
print(f"GRU has {(lstm_params - gru_params)/lstm_params*100:.1f}% fewer parameters")

import torch
import torch.nn as nn

# Bidirectional LSTM processes sequence both forwards and backwards
bilstm = nn.LSTM(
    input_size=20,
    hidden_size=64,
    num_layers=2,
    batch_first=True,
    dropout=0.3,
    bidirectional=True   # ← key parameter
)

x = torch.randn(4, 10, 20)
h0 = torch.zeros(4, 4, 64)  # num_layers * num_directions = 2 * 2 = 4
c0 = torch.zeros(4, 4, 64)

output, (h_n, c_n) = bilstm(x, (h0, c0))

# output: concatenation of forward and backward hidden states at each step
print(f"Bidirectional output: {output.shape}")   # (4, 10, 128) — 64*2 = 128!
print(f"h_n: {h_n.shape}")                       # (4, 4, 64)   — 4 = 2*2

# The first 64 dims = forward direction; last 64 dims = backward direction
forward_out  = output[:, :, :64]   # (4, 10, 64)
backward_out = output[:, :, 64:]   # (4, 10, 64)

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.LSTM(20, 64, batch_first=True)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# In the training loop:
optimizer.zero_grad()
# ... forward pass, compute loss ...
# loss.backward()

# Clip gradients: rescale so total grad norm <= max_norm
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# THEN step the optimizer
optimizer.step()

# Rule of thumb: max_norm = 1.0 is standard for most RNN tasks
# Check if clipping is triggering: total_norm > max_norm means clipping activated

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
import torch
import torch.nn as nn

lstm = nn.LSTM(input_size=20, hidden_size=64, batch_first=True)

# Batch of 3 sequences with lengths 5, 3, 2 (padded to 5)
batch = torch.randn(3, 5, 20)        # (batch=3, max_len=5, features=20)
lengths = torch.tensor([5, 3, 2])    # actual lengths

# Pack: tells PyTorch which positions are real vs padding
packed = pack_padded_sequence(batch, lengths, batch_first=True, enforce_sorted=True)

# Run LSTM on packed sequence (skips padding automatically)
packed_output, (h_n, c_n) = lstm(packed)

# Unpack back to padded tensor (for position-level outputs)
output, output_lengths = pad_packed_sequence(packed_output, batch_first=True)
print(f"Output shape: {output.shape}")   # (3, 5, 64) — padded back to max_len