"""
uncertainty.py
Monte Carlo Dropout uncertainty estimation.
Exact implementation from [IDSC]_D4.ipynb Cell 11.
"""

import torch
import torch.nn as nn
import numpy as np
from typing import Tuple

UNCERTAINTY_THRESHOLD = 0.05  # From notebook: UNCERTAINTY_THRESHOLD = 0.05
MC_SAMPLES = 1000             # Updated to 1000 samples for prediction

def enable_dropout(model: nn.Module):
    """
    Enable dropout layers during inference for MC Dropout.
    (Sets all Dropout modules to training mode.)
    """
    for m in model.modules():
        if isinstance(m, nn.Dropout):
            m.train()


def mc_dropout_predict(model: nn.Module,
                        x_tensor: torch.Tensor,
                        n_samples: int = MC_SAMPLES,
                        device: str = 'cpu') -> Tuple[float, float, np.ndarray]:
    """
    Run model n_samples forward passes with dropout active.

    Args:
        model:    MLPClassifier with dropout layers
        x_tensor: input feature tensor shape (1, input_dim)
        n_samples: number of MC passes (default 1000)
        device:   'cpu' or 'cuda'

    Returns:
        mean_prob: float, mean prediction probability
        variance:  float, prediction variance (= uncertainty)
        all_preds: np.ndarray shape (n_samples,)
    """
    x_tensor = x_tensor.to(device)
    model.eval()
    enable_dropout(model)  # activate dropout

    with torch.no_grad():
        # Duplicate the input n_samples times to run as a single batch
        x_batch = x_tensor.repeat(n_samples, 1)
        
        # Forward pass for all samples at once
        preds = model(x_batch).cpu().numpy().flatten()

    mean_prob = float(preds.mean())
    variance = float(preds.var())

    return mean_prob, variance, preds


def interpret_uncertainty(variance: float) -> dict:
    """
    Interpret the uncertainty level for display.
    """
    is_ambiguous = variance > UNCERTAINTY_THRESHOLD
    level = 'HIGH' if variance > 0.10 else ('MEDIUM' if variance > 0.05 else 'LOW')

    return {
        'variance': round(variance, 6),
        'threshold': UNCERTAINTY_THRESHOLD,
        'is_ambiguous': is_ambiguous,
        'uncertainty_level': level,
        'uncertainty_pct': round(min(variance / 0.15, 1.0) * 100, 1),  # 0–100% for UI
        'refer_to_specialist': is_ambiguous,
    }