"""Decoder module for scDistill.

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

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple, Optional


class Decoder(nn.Module):
    """Latent → Batch-corrected expression decoder with NB likelihood.

    Maps batch-corrected latent embeddings back to gene expression space.
    Outputs parameters for Negative Binomial distribution:
    - rho: normalized rate (softmax over genes)
    - theta: gene-wise dispersion parameter

    Architecture:
        Z (n_latent) → [Linear → ReLU → Dropout] × n_layers → rho (n_output)
        theta: learnable gene-wise parameter

    Parameters
    ----------
    n_latent : int, default=50
        Dimension of latent space.
    n_output : int
        Number of output features (genes).
    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.
    n_batches : int, default=1
        Number of batches for batch conditioning.
    n_batch_embedding : int, default=10
        Dimension of batch embedding. Uses nn.Embedding for scalability
        when n_batches is large (atlas integration scenarios).
    use_batch_conditioning : bool, default=False
        If True, concatenate batch embedding to latent embeddings.
    """

    def __init__(
        self,
        n_latent: int,
        n_output: int,
        n_hidden: int = 128,
        n_layers: int = 2,
        dropout: float = 0.1,
        n_batches: int = 1,
        n_batch_embedding: int = 10,
        use_batch_conditioning: bool = False,
    ):
        super().__init__()

        self.n_latent = n_latent
        self.n_output = n_output
        self.n_hidden = n_hidden
        self.n_layers = n_layers
        self.dropout = dropout
        self.n_batches = n_batches
        self.n_batch_embedding = n_batch_embedding
        self.use_batch_conditioning = use_batch_conditioning

        # Batch embedding layer (scalable for large n_batches)
        if use_batch_conditioning:
            self.batch_embedding = nn.Embedding(n_batches, n_batch_embedding)

        # Build decoder layers (outputs unnormalized log-rates)
        self.decoder = self._build_decoder()

        # Gene-wise dispersion parameter (log-scale for numerical stability)
        # Initialize to log(1) = 0, so theta starts at 1
        self.log_theta = nn.Parameter(torch.zeros(n_output))

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

        # Input dimension: n_latent + n_batch_embedding if batch conditioning enabled
        input_dim = self.n_latent + self.n_batch_embedding if self.use_batch_conditioning else self.n_latent

        # Input layer
        layers.append(nn.Linear(input_dim, 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 - will apply softmax for rho)
        layers.append(nn.Linear(self.n_hidden, self.n_output))

        return nn.Sequential(*layers)

    def _apply_batch_conditioning(
        self, z: torch.Tensor, batch: Optional[torch.Tensor]
    ) -> torch.Tensor:
        """Apply batch conditioning to latent embeddings.

        Uses nn.Embedding for scalability when n_batches is large.

        Parameters
        ----------
        z : torch.Tensor
            Latent embeddings (N, n_latent).
        batch : torch.Tensor, optional
            Batch indices (N,). Required if use_batch_conditioning=True.

        Returns
        -------
        z_conditioned : torch.Tensor
            Conditioned latent embeddings (N, n_latent + n_batch_embedding) or (N, n_latent).
        """
        if not self.use_batch_conditioning:
            return z
        if batch is None:
            raise ValueError("batch is required when use_batch_conditioning=True")
        batch_emb = self.batch_embedding(batch.long())
        return torch.cat([z, batch_emb], dim=-1)

    def forward(
        self,
        z: torch.Tensor,
        library_size: Optional[torch.Tensor] = None,
        batch: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Decode latent embeddings to normalized expression (log-scale).

        For backward compatibility, returns log-normalized expression
        when library_size is not provided.

        Parameters
        ----------
        z : torch.Tensor
            Latent embeddings (N, n_latent).
        library_size : torch.Tensor, optional
            Library size per cell (N,). If provided, returns log2(mu + 1).
        batch : torch.Tensor, optional
            Batch indices (N,). Required if use_batch_conditioning=True.

        Returns
        -------
        x : torch.Tensor
            Batch-corrected expression in log2(CPM + 1) scale (N, n_output).
        """
        # Apply batch conditioning
        z = self._apply_batch_conditioning(z, batch)

        # Get unnormalized log-rates
        log_rho_unnorm = self.decoder(z)

        # Normalized rate (softmax over genes)
        rho = F.softmax(log_rho_unnorm, dim=-1)

        if library_size is not None:
            # mu = library_size * rho (expected counts)
            mu = library_size.unsqueeze(-1) * rho

            # Return in log2(CPM + 1) scale for compatibility
            # CPM = (mu / library_size) * 1e6 = rho * 1e6
            cpm = rho * 1e6
            return torch.log2(cpm + 1)
        else:
            # Legacy mode: return log2(rho * 1e6 + 1)
            cpm = rho * 1e6
            return torch.log2(cpm + 1)

    def forward_nb(
        self,
        z: torch.Tensor,
        library_size: torch.Tensor,
        batch: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Forward pass returning NB parameters.

        Parameters
        ----------
        z : torch.Tensor
            Latent embeddings (N, n_latent).
        library_size : torch.Tensor
            Library size per cell (N,).
        batch : torch.Tensor, optional
            Batch indices (N,). Required if use_batch_conditioning=True.

        Returns
        -------
        mu : torch.Tensor
            Expected counts (N, n_output).
        theta : torch.Tensor
            Dispersion parameter (n_output,).
        """
        # Apply batch conditioning
        z = self._apply_batch_conditioning(z, batch)

        # Get unnormalized log-rates
        log_rho_unnorm = self.decoder(z)

        # Normalized rate (softmax over genes)
        rho = F.softmax(log_rho_unnorm, dim=-1)

        # mu = library_size * rho
        mu = library_size.unsqueeze(-1) * rho

        # Dispersion (always positive)
        theta = torch.exp(self.log_theta)

        return mu, theta

    def get_rho(
        self, z: torch.Tensor, batch: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Get normalized rate (rho) from latent embeddings.

        Parameters
        ----------
        z : torch.Tensor
            Latent embeddings (N, n_latent).
        batch : torch.Tensor, optional
            Batch indices (N,). Required if use_batch_conditioning=True.

        Returns
        -------
        rho : torch.Tensor
            Normalized rate (N, n_output). Sum over genes = 1.
        """
        # Apply batch conditioning
        z = self._apply_batch_conditioning(z, batch)

        log_rho_unnorm = self.decoder(z)
        return F.softmax(log_rho_unnorm, dim=-1)

    def get_theta(self) -> torch.Tensor:
        """Get dispersion parameter."""
        return torch.exp(self.log_theta)
