from __future__ import annotations

import torch
from torch import nn


class StochasticGate(nn.Module):
    def __init__(
        self,
        input_dim: int,
        sigma: float,
        temperature: float,
        top_k: int,
    ) -> None:
        super().__init__()
        self.logits = nn.Parameter(torch.zeros(input_dim))
        self.sigma = sigma
        self.temperature = temperature
        self.top_k = top_k

    def gate_probabilities(self) -> torch.Tensor:
        return torch.sigmoid(self.logits / self.temperature)

    def forward(self, x: torch.Tensor, hard_topk: bool = False) -> torch.Tensor:
        probabilities = self.gate_probabilities()
        gates = probabilities
        if hard_topk and 0 < self.top_k < gates.numel():
            _, top_indices = torch.topk(probabilities, self.top_k)
            mask = torch.zeros_like(gates)
            mask[top_indices] = 1.0
            gates = mask
        return x * gates


class ResidualMLPBlock(nn.Module):
    def __init__(self, dim: int, dropout: float) -> None:
        super().__init__()
        self.block = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, dim * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim * 2, dim),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.block(x)


class KGStableGateDANN(nn.Module):
    def __init__(
        self,
        input_dim: int,
        n_classes: int,
        n_domains: int,
        hidden_dim: int,
        projection_dim: int,
        encoder_blocks: int,
        dropout: float,
        top_k: int,
        gate_sigma: float,
        gate_temperature: float,
    ) -> None:
        super().__init__()
        self.gates = StochasticGate(input_dim, gate_sigma, gate_temperature, top_k)
        self.input_projection = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
        )
        self.encoder = nn.Sequential(
            *[ResidualMLPBlock(hidden_dim, dropout) for _ in range(encoder_blocks)]
        )
        self.projection_head = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, projection_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(projection_dim, projection_dim),
            nn.LayerNorm(projection_dim),
        )
        self.disease_classifier = nn.Sequential(
            nn.Linear(projection_dim, max(2, projection_dim // 2)),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(max(2, projection_dim // 2), n_classes),
        )
        self.domain_classifier = nn.Sequential(
            nn.Linear(projection_dim, max(2, projection_dim // 2)),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(max(2, projection_dim // 2), n_domains),
        )
        self.register_buffer("gate_prior_target", torch.zeros(input_dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.gates(x, hard_topk=True)
        hidden = self.input_projection(x)
        hidden = self.encoder(hidden)
        embedding = self.projection_head(hidden)
        return self.disease_classifier(embedding)