"""Frozen HuBERT features and a four-label sound-event head.

This path deliberately has no PANNs runtime dependency.  It shares the pinned
HuBERT encoder used for audio/context decisions, while keeping sound labels in
the fixed ``sound-v1`` taxonomy.
"""
from __future__ import annotations

from dataclasses import dataclass
import math
from pathlib import Path

import numpy as np
import torch
from torch import nn
from transformers import AutoFeatureExtractor, HubertModel

from .schema import SOUND_LABELS


HUBERT_MODEL = "facebook/hubert-base-ls960"
HUBERT_REVISION = "dba3bb02fda4248b6e082697eee756de8fe8aa8a"
SAMPLE_RATE = 16_000
WINDOW_SECONDS = 5.0
HOP_SECONDS = 2.5
TAIL_SECONDS = 0.5
FRAME_SECONDS = 0.020
FEATURE_DIM = 1_536


@dataclass(frozen=True)
class SoundWindows:
    waveform: torch.Tensor
    intervals: tuple[tuple[float, float], ...]
    observed_samples: tuple[int, ...]


def make_sound_windows(waveform: np.ndarray | torch.Tensor) -> SoundWindows:
    """Create observed 5-s windows with 2.5-s hop, padding only after audio."""
    wave = torch.as_tensor(np.asarray(waveform, dtype=np.float32))
    if wave.ndim != 1 or wave.numel() == 0:
        raise ValueError("waveform must be finite nonempty mono 16-kHz audio")
    window, hop = round(WINDOW_SECONDS * SAMPLE_RATE), round(HOP_SECONDS * SAMPLE_RATE)
    last_start = max(0, wave.numel() - window)
    starts = list(range(0, last_start + 1, hop))
    if not starts or starts[-1] != last_start:
        starts.append(last_start)
    frames, intervals, observed = [], [], []
    for start in starts:
        end = min(start + window, wave.numel())
        frames.append(torch.nn.functional.pad(wave[start:end], (0, window - (end - start))))
        intervals.append((start / SAMPLE_RATE, end / SAMPLE_RATE))
        observed.append(end - start)
    return SoundWindows(torch.stack(frames), tuple(intervals), tuple(observed))


def aggregate_window_scores(scores: torch.Tensor) -> torch.Tensor:
    if scores.ndim != 2 or scores.shape[0] == 0 or scores.shape[1] != len(SOUND_LABELS):
        raise ValueError("expected [nonempty windows, four-label] scores")
    return scores.max(dim=0).values


def aggregate_hubert_hidden(hidden: torch.Tensor, output_lengths: torch.Tensor) -> torch.Tensor:
    """Observed-frame mean plus the observed final 0.5-s mean for every window."""
    if hidden.ndim != 3 or output_lengths.ndim != 1 or hidden.shape[0] != output_lengths.numel():
        raise ValueError("hidden state and HuBERT output lengths are inconsistent")
    tail_frames = math.ceil(TAIL_SECONDS / FRAME_SECONDS)
    values = []
    for row, length in zip(hidden, output_lengths.tolist()):
        observed = min(row.shape[0], max(1, int(length)))
        values.append(torch.cat((row[:observed].mean(0), row[max(0, observed - tail_frames):observed].mean(0))))
    return torch.stack(values)


@torch.inference_mode()
def encode_hubert_windows(encoder: HubertModel, feature_extractor: AutoFeatureExtractor,
                          windows: torch.Tensor, observed_samples: tuple[int, ...] | None = None,
                          batch_size: int = 32) -> torch.Tensor:
    """Return frozen [windows, 1536] features; no padded frames enter pooling."""
    if windows.ndim != 2 or windows.shape[1] != round(WINDOW_SECONDS * SAMPLE_RATE) or batch_size < 1:
        raise ValueError("expected 5-s 16-kHz windows and a positive batch size")
    if observed_samples is None:
        observed_samples = tuple([windows.shape[1]] * windows.shape[0])
    if len(observed_samples) != windows.shape[0] or any(length < 1 or length > windows.shape[1] for length in observed_samples):
        raise ValueError("each window needs one valid observed-sample length")
    device = next(encoder.parameters()).device
    encoded = []
    for start in range(0, windows.shape[0], batch_size):
        batch = windows[start:start + batch_size]
        lengths = observed_samples[start:start + batch_size]
        values = [item[:length].numpy() for item, length in zip(batch.cpu(), lengths)]
        inputs = feature_extractor(values, sampling_rate=SAMPLE_RATE, padding=True,
                                   return_attention_mask=True, return_tensors="pt")
        mask = inputs["attention_mask"]
        hidden = encoder(input_values=inputs["input_values"].to(device), attention_mask=mask.to(device)).last_hidden_state
        lengths = encoder._get_feat_extract_output_lengths(mask.sum(dim=1)).to(device)
        encoded.append(aggregate_hubert_hidden(hidden, lengths).float().cpu())
    return torch.cat(encoded)


class WindowMILHead(nn.Module):
    """Small head shared by cache adapters with any fixed window feature size."""

    def __init__(self, feature_dim: int):
        super().__init__()
        if feature_dim < 1:
            raise ValueError("feature_dim must be positive")
        self.feature_dim = feature_dim
        self.layers = nn.Sequential(
            nn.LayerNorm(feature_dim), nn.Linear(feature_dim, 128), nn.GELU(), nn.Dropout(0.1),
            nn.Linear(128, len(SOUND_LABELS)),
        )

    def forward(self, window_features: torch.Tensor) -> torch.Tensor:
        if window_features.shape[-1] != self.feature_dim:
            raise ValueError(f"expected features ending in {self.feature_dim} dimensions")
        return self.layers(window_features)


class HubertSoundHead(WindowMILHead):
    """The 1536-d shared-HuBERT specialization of ``WindowMILHead``."""

    def __init__(self):
        super().__init__(FEATURE_DIM)


@torch.inference_mode()
def score_window_features(head: WindowMILHead, window_features: torch.Tensor) -> dict[str, torch.Tensor]:
    """Score cached/runtime HuBERT window features and apply documented clip max MIL."""
    if window_features.ndim != 2:
        raise ValueError("window_features must have shape [windows, 1536]")
    device = next(head.parameters()).device
    window_scores = torch.sigmoid(head(window_features.to(device))).cpu()
    return {"window_scores": window_scores, "clip_scores": aggregate_window_scores(window_scores)}


def padded_mil_logits(head: WindowMILHead, entries: list[dict], device: str | torch.device) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply a window head then clipwise max MIL to a cache-adapter batch.

    The stable adapter contract is ``entry['features'] -> float [W, D]`` and
    ``entry['target'] -> four-label float vector``.  Extra provenance such as
    intervals or duration is intentionally preserved but not interpreted by
    this fixed-label trainer.
    """
    if not entries:
        raise ValueError("at least one cache entry is required")
    maximum = max(entry["features"].shape[0] for entry in entries)
    features = torch.zeros((len(entries), maximum, head.feature_dim), device=device)
    valid = torch.zeros((len(entries), maximum), dtype=torch.bool, device=device)
    targets = torch.stack([torch.as_tensor(entry["target"], dtype=torch.float32) for entry in entries]).to(device)
    for index, entry in enumerate(entries):
        value = torch.as_tensor(entry["features"], dtype=torch.float32)
        if value.ndim != 2 or value.shape[0] < 1 or value.shape[1] != head.feature_dim:
            raise ValueError("cache entry features must have shape [nonempty windows, configured feature_dim]")
        features[index, :value.shape[0]] = value.to(device)
        valid[index, :value.shape[0]] = True
    return head(features).masked_fill(~valid.unsqueeze(-1), -torch.inf).amax(dim=1), targets


def load_hubert_sound_head(checkpoint: str | Path, device: str | torch.device = "cpu") -> tuple[HubertSoundHead, torch.Tensor, dict]:
    payload = torch.load(checkpoint, map_location="cpu", weights_only=False)
    if not isinstance(payload, dict) or payload.get("format") != "whisper-jev-unified-hubert-sound-v1":
        raise ValueError("expected a unified HuBERT sound checkpoint")
    if tuple(payload.get("taxonomy", ())) != tuple(SOUND_LABELS):
        raise ValueError("checkpoint taxonomy differs from the fixed sound taxonomy")
    head = HubertSoundHead().to(device).eval()
    head.load_state_dict(payload["head"], strict=True)
    thresholds = torch.as_tensor(payload["thresholds"], dtype=torch.float32)
    return head, thresholds, payload