from __future__ import annotations

import torch
from torch import nn


class ChronosMicroGRU(nn.Module):
    def __init__(
        self,
        channels: int = 6,
        hidden_size: int = 24,
        horizon: int = 8,
    ) -> None:
        super().__init__()
        self.channels = channels
        self.horizon = horizon
        self.recurrent = nn.GRU(
            input_size=channels,
            hidden_size=hidden_size,
            batch_first=True,
        )
        self.head = nn.Sequential(
            nn.Linear(hidden_size, 48),
            nn.GELU(),
            nn.Linear(48, horizon * channels * 2),
        )

    def forward(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        _, hidden = self.recurrent(context)
        output = self.head(hidden[-1]).reshape(
            len(context),
            self.horizon,
            self.channels,
            2,
        )
        mean = output[..., 0]
        log_variance = torch.clamp(output[..., 1], min=-6, max=3)
        return mean, log_variance


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