import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

CLASS_NAMES = [
    'Healthy (0)',
    'Mild NPDR (1)',
    'Moderate NPDR (2)',
    'Severe NPDR (3)',
    'Proliferative DR (4)',
]

NUM_CLASSES = 5


def set_seed(seed: int = 42) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = True


def compute_class_weights(labels) -> torch.FloatTensor:
    """Inverse-frequency weighting computed from training labels."""
    labels = np.array(labels)
    class_counts = np.bincount(labels, minlength=NUM_CLASSES)
    total = len(labels)
    weights = total / (NUM_CLASSES * class_counts.astype(float))
    return torch.FloatTensor(weights)


def get_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device('cuda')
    return torch.device('cpu')


class FocalLoss(nn.Module):
    """Focal loss with optional per-class alpha weighting.

    Focuses learning on hard, misclassified examples by down-weighting
    easy examples (high pt). Good for class imbalance in medical imaging.

    Args:
        alpha: Per-class weight tensor (same shape as class weights for CE).
               If None, no per-class weighting is applied.
        gamma: Focusing parameter. gamma=0 reduces to standard CE.
               gamma=2 is the standard value from the RetinaNet paper.
    """

    def __init__(self, alpha=None, gamma: float = 2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        ce_loss = F.cross_entropy(inputs, targets, weight=self.alpha, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = ((1.0 - pt) ** self.gamma) * ce_loss
        return focal_loss.mean()