"""
Audio classifier for Electrical Outlets. Expects spectrogram or waveform; outputs class logits.
Severity from label_mapping. Small CNN for 100-sample regime.
"""
from pathlib import Path
from typing import Dict, Any, Optional
import json
import torch
import torch.nn as nn


class SpectrogramCNN(nn.Module):
    """Lightweight CNN on mel spectrogram (n_mels x time)."""

    def __init__(self, n_mels: int = 64, time_steps: int = 128, num_classes: int = 4):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.fc = nn.Linear(128, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.dim() == 2:
            x = x.unsqueeze(0).unsqueeze(0)
        elif x.dim() == 3:
            x = x.unsqueeze(1)
        x = self.conv(x)
        x = x.flatten(1)
        return self.fc(x)


class ElectricalOutletsAudioModel(nn.Module):
    """Wrapper: optional mel transform then SpectrogramCNN. Severity from mapping."""

    def __init__(
        self,
        num_classes: int = 4,
        label_mapping_path: Optional[Path] = None,
        n_mels: int = 64,
        time_steps: int = 128,
    ):
        super().__init__()
        self.num_classes = num_classes
        self.n_mels = n_mels
        self.time_steps = time_steps
        self.backbone = SpectrogramCNN(n_mels=n_mels, time_steps=time_steps, num_classes=num_classes)
        self.idx_to_label = None
        self.idx_to_issue_type = None
        self.idx_to_severity = None
        if label_mapping_path and Path(label_mapping_path).exists():
            with open(label_mapping_path) as f:
                lm = json.load(f)
            self.idx_to_label = lm["audio"]["idx_to_label"]
            self.idx_to_issue_type = [lm["audio"]["label_to_issue_type"].get(lbl, "normal") for lbl in lm["audio"]["idx_to_label"]]
            self.idx_to_severity = [lm["audio"]["label_to_severity"].get(lm["audio"]["idx_to_label"][i], "medium") for i in range(num_classes)]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.backbone(x)

    def predict_to_schema(self, logits: torch.Tensor) -> Dict[str, Any]:
        probs = torch.softmax(logits, dim=-1)
        if logits.dim() == 1:
            probs = probs.unsqueeze(0)
        conf, pred = probs.max(dim=-1)
        pred = pred.item() if pred.numel() == 1 else pred
        conf = conf.item() if conf.numel() == 1 else conf
        issue_type = (self.idx_to_issue_type or ["normal"] * self.num_classes)[pred]
        severity = (self.idx_to_severity or ["medium"] * self.num_classes)[pred]
        result = "normal" if issue_type == "normal" else "issue_detected"
        return {
            "result": result,
            "issue_type": issue_type,
            "severity": severity,
            "confidence": float(conf),
            "class_idx": int(pred),
        }