#!/usr/bin/env python3
"""Shared mobility-classification utilities used across Task 2 helpers.

This module provides the lightweight LWM classifier head plus supporting
sampling and normalization helpers that were previously bundled inside the
stand-alone mobility fine-tuning scripts.  They remain available so that
benchmarking, router training, and visualisation pipelines can reuse the same
logic without depending on a separate CLI.
"""

from __future__ import annotations

import glob
import json
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence, Tuple

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

from pretraining.pretrained_model import lwm as lwm_model
from task1.train_mcs_models import (
    _extract_metadata,
    identify_modulation,
    load_all_samples,
)

MOBILITY_LABELS = ["static", "pedestrian", "vehicular"]
BINARY_MOBILITY_LABELS = ["vehicular", "pedestrian"]


def load_dataset_stats(models_root: Path) -> Dict[str, float | str]:
    """Load dataset statistics (mean/std/normalization mode) from a models directory."""
    stats_path = models_root / "dataset_stats.json"
    if not stats_path.exists():
        print(
            f"[WARN] dataset_stats.json not found under {models_root}; "
            "falling back to per-sample normalization with mean=0/std=1.",
            flush=True,
        )
        return {"mean": 0.0, "std": 1.0, "normalization": "per_sample"}
    with open(stats_path, "r", encoding="utf-8") as f:
        stats = json.load(f)
    mean = float(stats.get("mean", 0.0))
    std = float(stats.get("std", 1.0))
    if std == 0.0:
        std = 1.0
    normalization = str(stats.get("normalization", stats.get("mode", "dataset")))
    return {
        "mean": mean,
        "std": std,
        "normalization": normalization,
    }


def gather_controlled_groups(
    data_root: Path,
    cities: Sequence[str],
    comm: str,
    mobilities: Sequence[str],
    snrs: Sequence[str] | None,
    fft_whitelist: Sequence[str] | None,
) -> Dict[Tuple[str, str, str, str, str], Dict[str, List[str]]]:
    """Group spectrogram paths by (city, modulation, rate, SNR, FFT) while balancing mobilities."""
    groups: Dict[Tuple[str, str, str, str, str], Dict[str, List[str]]] = defaultdict(lambda: defaultdict(list))
    mobility_set = set(mobilities)
    snr_set = set(snrs) if snrs else None
    fft_set = set(fft_whitelist) if fft_whitelist else None

    for city in cities:
        base = data_root / city / comm
        if not base.exists():
            continue
        pattern = str(base / "**" / "spectrograms" / "*.pkl")
        for path_str in glob.iglob(pattern, recursive=True):
            path = Path(path_str)
            rate, snr, mobility = _extract_metadata(path.parts)
            if mobility not in mobility_set:
                continue
            if snr_set is not None and snr not in snr_set:
                continue
            fft = next((part for part in path.parts if part.startswith("win")), "fft_unknown")
            if fft_set is not None and fft not in fft_set:
                continue
            _, modulation = identify_modulation(path_str)
            if modulation is None:
                continue
            key = (city, modulation, rate, snr, fft)
            groups[key][mobility].append(str(path))
    return {key: dict(mob_map) for key, mob_map in groups.items()}


def _collect_balanced_arrays(
    groups: Dict[Tuple[str, str, str, str, str], Dict[str, List[str]]],
    mobilities: Sequence[str],
    max_per_config: int,
    rng: np.random.Generator,
) -> Tuple[np.ndarray, np.ndarray, Dict[str, Any]]:
    """Load spectrogram arrays with per-configuration balance across mobilities."""
    features: List[np.ndarray] = []
    labels: List[np.ndarray] = []
    mobility_to_idx = {mob: idx for idx, mob in enumerate(mobilities)}
    per_mobility_totals = {mob: 0 for mob in mobilities}
    matched_configs = 0
    preview_configs: List[Tuple[str, str, str, str, str]] = []

    for key, mobility_map in groups.items():
        if not all(mob in mobility_map for mob in mobilities):
            continue

        cached_arrays: Dict[str, np.ndarray] = {}
        per_mobility_counts: List[int] = []
        for mobility in mobilities:
            paths = mobility_map[mobility]
            collected: List[np.ndarray] = []
            for path in paths:
                arr = load_all_samples(path)
                if arr.size == 0:
                    continue
                collected.append(arr)
            if not collected:
                cached_arrays = {}
                break
            stacked = np.concatenate(collected, axis=0)
            cached_arrays[mobility] = stacked
            per_mobility_counts.append(stacked.shape[0])

        if len(cached_arrays) != len(mobilities):
            continue

        limit = min(per_mobility_counts)
        if max_per_config > 0:
            limit = min(limit, max_per_config)
        if limit == 0:
            continue

        for mobility in mobilities:
            arr = cached_arrays[mobility]
            if arr.shape[0] > limit:
                indices = rng.permutation(arr.shape[0])[:limit]
                arr = arr[indices]
            features.append(arr)
            labels.append(np.full(arr.shape[0], mobility_to_idx[mob], dtype=np.int64))
            per_mobility_totals[mobility] += arr.shape[0]

        if matched_configs < 5:
            preview_configs.append(key)
        matched_configs += 1

    if not features:
        return (
            np.empty((0, 128, 128), dtype=np.float32),
            np.empty((0,), dtype=np.int64),
            {"per_mobility": per_mobility_totals, "matched_configs": matched_configs, "preview_configs": preview_configs},
        )

    stacked_features = np.concatenate(features, axis=0).astype(np.float32, copy=False)
    stacked_labels = np.concatenate(labels, axis=0).astype(np.int64, copy=False)
    return stacked_features, stacked_labels, {
        "per_mobility": per_mobility_totals,
        "matched_configs": matched_configs,
        "preview_configs": preview_configs,
    }


class ResidualBlock1D(nn.Module):
    """1D Residual block used by the Res1DCNN classification head."""

    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm1d(out_channels)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm1d(out_channels)
        self.shortcut = nn.Sequential()
        if in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv1d(in_channels, out_channels, kernel_size=1),
                nn.BatchNorm1d(out_channels),
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.bn2(self.conv2(x))
        x += self.shortcut(residual)
        return F.relu(x)


class Res1DCNNHead(nn.Module):
    """Compact ResNet-style 1D head for classifying 128-d embeddings."""

    def __init__(self, input_dim: int, num_classes: int, dropout: float = 0.5) -> None:
        super().__init__()
        hidden_dim = 64
        self.conv1 = nn.Conv1d(1, hidden_dim, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm1d(hidden_dim)
        self.res_block = ResidualBlock1D(hidden_dim, hidden_dim)
        self.fc = nn.Linear(hidden_dim, num_classes)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.unsqueeze(1)
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.res_block(x)
        x = F.adaptive_avg_pool1d(x, 1).squeeze(-1)
        x = self.dropout(x)
        return self.fc(x)


class LWMClassifierMinimal(nn.Module):
    """LWM backbone wrapper with configurable classifier and optional projection head."""

    def __init__(
        self,
        backbone: nn.Module,
        num_classes: int,
        classifier_dim: int,
        dropout: float,
        trainable_layers: int,
        projection_dim: int,
        append_input_stats: bool,
        normalization_stats: Dict[str, object] | None,
        head_type: str = "mlp",
    ) -> None:
        super().__init__()
        self.backbone = backbone
        self.patch_size = 4
        self.unfold = nn.Unfold(kernel_size=self.patch_size, stride=self.patch_size)
        self.head_type = head_type

        self.append_input_stats = bool(append_input_stats)
        stats_info = normalization_stats or {}
        self.normalization_mode = str(stats_info.get("normalization", "dataset")).lower()
        self.dataset_mean = float(stats_info.get("mean", 0.0))
        self.dataset_std = float(stats_info.get("std", 1.0))
        if abs(self.dataset_std) < 1e-6:
            self.dataset_std = 1e-6
        base_dim = 128
        stats_dim = 2 if self.append_input_stats else 0
        input_dim = base_dim + stats_dim

        classifier_dim = max(32, int(classifier_dim))
        dropout = max(0.0, float(dropout))

        if head_type == "linear":
            self.classifier = nn.Sequential(
                nn.LayerNorm(input_dim),
                nn.Linear(input_dim, num_classes),
            )
        elif head_type == "res1dcnn":
            self.classifier = nn.Sequential(
                nn.LayerNorm(input_dim),
                Res1DCNNHead(input_dim, num_classes, dropout=dropout),
            )
        else:
            head_layers: List[nn.Module] = [
                nn.LayerNorm(input_dim),
                nn.Linear(input_dim, classifier_dim),
                nn.GELU(),
            ]
            if dropout > 0:
                head_layers.append(nn.Dropout(dropout))
            head_layers.append(nn.Linear(classifier_dim, num_classes))
            self.classifier = nn.Sequential(*head_layers)

        proj_dim = int(projection_dim)
        if proj_dim > 0:
            self.projection_head = nn.Sequential(
                nn.Linear(128, proj_dim),
                nn.ReLU(inplace=True),
                nn.Linear(proj_dim, proj_dim),
            )
        else:
            self.projection_head = None

        for param in self.backbone.parameters():
            param.requires_grad = False

        if trainable_layers > 0:
            layers = getattr(self.backbone, "layers", None)
            if layers is not None:
                trainable_layers = min(trainable_layers, len(layers))
                for layer in layers[-trainable_layers:]:
                    for param in layer.parameters():
                        param.requires_grad = True

    def spectrogram_to_tokens(self, x: torch.Tensor) -> torch.Tensor:
        x = x.unsqueeze(1)
        patches = self.unfold(x).transpose(1, 2)
        cls_token = torch.full(
            (patches.size(0), 1, patches.size(-1)),
            0.2,
            dtype=patches.dtype,
            device=patches.device,
        )
        return torch.cat([cls_token, patches], dim=1)

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        tokens = self.spectrogram_to_tokens(x)
        outputs = self.backbone(tokens)
        if outputs.size(1) <= 1:
            return outputs[:, 0, :]
        return outputs[:, 1:, :].mean(dim=1)

    def _collect_input_stats(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=(1, 2))
        std = x.std(dim=(1, 2), unbiased=False)
        if self.normalization_mode == "dataset":
            mean = mean * self.dataset_std + self.dataset_mean
            std = std * self.dataset_std
        return torch.stack([mean, std], dim=1)

    def forward(
        self,
        x: torch.Tensor,
        *,
        input_stats: torch.Tensor | None = None,
        return_projection: bool = False,
    ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
        features = self.forward_features(x)
        classifier_input = features
        if self.append_input_stats:
            stats = input_stats if input_stats is not None else self._collect_input_stats(x)
            if stats.dtype != classifier_input.dtype:
                stats = stats.to(classifier_input.dtype)
            stats = stats.to(classifier_input.device)
            classifier_input = torch.cat([classifier_input, stats], dim=1)
        logits = self.classifier(classifier_input)
        if return_projection:
            projection = self.projection_head(features) if self.projection_head is not None else None
            return logits, projection
        return logits


def prepare_model(
    checkpoint: Path,
    num_classes: int,
    classifier_dim: int,
    dropout: float,
    trainable_layers: int,
    projection_dim: int,
    *,
    append_input_stats: bool = False,
    normalization_stats: Dict[str, object] | None = None,
    head_type: str = "mlp",
) -> nn.Module:
    """Instantiate an LWM backbone with the minimal classifier head."""
    backbone = lwm_model(element_length=16, d_model=128, n_layers=12, max_len=1025, n_heads=8, dropout=0.1)
    state = torch.load(checkpoint, map_location="cpu")
    if any(k.startswith("module.") for k in state):
        state = {k.replace("module.", ""): v for k, v in state.items()}
    backbone.load_state_dict(state, strict=False)
    return LWMClassifierMinimal(
        backbone,
        num_classes=num_classes,
        classifier_dim=classifier_dim,
        dropout=dropout,
        trainable_layers=trainable_layers,
        projection_dim=projection_dim,
        append_input_stats=append_input_stats,
        normalization_stats=normalization_stats,
        head_type=head_type,
    )


def supervised_contrastive_loss(
    features: torch.Tensor,
    labels: torch.Tensor,
    temperature: float,
) -> torch.Tensor:
    """Supervised contrastive loss over a batch of feature embeddings."""
    batch_size = features.size(0)
    if batch_size < 2:
        return features.new_tensor(0.0)

    features = F.normalize(features, dim=1)
    similarity = torch.div(torch.matmul(features, features.T), max(temperature, 1e-6))
    logits_max, _ = similarity.max(dim=1, keepdim=True)
    similarity = similarity - logits_max.detach()

    device = features.device
    labels = labels.contiguous().view(-1, 1)
    mask = torch.eq(labels, labels.T).float().to(device)
    logits_mask = torch.ones_like(mask) - torch.eye(batch_size, device=device)
    mask = mask * logits_mask

    exp_logits = torch.exp(similarity) * logits_mask
    log_prob = similarity - torch.log(exp_logits.sum(dim=1, keepdim=True) + 1e-12)

    mask_sum = mask.sum(dim=1)
    valid = mask_sum > 0
    if not torch.any(valid):
        return features.new_tensor(0.0)

    mean_log_prob_pos = (mask * log_prob).sum(dim=1) / mask_sum.clamp_min(1e-12)
    loss = -mean_log_prob_pos[valid].mean()
    return loss


__all__ = [
    "BINARY_MOBILITY_LABELS",
    "LWMClassifierMinimal",
    "MOBILITY_LABELS",
    "Res1DCNNHead",
    "_collect_balanced_arrays",
    "gather_controlled_groups",
    "load_dataset_stats",
    "prepare_model",
    "supervised_contrastive_loss",
]