# utils/adaptive_classifiers.py
import torch
import torch.nn as nn
from typing import Optional

CLASS_NAMES = ['gore', 'hate', 'medical', 'safe', 'sexual']

class SafetyClassifier1280(nn.Module):
    """
    Unified safety classifier for mid-UNet features of shape (B, 1280, H, W).
    Robust to variable HxW via AdaptiveAvgPool2d((8,8)) before the head.
    """
    def __init__(self, num_classes: int = 5):
        super().__init__()
        self.pre = nn.AdaptiveAvgPool2d((8, 8))
        self.net = nn.Sequential(
            nn.Conv2d(1280, 512, 3, padding=1),
            nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.MaxPool2d(2),    # 512 x 4 x 4
            nn.Conv2d(512, 256, 3, padding=1),
            nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(2),    # 256 x 2 x 2
            nn.AdaptiveAvgPool2d(1), nn.Flatten(),                          # 256
            nn.Linear(256, 128), nn.ReLU(inplace=True), nn.Dropout(0.3),
            nn.Linear(128, num_classes)
        )
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(m):
        if isinstance(m, nn.Linear):
            nn.init.xavier_uniform_(m.weight); nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if m.bias is not None: nn.init.zeros_(m.bias)
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.ones_(m.weight); nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.pre(x)  # (B, 1280, 8, 8)
        return self.net(x)

def load_classifier_1280(
    weights_path: str,
    device: Optional[torch.device] = None,
    dtype: torch.dtype = torch.float32
) -> SafetyClassifier1280:
    model = SafetyClassifier1280().to(device or "cpu", dtype=dtype)
    state = torch.load(weights_path, map_location="cpu")
    if isinstance(state, dict) and "model_state_dict" in state:
        state = state["model_state_dict"]
    model.load_state_dict(state, strict=True)
    model.eval()
    return model

def pick_weights_for_pipe(pipe) -> str:
    """
    Optional helper: return a default weights file based on the base SD pipeline id.
    You can also use a single shared file 'classifiers/safety_classifier_1280.pth'.
    """
    name = str(getattr(pipe, "_internal_dict", {}).get("_name_or_path", "")).lower()
    # Adjust logic as you like — default to a single shared file:
    return "classifiers/safety_classifier_1280.pth"