from __future__ import annotations

import torch
from torch import nn


class HighwayNetwork(nn.Module):
    def __init__(self, width: int = 32, layers: int = 8) -> None:
        super().__init__()
        self.input = nn.Linear(64, width)
        self.transforms = nn.ModuleList([nn.Linear(width, width) for _ in range(layers)])
        self.gates = nn.ModuleList([nn.Linear(width, width) for _ in range(layers)])
        self.output = nn.Linear(width, 10)
        for gate in self.gates:
            nn.init.constant_(gate.bias, -2.0)

    def forward(
        self,
        pixels: torch.Tensor,
        *,
        return_gates: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        hidden = torch.tanh(self.input(pixels))
        gate_values = []
        for transform, gate in zip(self.transforms, self.gates, strict=True):
            transformed = torch.tanh(transform(hidden))
            carry = torch.sigmoid(gate(hidden))
            hidden = carry * transformed + (1 - carry) * hidden
            gate_values.append(carry)
        logits = self.output(hidden)
        if return_gates:
            return logits, torch.stack(gate_values, dim=1)
        return logits


class PlainDeepNetwork(nn.Module):
    def __init__(self, width: int = 32, layers: int = 16) -> None:
        super().__init__()
        self.input = nn.Linear(64, width)
        self.layers = nn.ModuleList([nn.Linear(width, width) for _ in range(layers)])
        self.output = nn.Linear(width, 10)

    def forward(self, pixels: torch.Tensor) -> torch.Tensor:
        hidden = torch.tanh(self.input(pixels))
        for layer in self.layers:
            hidden = torch.tanh(layer(hidden))
        return self.output(hidden)


def parameter_count(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())