"""
Autism Detector Model

A feedforward neural network for ASD risk classification
from structured clinical data.
"""

import torch
import torch.nn as nn


class AutismDetector(nn.Module):
    """
    Binary classifier for autism spectrum disorder screening.

    Input: 8 preprocessed clinical features
    Output: 2 logits (Healthy, ASD)

    Features (in order):
        1. developmental_milestones - N/G/M/C (encoded 0-3)
        2. iq_dq - numeric, normalized 0-1
        3. intellectual_disability - N/F70.0/F71/F72 (encoded 0-3)
        4. language_disorder - N/Y (encoded 0-1)
        5. language_development - N/delay/A (encoded 0-2)
        6. dysmorphism - NO/Y (encoded 0-1)
        7. behaviour_disorder - N/Y (encoded 0-1)
        8. neurological_exam - N/abnormal (encoded 0-1)
    """

    def __init__(self, input_size=8, hidden_sizes=None, num_classes=2, dropout=0.3):
        super().__init__()

        if hidden_sizes is None:
            hidden_sizes = [64, 32]

        layers = []
        prev_size = input_size

        for hidden_size in hidden_sizes:
            layers.extend([
                nn.Linear(prev_size, hidden_size),
                nn.ReLU(),
                nn.Dropout(dropout),
            ])
            prev_size = hidden_size

        layers.append(nn.Linear(prev_size, num_classes))
        self.classifier = nn.Sequential(*layers)

        # Store config
        self.input_size = input_size
        self.hidden_sizes = hidden_sizes
        self.num_classes = num_classes
        self.dropout = dropout

    def forward(self, x):
        """
        Forward pass.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape (batch_size, 8)

        Returns
        -------
        torch.Tensor
            Output logits of shape (batch_size, num_classes)
        """
        return self.classifier(x)

    def predict(self, x):
        """
        Make predictions with probabilities.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape (batch_size, 8)

        Returns
        -------
        dict with 'prediction', 'probability', 'logits'
        """
        self.eval()
        with torch.no_grad():
            logits = self.forward(x)
            probs = torch.softmax(logits, dim=-1)
            pred_class = torch.argmax(probs, dim=-1)

        return {
            'prediction': pred_class,
            'probabilities': probs,
            'logits': logits
        }


def load_model(model_path, device='cpu'):
    """Load TorchScript model."""
    model = torch.jit.load(model_path, map_location=device)
    model.eval()
    return model


def preprocess(data, config):
    """
    Preprocess input data using JSON config.

    Parameters
    ----------
    data : dict
        Input features as dictionary
    config : dict
        Preprocessor configuration from preprocessor_config.json

    Returns
    -------
    torch.Tensor
        Preprocessed features tensor of shape (1, 8)
    """
    features = []

    for feature_name in config["feature_order"]:
        if feature_name in config["categorical_features"]:
            feat_config = config["categorical_features"][feature_name]

            if feat_config["type"] == "text_binary":
                # For neurological_exam: N -> 0, anything else -> 1
                raw_value = str(data[feature_name]).strip().upper()
                value = 0 if raw_value == feat_config["normal_value"] else 1
            else:
                # Standard categorical/binary mapping
                raw_value = data[feature_name]
                value = feat_config["mapping"].get(raw_value, 0)

        elif feature_name in config["numeric_features"]:
            feat_config = config["numeric_features"][feature_name]
            raw = float(data[feature_name])
            # Min-max normalization
            value = (raw - feat_config["min"]) / (feat_config["max"] - feat_config["min"])
            value = max(0, min(1, value))  # Clamp to [0, 1]

        features.append(value)

    return torch.tensor([features], dtype=torch.float32)


def get_risk_level(probability):
    """
    Get risk level from ASD probability.

    Returns
    -------
    str: 'low', 'medium', or 'high'
    """
    if probability < 0.4:
        return "low"
    elif probability < 0.7:
        return "medium"
    else:
        return "high"


if __name__ == '__main__':
    # Test model creation
    model = AutismDetector()
    print(f"Model architecture:\n{model}")

    # Test forward pass
    x = torch.randn(2, 8)
    output = model(x)
    print(f"\nInput shape: {x.shape}")
    print(f"Output shape: {output.shape}")
    print(f"Output (logits): {output}")

    probs = torch.softmax(output, dim=-1)
    print(f"Probabilities: {probs}")