new_memory = tanh(current_input + old_memory)

hidden = h0
for token in tokens:
    hidden = update(token, hidden)

#!/usr/bin/env python
"""Chapter 2.5: hands-on RNN and CNN examples."""

from __future__ import annotations

import argparse
import math
import random

import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset


def demo_manual_rnn() -> None:
    print("=== 1. manual RNN: h_t = tanh(0.8*x_t + 0.6*h_{t-1}) ===")
    inputs = [1.0, 0.5, 1.2]
    hidden = 0.0

    for step, value in enumerate(inputs, start=1):
        old_hidden = hidden
        hidden = math.tanh(0.8 * value + 0.6 * old_hidden)
        print(
            f"step {step} | x={value:.3f} | old h={old_hidden:+.6f} | "
            f"new h={hidden:+.6f}"
        )


def demo_pytorch_rnn() -> None:
    print("\n=== 2. PyTorch nn.RNN shapes ===")
    model = nn.RNN(input_size=3, hidden_size=4, batch_first=True)
    inputs = torch.randn(2, 5, 3)
    outputs, final_hidden = model(inputs)

    print(f"input shape:        {tuple(inputs.shape)}")
    print(f"all outputs shape:  {tuple(outputs.shape)}")
    print(f"final hidden shape: {tuple(final_hidden.shape)}")
    print("meaning: (batch=2, seq_len=5, input_dim=3) -> (batch=2, seq_len=5, hidden_dim=4)")


def demo_conv1d() -> None:
    print("\n=== 3. manual Conv1d example ===")
    inputs = torch.tensor([[[1.0, 0.0, -1.0, 2.0, 1.0]]])
    weight = torch.tensor([[[0.25, 0.50, 0.25]]])
    bias = torch.tensor([0.0])
    output = F.conv1d(inputs, weight, bias=bias, stride=1, padding=0)

    print(f"input shape:  {tuple(inputs.shape)}")
    print(f"weight shape: {tuple(weight.shape)}")
    print(f"output shape: {tuple(output.shape)}")
    print(f"input values:  {inputs[0, 0].tolist()}")
    print(f"kernel values: {weight[0, 0].tolist()}")
    print(f"output values: {output[0, 0].tolist()}")


def demo_conv2d_edge_detector() -> None:
    print("\n=== 4. Conv2d edge detector ===")
    image = torch.tensor(
        [
            [
                [
                    [1.0, 1.0, 1.0, 0.0, 0.0],
                    [1.0, 1.0, 1.0, 0.0, 0.0],
                    [1.0, 1.0, 1.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0],
                ]
            ]
        ]
    )
    vertical_edge_kernel = torch.tensor(
        [
            [
                [
                    [1.0, 0.0, -1.0],
                    [1.0, 0.0, -1.0],
                    [1.0, 0.0, -1.0],
                ]
            ]
        ]
    )
    output = F.conv2d(image, vertical_edge_kernel)

    print(f"image shape:  {tuple(image.shape)}")
    print(f"kernel shape: {tuple(vertical_edge_kernel.shape)}")
    print(f"output shape: {tuple(output.shape)}")
    print("feature map:")
    for row in output[0, 0]:
        print(" ".join(f"{value:+.1f}" for value in row.tolist()))


class RNNClassifier(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, num_classes: int) -> None:
        super().__init__()
        self.rnn = nn.RNN(input_dim, hidden_dim, batch_first=True)
        self.classifier = nn.Linear(hidden_dim, num_classes)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        _, final_hidden = self.rnn(inputs)
        return self.classifier(final_hidden.squeeze(0))


class CNNClassifier(nn.Module):
    def __init__(self, input_dim: int, num_channels: int, num_classes: int) -> None:
        super().__init__()
        self.conv = nn.Conv1d(input_dim, num_channels, kernel_size=3, padding=1)
        self.activation = nn.ReLU()
        self.pooling = nn.AdaptiveAvgPool1d(1)
        self.classifier = nn.Linear(num_channels, num_classes)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        # Text tensors are (batch, seq_len, features), Conv1d wants
        # (batch, features, seq_len).
        features = inputs.transpose(1, 2)
        features = self.pooling(self.activation(self.conv(features))).squeeze(-1)
        return self.classifier(features)


def make_sequence_dataset(
    sample_count: int,
    sequence_length: int,
    seed: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Create a toy binary sequence classification dataset.

    Label 1 sequences contain the consecutive pattern 0 -> 1.
    Label 0 sequences contain the consecutive pattern 2 -> 3.
    """
    rng = random.Random(seed)
    token_ids: list[list[int]] = []
    labels: list[int] = []

    for index in range(sample_count):
        label = index % 2
        sequence = [rng.choice([0, 2]) for _ in range(sequence_length)]
        position = rng.randrange(0, sequence_length - 1)
        if label == 1:
            sequence[position] = 0
            sequence[position + 1] = 1
        else:
            sequence[position] = 2
            sequence[position + 1] = 3
        token_ids.append(sequence)
        labels.append(label)

    inputs = F.one_hot(torch.tensor(token_ids, dtype=torch.long), num_classes=4).float()
    targets = torch.tensor(labels, dtype=torch.long)
    return inputs, targets


def train_classifier(
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
    epochs: int,
    learning_rate: float,
) -> list[float]:
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    loss_fn = nn.CrossEntropyLoss()
    losses: list[float] = []

    for _ in range(epochs):
        model.train()
        epoch_losses: list[float] = []
        for inputs, targets in loader:
            inputs = inputs.to(device)
            targets = targets.to(device)
            optimizer.zero_grad(set_to_none=True)
            logits = model(inputs)
            loss = loss_fn(logits, targets)
            loss.backward()
            optimizer.step()
            epoch_losses.append(loss.item())
        losses.append(sum(epoch_losses) / len(epoch_losses))

    return losses


@torch.inference_mode()
def evaluate_classifier(model: nn.Module, loader: DataLoader, device: torch.device) -> float:
    correct = 0
    total = 0
    model.eval()
    for inputs, targets in loader:
        inputs = inputs.to(device)
        targets = targets.to(device)
        predictions = model(inputs).argmax(dim=1)
        correct += (predictions == targets).sum().item()
        total += targets.numel()
    return correct / total


def compare_rnn_cnn(device: torch.device, epochs: int) -> None:
    print("\n=== 5. toy sequence classification: RNN vs CNN ===")
    inputs, targets = make_sequence_dataset(sample_count=800, sequence_length=16, seed=1337)
    split = 640
    train_dataset = TensorDataset(inputs[:split], targets[:split])
    test_dataset = TensorDataset(inputs[split:], targets[split:])
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=64)

    models: dict[str, nn.Module] = {
        "RNN": RNNClassifier(input_dim=4, hidden_dim=16, num_classes=2).to(device),
        "CNN": CNNClassifier(input_dim=4, num_channels=16, num_classes=2).to(device),
    }

    for name, model in models.items():
        losses = train_classifier(model, train_loader, device, epochs, learning_rate=0.01)
        accuracy = evaluate_classifier(model, test_loader, device)
        parameter_count = sum(parameter.numel() for parameter in model.parameters())
        print(
            f"{name} | final loss {losses[-1]:.4f} | "
            f"test accuracy {accuracy:.4f} | parameters {parameter_count:,}"
        )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--seed", type=int, default=1337)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device(args.device)

    demo_manual_rnn()
    demo_pytorch_rnn()
    demo_conv1d()
    demo_conv2d_edge_detector()
    compare_rnn_cnn(device, args.epochs)
    print(f"\ndevice: {device}")


if __name__ == "__main__":
    main()