import torch
import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig

# 1. Definiamo la configurazione con i tuoi valori di default
class AITextClassifierConfig(PretrainedConfig):
    model_type = "ai_text_classifier"
    
    def __init__(
        self, 
        input_dim: int = 1024, 
        dropout_high: float = 0.25, 
        dropout_low: float = 0.1, 
        **kwargs
    ):
        super().__init__(**kwargs)
        self.input_dim = input_dim
        self.dropout_high = dropout_high
        self.dropout_low = dropout_low


# 2. Adattiamo il tuo modello ereditando da PreTrainedModel
class AITextClassifier(PreTrainedModel):
    config_class = AITextClassifierConfig

    # Il costruttore ora accetta l'oggetto 'config' di Hugging Face
    def __init__(self, config: AITextClassifierConfig):
        super().__init__(config)

        # Ricreiamo esattamente la tua rete usando i parametri estratti dal config
        self.network = nn.Sequential(
            # --- Layer 1 ---
            nn.Linear(config.input_dim, 512),
            nn.BatchNorm1d(512),
            nn.GELU(),
            nn.Dropout(config.dropout_high),

            # --- Layer 2 ---
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.GELU(),
            nn.Dropout(config.dropout_high),

            # --- Layer 3 ---
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.GELU(),
            nn.Dropout(config.dropout_low),

            # --- Output ---
            nn.Linear(128, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x).squeeze(-1)

    @torch.no_grad()
    def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
        self.eval()
        return torch.sigmoid(self.forward(x))