import torch
import torch.nn as nn
from gemmasight.config import DIM_FUSED

class MSIClassifier(nn.Module):
    def __init__(self, input_dim=DIM_FUSED):
        super().__init__()
        
        # Dense(512) -> BatchNorm -> ReLU -> Dropout(0.4)
        self.layer1 = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.4)
        )
        
        # Dense(256) -> BatchNorm -> ReLU -> Dropout(0.3)
        self.layer2 = nn.Sequential(
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3)
        )
        
        # Dense(128) -> BatchNorm -> ReLU -> Dropout(0.2)
        self.layer3 = nn.Sequential(
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.2)
        )
        
        # Dense(1) -> Sigmoid
        self.out = nn.Sequential(
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Input: Tensor of shape (Batch, 1536)
        Output: Tensor of shape (Batch, 1) representing MSI-High probability
        """
        # Handle single batch case for input format compatibility
        is_single = len(x.shape) == 1
        if is_single:
            x = x.unsqueeze(0)
            
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        prob = self.out(x)
        
        if is_single:
            prob = prob.squeeze(0)
        return prob