"""
Personality classifier model definition.

The classifier is a simple MLP that maps 384-dimensional MiniLM embeddings
to 10 personality dimension scores in the range [0, 1].

Architecture:
    Input: 384-dim embedding
    Hidden 1: 256 units + ReLU + Dropout
    Hidden 2: 128 units + ReLU + Dropout
    Output: 10 units + Sigmoid
"""

from __future__ import annotations

import torch
import torch.nn as nn


PERSONALITY_DIMENSIONS = [
    "openness",
    "conscientiousness",
    "extraversion",
    "agreeableness",
    "neuroticism",
    "assertiveness",
    "ambition",
    "resilience",
    "integrity",
    "curiosity",
]


class PersonalityClassifier(nn.Module):
    """
    MLP classifier that maps text embeddings to personality scores.

    The model takes 384-dimensional embeddings from MiniLM and outputs
    10 scores (one per personality dimension) in the range [0, 1].

    Args:
        input_dim: Dimension of input embeddings (default: 384 for MiniLM)
        hidden_dims: Tuple of hidden layer dimensions (default: (256, 128))
        output_dim: Number of personality dimensions (default: 10)
        dropout: Dropout probability (default: 0.1)
    """

    def __init__(
        self,
        input_dim: int = 384,
        hidden_dims: tuple[int, int] = (256, 128),
        output_dim: int = 10,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()

        self.input_dim = input_dim
        self.hidden_dims = hidden_dims
        self.output_dim = output_dim
        self.dropout_rate = dropout

        # Build MLP layers
        self.layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dims[0]),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dims[0], hidden_dims[1]),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dims[1], output_dim),
            nn.Sigmoid(),  # Output scores in [0, 1]
        )

        # Initialize weights
        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize weights using Xavier uniform initialization."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass.

        Args:
            x: Input embeddings of shape (batch_size, input_dim)

        Returns:
            Personality scores of shape (batch_size, output_dim)
        """
        return self.layers(x)

    def predict(self, embeddings: torch.Tensor) -> dict[str, float]:
        """
        Make predictions and return as a dictionary.

        Args:
            embeddings: Single embedding of shape (input_dim,) or (1, input_dim)

        Returns:
            Dictionary mapping dimension names to scores
        """
        self.eval()
        with torch.no_grad():
            if embeddings.dim() == 1:
                embeddings = embeddings.unsqueeze(0)
            scores = self.forward(embeddings)
            scores = scores.squeeze(0).cpu().numpy()

        return {dim: float(scores[i]) for i, dim in enumerate(PERSONALITY_DIMENSIONS)}

    def get_config(self) -> dict:
        """Return model configuration for serialization."""
        return {
            "input_dim": self.input_dim,
            "hidden_dims": list(self.hidden_dims),
            "output_dim": self.output_dim,
            "dropout": self.dropout_rate,
        }

    @classmethod
    def from_config(cls, config: dict) -> "PersonalityClassifier":
        """Create model from configuration dictionary."""
        return cls(
            input_dim=config.get("input_dim", 384),
            hidden_dims=tuple(config.get("hidden_dims", [256, 128])),
            output_dim=config.get("output_dim", 10),
            dropout=config.get("dropout", 0.1),
        )
