"""Speech Artifact Detectors — 10 CNN classifiers for detecting audio processing artifacts.

Three architecture families, 10 specialized models:

**Comb/TTS Artifact Detectors** (trained on real vs TTS-predicted audio + comb-filtered augmentations):
  - stft_classifier: Multi-resolution STFT CNN (~1.69M params, 99.88% val acc)
  - waveform_1d: Raw waveform 1D CNN (~1.90M params, 97.44% val acc)
  - mel_classifier: Mel spectrogram 2D CNN (~1.01M params, 99.75% val acc)

**Augmentation Artifact Detectors** (each trained on clean vs specifically-corrupted audio):
  - spectral_denoising: Detects spectral gating denoiser artifacts (99.83% val acc)
  - pitch_correction: Detects autotune/pitch correction artifacts (100% val acc)
  - codec_compression: Detects MP3 codec compression artifacts (100% val acc)
  - comb_filtering: Detects comb filter resonance/notches (100% val acc)
  - phase_vocoder: Detects phase vocoder time-stretch smearing (100% val acc)
  - bandwidth_limitation: Detects missing high frequencies from lowpass filtering (100% val acc)
  - clipping_distortion: Detects hard/soft clipping and harmonic distortion (98.33% val acc)

All models operate on 16 kHz mono waveforms and output an artifact score in [0, 1]:
  0 = clean,  1 = artifact detected.

Usage:
    from speech_artifact_detector import load_model, load_all_models, score_file

    # Single model
    model = load_model("stft_classifier.pt")
    score = score_file(model, "audio.wav")
    print(f"Artifact score: {score:.4f}")

    # All models
    models = load_all_models(".")
    scores = score_file_all(models, "audio.wav")
    for name, s in scores.items():
        print(f"  {name}: {s:.4f}")

    # Batch scoring
    scores = score_batch(models, ["audio1.wav", "audio2.wav", "audio3.wav"])
"""

import os
from typing import Dict, List, Optional, Union

import numpy as np
import soundfile as sf
import torch
import torch.nn as nn
import torchaudio

TARGET_SR = 16000
MAX_SAMPLES = 10 * TARGET_SR  # 10 seconds


# ═══════════════════════════════════════════════════════════════════════════════
# Architecture 1: Multi-Resolution STFT Classifier (~1.69M params)
# Used by: stft_classifier, all 7 augmentation detectors
# ═══════════════════════════════════════════════════════════════════════════════

class STFTResolutionBlock(nn.Module):
    """Single-resolution STFT -> 2D CNN -> pooled features."""

    def __init__(self, n_fft, hop_length, out_dim=128):
        super().__init__()
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.convs = nn.Sequential(
            nn.Conv2d(1, 32, (5, 5), (2, 2), (2, 2)),
            nn.BatchNorm2d(32), nn.GELU(),
            nn.Conv2d(32, 64, (3, 3), (2, 2), (1, 1)),
            nn.BatchNorm2d(64), nn.GELU(),
            nn.Conv2d(64, 128, (3, 3), (2, 2), (1, 1)),
            nn.BatchNorm2d(128), nn.GELU(),
            nn.Conv2d(128, out_dim, (3, 3), (2, 2), (1, 1)),
            nn.BatchNorm2d(out_dim), nn.GELU(),
            nn.Conv2d(out_dim, out_dim, (3, 3), (1, 1), (1, 1)),
            nn.BatchNorm2d(out_dim), nn.GELU(),
        )
        self.pool = nn.AdaptiveAvgPool2d(1)

    def forward(self, waveform):
        # waveform: [B, T]
        window = torch.hann_window(self.n_fft, device=waveform.device, dtype=waveform.dtype)
        stft = torch.stft(
            waveform, self.n_fft, self.hop_length,
            window=window, return_complex=True,
        )
        mag = torch.log(stft.abs().clamp(min=1e-7))  # [B, F, T']
        mag = mag.unsqueeze(1)  # [B, 1, F, T']
        h = self.convs(mag)
        h = self.pool(h).flatten(1)  # [B, out_dim]
        return h


class MultiResSTFTClassifier(nn.Module):
    """Multi-resolution STFT classifier (1,688,833 parameters).

    Computes STFT at 4 resolutions (256/512/1024/2048 FFT), runs each
    through a 5-layer 2D CNN, concatenates the 128-dim features from
    each resolution, and classifies through an MLP head.
    """

    def __init__(self):
        super().__init__()
        self.blocks = nn.ModuleList([
            STFTResolutionBlock(n_fft=256, hop_length=64),
            STFTResolutionBlock(n_fft=512, hop_length=128),
            STFTResolutionBlock(n_fft=1024, hop_length=256),
            STFTResolutionBlock(n_fft=2048, hop_length=512),
        ])
        self.head = nn.Sequential(
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(256, 1),
        )

    def forward(self, waveform):
        """waveform: [B, 1, T] -> logits [B, 1]"""
        x = waveform.squeeze(1)  # [B, T]
        feats = [block(x) for block in self.blocks]
        h = torch.cat(feats, dim=1)  # [B, 512]
        return self.head(h)

    def artifact_score(self, waveform):
        """Returns artifact probability in [0, 1]."""
        return torch.sigmoid(self.forward(waveform))


# ═══════════════════════════════════════════════════════════════════════════════
# Architecture 2: Waveform 1D CNN (~1.90M params)
# Used by: waveform_1d
# ═══════════════════════════════════════════════════════════════════════════════

class Waveform1DCNN(nn.Module):
    """Direct 1D convolution classifier on raw waveform (1,898,753 parameters).

    6 conv blocks with increasing channels [1->64->128->256->256->512->512],
    BatchNorm, GELU, and MaxPool. Global average pooling then MLP head.
    """

    def __init__(self):
        super().__init__()
        channels = [1, 64, 128, 256, 256, 512, 512]
        kernels = [15, 11, 7, 5, 3, 3]
        strides = [4, 2, 2, 2, 2, 1]

        blocks = []
        for i in range(6):
            blocks.extend([
                nn.Conv1d(channels[i], channels[i + 1], kernels[i],
                          stride=strides[i], padding=kernels[i] // 2),
                nn.BatchNorm1d(channels[i + 1]),
                nn.GELU(),
            ])
            if strides[i] > 1:
                blocks.append(nn.MaxPool1d(2))
        self.features = nn.Sequential(*blocks)
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(512, 128),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(128, 1),
        )

    def forward(self, waveform):
        """waveform: [B, 1, T] -> logits [B, 1]"""
        h = self.features(waveform)
        return self.head(h)

    def artifact_score(self, waveform):
        """Returns artifact probability in [0, 1]."""
        return torch.sigmoid(self.forward(waveform))


# ═══════════════════════════════════════════════════════════════════════════════
# Architecture 3: Mel Spectrogram CNN (~1.01M params)
# Used by: mel_classifier
# ═══════════════════════════════════════════════════════════════════════════════

class MelCNNClassifier(nn.Module):
    """Mel spectrogram -> 2D CNN classifier (1,012,417 parameters).

    Computes 80-band mel spectrogram from raw waveform (differentiable),
    then runs through 5-layer 2D CNN with BatchNorm and GELU.
    """

    def __init__(self, sr=TARGET_SR, n_fft=1024, hop_length=256, n_mels=80):
        super().__init__()
        self.mel_spec = torchaudio.transforms.MelSpectrogram(
            sample_rate=sr, n_fft=n_fft, hop_length=hop_length,
            n_mels=n_mels, f_min=0, f_max=sr // 2, power=2.0,
        )
        self.convs = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.BatchNorm2d(32), nn.GELU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64), nn.GELU(), nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128), nn.GELU(), nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256), nn.GELU(), nn.MaxPool2d(2),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256), nn.GELU(),
        )
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256, 128),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(128, 1),
        )

    def forward(self, waveform):
        """waveform: [B, 1, T] -> logits [B, 1]"""
        x = waveform.squeeze(1)  # [B, T]
        mel = self.mel_spec(x)  # [B, n_mels, T']
        mel = torch.log(mel.clamp(min=1e-7))
        mel = mel.unsqueeze(1)  # [B, 1, n_mels, T']
        h = self.convs(mel)
        return self.head(h)

    def artifact_score(self, waveform):
        """Returns artifact probability in [0, 1]."""
        return torch.sigmoid(self.forward(waveform))


# ═══════════════════════════════════════════════════════════════════════════════
# Model registry and metadata
# ═══════════════════════════════════════════════════════════════════════════════

# Maps architecture name -> class
ARCHITECTURE_REGISTRY = {
    "MultiResSTFTClassifier": MultiResSTFTClassifier,
    "stft_classifier": MultiResSTFTClassifier,
    "waveform_1d": Waveform1DCNN,
    "mel_classifier": MelCNNClassifier,
}

# All 10 models with their checkpoint filenames and descriptions
MODEL_CATALOG = {
    # --- Comb/TTS artifact detectors (3) ---
    "stft_classifier": {
        "file": "stft_classifier.pt",
        "architecture": "stft_classifier",
        "category": "tts_artifact",
        "description": "Multi-resolution STFT classifier for TTS vocoder artifacts",
        "detects": "Vocoder synthesis artifacts, comb filtering, metallic resonance",
        "params": "1.69M",
        "val_acc": 99.88,
    },
    "waveform_1d": {
        "file": "waveform_1d.pt",
        "architecture": "waveform_1d",
        "category": "tts_artifact",
        "description": "Raw waveform 1D CNN for TTS vocoder artifacts",
        "detects": "Vocoder synthesis artifacts, waveform-level anomalies",
        "params": "1.90M",
        "val_acc": 97.44,
    },
    "mel_classifier": {
        "file": "mel_classifier.pt",
        "architecture": "mel_classifier",
        "category": "tts_artifact",
        "description": "Mel spectrogram 2D CNN for TTS vocoder artifacts",
        "detects": "Vocoder mel-domain artifacts, spectral smoothing",
        "params": "1.01M",
        "val_acc": 99.75,
    },
    # --- Augmentation artifact detectors (7) ---
    "spectral_denoising": {
        "file": "spectral_denoising.pt",
        "architecture": "MultiResSTFTClassifier",
        "category": "augmentation",
        "description": "Spectral gating denoiser artifact detector",
        "detects": "Artifacts from spectral gating noise reduction (metallic, underwater sound)",
        "params": "1.69M",
        "val_acc": 99.83,
    },
    "pitch_correction": {
        "file": "pitch_correction.pt",
        "architecture": "MultiResSTFTClassifier",
        "category": "augmentation",
        "description": "Pitch correction / autotune artifact detector",
        "detects": "Autotune stepping artifacts, unnatural pitch quantization to semitone grid",
        "params": "1.69M",
        "val_acc": 100.0,
    },
    "codec_compression": {
        "file": "codec_compression.pt",
        "architecture": "MultiResSTFTClassifier",
        "category": "augmentation",
        "description": "Codec compression artifact detector",
        "detects": "MP3 compression artifacts at low bitrates (8-48 kbps): warbling, pre-echo, bandwidth loss",
        "params": "1.69M",
        "val_acc": 100.0,
    },
    "comb_filtering": {
        "file": "comb_filtering.pt",
        "architecture": "MultiResSTFTClassifier",
        "category": "augmentation",
        "description": "Comb filter artifact detector",
        "detects": "Comb filter resonance peaks and notches from delay+feedback mixing",
        "params": "1.69M",
        "val_acc": 100.0,
    },
    "phase_vocoder": {
        "file": "phase_vocoder.pt",
        "architecture": "MultiResSTFTClassifier",
        "category": "augmentation",
        "description": "Phase vocoder time-stretch artifact detector",
        "detects": "Phase smearing and transient blurring from STFT-based time-stretching",
        "params": "1.69M",
        "val_acc": 100.0,
    },
    "bandwidth_limitation": {
        "file": "bandwidth_limitation.pt",
        "architecture": "MultiResSTFTClassifier",
        "category": "augmentation",
        "description": "Bandwidth limitation artifact detector",
        "detects": "Missing high-frequency content from lowpass filtering (cutoff 2-6 kHz)",
        "params": "1.69M",
        "val_acc": 100.0,
    },
    "clipping_distortion": {
        "file": "clipping_distortion.pt",
        "architecture": "MultiResSTFTClassifier",
        "category": "augmentation",
        "description": "Clipping/distortion artifact detector",
        "detects": "Hard/soft clipping and harmonic distortion from overdrive",
        "params": "1.69M",
        "val_acc": 98.33,
    },
}


# ═══════════════════════════════════════════════════════════════════════════════
# Loading utilities
# ═══════════════════════════════════════════════════════════════════════════════

def load_model(
    checkpoint_path: str,
    device: Optional[str] = None,
) -> nn.Module:
    """Load a single artifact detector from a checkpoint file.

    Args:
        checkpoint_path: Path to .pt checkpoint file.
        device: "cuda", "cpu", or None for auto-detect.

    Returns:
        Model in eval mode on the specified device.
    """
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(device)

    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=True)
    arch = ckpt["architecture"]
    cls = ARCHITECTURE_REGISTRY[arch]
    model = cls()
    model.load_state_dict(ckpt["model_state_dict"])
    model.to(device).eval()
    return model


def load_all_models(
    model_dir: str,
    device: Optional[str] = None,
    categories: Optional[List[str]] = None,
) -> Dict[str, nn.Module]:
    """Load all available artifact detectors from a directory.

    Args:
        model_dir: Directory containing .pt checkpoint files.
        device: "cuda", "cpu", or None for auto-detect.
        categories: Optional filter. List of categories to load:
            "tts_artifact", "augmentation", or both. None = load all.

    Returns:
        Dict mapping model name to loaded model.
    """
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    models = {}
    for name, info in MODEL_CATALOG.items():
        if categories and info["category"] not in categories:
            continue
        path = os.path.join(model_dir, info["file"])
        if os.path.exists(path):
            models[name] = load_model(path, device)
    return models


def load_from_hub(
    repo_id: str = "laion/speech-artifact-detectors",
    device: Optional[str] = None,
    categories: Optional[List[str]] = None,
) -> Dict[str, nn.Module]:
    """Load all models directly from HuggingFace Hub.

    Args:
        repo_id: HuggingFace model repository ID.
        device: "cuda", "cpu", or None for auto-detect.
        categories: Optional filter. List of categories to load.

    Returns:
        Dict mapping model name to loaded model.
    """
    from huggingface_hub import hf_hub_download

    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    models = {}
    for name, info in MODEL_CATALOG.items():
        if categories and info["category"] not in categories:
            continue
        try:
            path = hf_hub_download(repo_id, info["file"])
            models[name] = load_model(path, device)
        except Exception as e:
            print(f"Warning: Could not load {name}: {e}")
    return models


# ═══════════════════════════════════════════════════════════════════════════════
# Audio loading
# ═══════════════════════════════════════════════════════════════════════════════

def load_audio(path: str, target_sr: int = TARGET_SR) -> torch.Tensor:
    """Load an audio file and convert to 16 kHz mono float32 tensor.

    Returns:
        Tensor of shape [T] (1-D).
    """
    audio_np, sr = sf.read(path, dtype="float32")
    if audio_np.ndim > 1:
        audio_np = audio_np[:, 0]  # take first channel
    wav = torch.from_numpy(audio_np).float()
    if sr != target_sr:
        wav = torchaudio.functional.resample(wav.unsqueeze(0), sr, target_sr).squeeze(0)
    return wav


def prepare_waveform(
    waveform: torch.Tensor,
    max_samples: int = MAX_SAMPLES,
) -> torch.Tensor:
    """Prepare a 1-D waveform tensor for model input.

    Truncates to max_samples and pads if shorter. Returns [1, 1, T].
    """
    wav = waveform[:max_samples]
    if wav.shape[0] < max_samples:
        wav = torch.cat([wav, torch.zeros(max_samples - wav.shape[0])])
    return wav.unsqueeze(0).unsqueeze(0)  # [1, 1, T]


# ═══════════════════════════════════════════════════════════════════════════════
# Inference utilities
# ═══════════════════════════════════════════════════════════════════════════════

def score_waveform(
    model: nn.Module,
    waveform: torch.Tensor,
    device: Optional[str] = None,
    max_samples: int = MAX_SAMPLES,
) -> float:
    """Score a waveform tensor with a single model.

    Args:
        model: Loaded artifact detector model.
        waveform: 1-D tensor [T] at 16 kHz.
        device: Device to run inference on.
        max_samples: Truncate waveform to this length.

    Returns:
        Artifact score in [0, 1].
    """
    if device is None:
        device = next(model.parameters()).device
    else:
        device = torch.device(device)

    x = prepare_waveform(waveform, max_samples).to(device)
    with torch.no_grad():
        return model.artifact_score(x).item()


def score_waveform_batch(
    model: nn.Module,
    waveforms: List[torch.Tensor],
    device: Optional[str] = None,
    max_samples: int = MAX_SAMPLES,
    batch_size: int = 16,
) -> List[float]:
    """Score multiple waveforms with a single model.

    Args:
        model: Loaded artifact detector model.
        waveforms: List of 1-D tensors [T] at 16 kHz.
        device: Device to run inference on.
        max_samples: Truncate waveforms to this length.
        batch_size: Batch size for inference.

    Returns:
        List of artifact scores in [0, 1].
    """
    if device is None:
        device = next(model.parameters()).device
    else:
        device = torch.device(device)

    all_scores = []
    for i in range(0, len(waveforms), batch_size):
        batch_wavs = waveforms[i:i + batch_size]
        # Prepare batch
        prepared = []
        for wav in batch_wavs:
            w = wav[:max_samples]
            if w.shape[0] < max_samples:
                w = torch.cat([w, torch.zeros(max_samples - w.shape[0])])
            prepared.append(w)
        x = torch.stack(prepared).unsqueeze(1).to(device)  # [B, 1, T]
        with torch.no_grad():
            scores = model.artifact_score(x).squeeze(1).tolist()
        all_scores.extend(scores)
    return all_scores


def score_file(
    model: nn.Module,
    audio_path: str,
    device: Optional[str] = None,
) -> float:
    """Score an audio file with a single model.

    Args:
        model: Loaded artifact detector model.
        audio_path: Path to audio file (any format supported by soundfile).
        device: Device to run inference on.

    Returns:
        Artifact score in [0, 1].
    """
    wav = load_audio(audio_path)
    return score_waveform(model, wav, device)


def score_file_all(
    models: Dict[str, nn.Module],
    audio_path: str,
    device: Optional[str] = None,
) -> Dict[str, float]:
    """Score an audio file with all loaded models.

    Args:
        models: Dict of {name: model} from load_all_models().
        audio_path: Path to audio file.
        device: Device to run inference on.

    Returns:
        Dict mapping model name to artifact score in [0, 1].
    """
    wav = load_audio(audio_path)
    return {name: score_waveform(m, wav, device) for name, m in models.items()}


def score_batch(
    models: Dict[str, nn.Module],
    audio_paths: List[str],
    device: Optional[str] = None,
    batch_size: int = 16,
) -> List[Dict[str, float]]:
    """Score multiple audio files with all loaded models.

    Args:
        models: Dict of {name: model} from load_all_models().
        audio_paths: List of paths to audio files.
        device: Device to run inference on.
        batch_size: Batch size for inference.

    Returns:
        List of dicts, each mapping model name to artifact score.
    """
    # Load all audio
    waveforms = [load_audio(p) for p in audio_paths]

    # Score with each model
    results = [{} for _ in audio_paths]
    for name, model in models.items():
        scores = score_waveform_batch(model, waveforms, device, batch_size=batch_size)
        for i, s in enumerate(scores):
            results[i][name] = s
    return results


# ═══════════════════════════════════════════════════════════════════════════════
# CLI
# ═══════════════════════════════════════════════════════════════════════════════

def main():
    import argparse
    import glob

    p = argparse.ArgumentParser(
        description="Score audio files for speech artifacts using 10 CNN detectors",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""Examples:
  # Score a single file with all models (from local directory)
  python speech_artifact_detector.py --model-dir . --input audio.wav

  # Score using models from HuggingFace Hub
  python speech_artifact_detector.py --hub --input audio.wav

  # Score only augmentation detectors
  python speech_artifact_detector.py --hub --category augmentation --input audio.wav

  # Score only TTS artifact detectors
  python speech_artifact_detector.py --hub --category tts_artifact --input audio.wav

  # Score a directory of files
  python speech_artifact_detector.py --model-dir . --input /path/to/wavs/ --ext wav

  # Score with a specific model
  python speech_artifact_detector.py --checkpoint stft_classifier.pt --input audio.wav
""",
    )
    p.add_argument("--model-dir", help="Directory with .pt checkpoints")
    p.add_argument("--hub", action="store_true",
                   help="Load models from HuggingFace Hub (laion/speech-artifact-detectors)")
    p.add_argument("--checkpoint", help="Path to a single checkpoint")
    p.add_argument("--input", required=True, help="Audio file or directory to score")
    p.add_argument("--ext", default="wav", help="File extension when --input is a directory")
    p.add_argument("--device", default=None, help="Device (cuda/cpu, auto-detect if omitted)")
    p.add_argument("--category", choices=["tts_artifact", "augmentation"],
                   help="Load only models from this category")
    p.add_argument("--threshold", type=float, default=0.5,
                   help="Threshold for artifact classification (default: 0.5)")
    p.add_argument("--batch-size", type=int, default=16,
                   help="Batch size for inference (default: 16)")
    args = p.parse_args()

    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    categories = [args.category] if args.category else None

    # Load model(s)
    if args.checkpoint:
        model = load_model(args.checkpoint, device)
        ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
        name = ckpt.get("augmentation", ckpt.get("architecture", "unknown"))
        models = {name: model}
        print(f"Loaded {name} (val_acc={ckpt.get('val_acc', '?')})")
    elif args.hub:
        models = load_from_hub(device=device, categories=categories)
        print(f"Loaded {len(models)} models from Hub: {list(models.keys())}")
    elif args.model_dir:
        models = load_all_models(args.model_dir, device, categories=categories)
        print(f"Loaded {len(models)} models: {list(models.keys())}")
    else:
        p.error("Provide --model-dir, --hub, or --checkpoint")

    # Collect files
    if os.path.isdir(args.input):
        files = sorted(glob.glob(os.path.join(args.input, f"*.{args.ext}")))
    else:
        files = [args.input]
    print(f"Scoring {len(files)} file(s)...\n")

    model_names = list(models.keys())
    header = f"{'File':<40s} " + " ".join(f"{n:>20s}" for n in model_names) + f" {'Mean':>8s}  {'Verdict'}"
    print(header)
    print("-" * len(header))

    n_artifact = 0
    for fpath in files:
        scores = score_file_all(models, fpath, device)
        mean_score = np.mean(list(scores.values()))
        verdict = "ARTIFACT" if mean_score >= args.threshold else "clean"
        if verdict == "ARTIFACT":
            n_artifact += 1

        fname = os.path.basename(fpath)
        row = f"{fname:<40s} " + " ".join(f"{scores[n]:>20.4f}" for n in model_names) + f" {mean_score:>8.4f}  {verdict}"
        print(row)

    print(f"\n{n_artifact}/{len(files)} classified as artifact (threshold={args.threshold})")


if __name__ == "__main__":
    main()