from __future__ import annotations

import torch
from torch import nn


class ConditionalNeuralProcess(nn.Module):
    def __init__(self, representation_dimensions: int = 64) -> None:
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(2, 64),
            nn.ReLU(),
            nn.Linear(64, representation_dimensions),
            nn.ReLU(),
        )
        self.decoder = nn.Sequential(
            nn.Linear(representation_dimensions + 1, 64),
            nn.ReLU(),
            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, 2),
        )

    def forward(
        self,
        context_x: torch.Tensor,
        context_y: torch.Tensor,
        target_x: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        pairs = torch.cat([context_x, context_y], dim=2)
        representation = self.encoder(pairs).mean(dim=1)
        expanded = representation[:, None].expand(-1, target_x.shape[1], -1)
        output = self.decoder(torch.cat([target_x, expanded], dim=2))
        mean = output[..., :1]
        standard_deviation = 0.03 + 0.97 * torch.nn.functional.softplus(
            output[..., 1:]
        )
        return mean, standard_deviation


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