import torch
import numpy as np


def mc_dropout_inference(model, images, num_samples=50, device='cpu'):
    model.eval()
    model.enable_mc_dropout()
    all_preds = []
    with torch.no_grad():
        for _ in range(num_samples):
            preds = model(images.to(device))
            all_preds.append(preds.cpu().numpy())
    all_preds = np.array(all_preds)
    mean_preds = all_preds.mean(axis=0)
    epistemic_std = all_preds.std(axis=0)
    return mean_preds, epistemic_std


def compute_calibration_curve(true_rs, pred_rs, pred_std, num_bins=20):
    confidence_levels = np.linspace(0.05, 0.95, num_bins)
    observed_coverage = []
    for conf in confidence_levels:
        z_score = conf
        lower = pred_rs - z_score * pred_std
        upper = pred_rs + z_score * pred_std
        coverage = np.mean((true_rs >= lower) & (true_rs <= upper))
        observed_coverage.append(coverage)
    return confidence_levels, np.array(observed_coverage)


def compute_uncertainty_quality_metrics(true_rs, pred_rs, pred_std):
    z_scores = np.abs(true_rs - pred_rs) / (pred_std + 1e-10)
    within_1std = np.mean(z_scores < 1.0)
    within_2std = np.mean(z_scores < 2.0)
    within_3std = np.mean(z_scores < 3.0)
    sharpness = pred_std.mean()
    nll = 0.5 * np.mean(np.log(2 * np.pi * pred_std ** 2) + (true_rs - pred_rs) ** 2 / pred_std ** 2)
    return {
        'within_1std': float(within_1std),
        'within_2std': float(within_2std),
        'within_3std': float(within_3std),
        'sharpness': float(sharpness),
        'nll': float(nll)
    }