from __future__ import annotations

import torch
from torch import nn


class JEncoder(nn.Module):
    def __init__(self, latent_dim: int = 32) -> None:
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(64, 96),
            nn.LayerNorm(96),
            nn.GELU(),
            nn.Linear(96, 64),
            nn.LayerNorm(64),
            nn.GELU(),
            nn.Linear(64, latent_dim),
        )

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        return self.network(images.flatten(1))


class JPredictor(nn.Module):
    def __init__(self, latent_dim: int = 32) -> None:
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(latent_dim, 64),
            nn.GELU(),
            nn.Linear(64, latent_dim),
        )

    def forward(self, embeddings: torch.Tensor) -> torch.Tensor:
        return self.network(embeddings)


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