"""Encoder module for scDistill.

Maps gene expression to batch-corrected latent space.
"""

import torch
import torch.nn as nn
from typing import List


class Encoder(nn.Module):
    """Expression → Latent (batch-free) encoder.

    Maps gene expression to a latent space that matches the teacher's
    batch-corrected embeddings.

    Architecture:
        X (n_input) → [Linear → ReLU → Dropout] × n_layers → LayerNorm → Z (n_latent)

    LayerNorm is applied at the output to normalize per-sample embeddings,
    which is critical for preserving batch mixing from the teacher.

    Parameters
    ----------
    n_input : int
        Number of input features (genes).
    n_latent : int, default=50
        Dimension of latent space.
    n_hidden : int, default=128
        Size of hidden layers.
    n_layers : int, default=2
        Number of hidden layers. n_layers=1 means input → hidden → output.
    dropout : float, default=0.1
        Dropout probability between layers.
    """

    def __init__(
        self,
        n_input: int,
        n_latent: int = 50,
        n_hidden: int = 128,
        n_layers: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()

        self.n_input = n_input
        self.n_latent = n_latent
        self.n_hidden = n_hidden
        self.n_layers = n_layers
        self.dropout = dropout

        # Build encoder layers
        self.encoder = self._build_encoder()

        # LayerNorm at output (normalizes per-sample embeddings)
        # Critical for preserving batch mixing from teacher
        self.layer_norm = nn.LayerNorm(n_latent)

    def _build_encoder(self) -> nn.Sequential:
        """Build encoder MLP."""
        layers: List[nn.Module] = []

        # Input layer
        layers.append(nn.Linear(self.n_input, self.n_hidden))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(self.dropout))

        # Hidden layers
        for _ in range(self.n_layers - 1):
            layers.append(nn.Linear(self.n_hidden, self.n_hidden))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(self.dropout))

        # Output layer (no activation)
        layers.append(nn.Linear(self.n_hidden, self.n_latent))

        return nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Encode expression to latent space.

        Parameters
        ----------
        x : torch.Tensor
            Gene expression matrix (N, n_input).
            Should be normalized (e.g., log2(CPM+1)).

        Returns
        -------
        z : torch.Tensor
            Latent embeddings (N, n_latent), LayerNorm normalized.
        """
        z = self.encoder(x)
        z = self.layer_norm(z)
        return z
