"""Paper-faithful PyTorch CNN-LSTM baseline for ClimateBench."""

from __future__ import annotations

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


TARGETS = ("tas", "dtr", "pr", "pr90")


class ReLULSTM(nn.Module):
    """Keras-compatible LSTM using sigmoid gates and ReLU activation."""

    def __init__(self, input_size: int = 20, hidden_size: int = 25):
        super().__init__()
        self.input_size = int(input_size)
        self.hidden_size = int(hidden_size)
        self.kernel = nn.Parameter(torch.empty(input_size, 4 * hidden_size))
        self.recurrent_kernel = nn.Parameter(torch.empty(hidden_size, 4 * hidden_size))
        self.bias = nn.Parameter(torch.zeros(4 * hidden_size))
        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.xavier_uniform_(self.kernel)
        nn.init.orthogonal_(self.recurrent_kernel)
        nn.init.zeros_(self.bias)
        with torch.no_grad():
            self.bias[self.hidden_size:2 * self.hidden_size].fill_(1.0)

    def forward(self, sequence: torch.Tensor) -> torch.Tensor:
        if sequence.ndim != 3 or sequence.shape[-1] != self.input_size:
            raise ValueError(f"expected [B,T,{self.input_size}], got {tuple(sequence.shape)}")
        batch = sequence.shape[0]
        hidden = sequence.new_zeros(batch, self.hidden_size)
        cell = sequence.new_zeros(batch, self.hidden_size)
        for step in range(sequence.shape[1]):
            gates = sequence[:, step] @ self.kernel + hidden @ self.recurrent_kernel + self.bias
            input_gate, forget_gate, candidate, output_gate = gates.chunk(4, dim=-1)
            input_gate = torch.sigmoid(input_gate)
            forget_gate = torch.sigmoid(forget_gate)
            candidate = F.relu(candidate)
            output_gate = torch.sigmoid(output_gate)
            cell = forget_gate * cell + input_gate * candidate
            hidden = output_gate * F.relu(cell)
        return hidden


class ClimateBenchBranch(nn.Module):
    """One 364,764-parameter paper CNN-LSTM target emulator."""

    PAPER_PARAMETER_COUNT = 364_764

    def __init__(self, height: int = 96, width: int = 144):
        super().__init__()
        self.height = int(height)
        self.width = int(width)
        self.conv = nn.Conv2d(4, 20, kernel_size=3, padding="same")
        self.pool = nn.AvgPool2d(kernel_size=2)
        self.lstm = ReLULSTM(20, 25)
        self.dense = nn.Linear(25, height * width)
        count = sum(parameter.numel() for parameter in self.parameters())
        if count != self.PAPER_PARAMETER_COUNT:
            raise RuntimeError(f"paper branch must have 364764 parameters, got {count}")

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        if inputs.ndim != 5 or inputs.shape[2:] != (4, self.height, self.width):
            raise ValueError(f"expected [B,T,4,{self.height},{self.width}], got {tuple(inputs.shape)}")
        batch, time = inputs.shape[:2]
        features = F.relu(self.conv(inputs.reshape(batch * time, 4, self.height, self.width)))
        features = self.pool(features).mean(dim=(-2, -1)).reshape(batch, time, 20)
        return self.dense(self.lstm(features)).reshape(batch, 1, self.height, self.width)


class ClimateBench(nn.Module):
    """Four independent paper branches ordered as tas, dtr, pr and pr90."""

    def __init__(self, height: int = 96, width: int = 144):
        super().__init__()
        self.height = int(height)
        self.width = int(width)
        self.branches = nn.ModuleDict({name: ClimateBenchBranch(height, width) for name in TARGETS})

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return torch.cat([self.branches[name](inputs) for name in TARGETS], dim=1)

    def parameter_counts(self) -> dict[str, int]:
        return {name: sum(parameter.numel() for parameter in branch.parameters())
                for name, branch in self.branches.items()}