#!/usr/bin/env python3
"""Evaluate FastOobleckDecoder checkpoints against the teacher.

Generates audio from random latents using both teacher and student,
computes SNR, multi-res STFT distance, and mel distance. Saves
example audio clips for listening comparison.

Usage:
    uv run python scripts/eval_fast_decoder.py --ckpt checkpoints/fast_decoder_v3/student_step500000.pt
    uv run python scripts/eval_fast_decoder.py --ckpt checkpoints/fast_decoder_v3/student_step520000.pt
"""

import argparse
import json
import math
import os
import sys
import time
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import weight_norm


# =========================================================================
# Student model (must match training script exactly)
# =========================================================================

class Snake1d(nn.Module):
    def __init__(self, hidden_dim, logscale=True):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.alpha.requires_grad = True
        self.beta.requires_grad = True
        self.logscale = logscale

    def forward(self, hidden_states):
        shape = hidden_states.shape
        alpha = self.alpha if not self.logscale else torch.exp(self.alpha)
        beta = self.beta if not self.logscale else torch.exp(self.beta)
        hidden_states = hidden_states.reshape(shape[0], shape[1], -1)
        hidden_states = hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2)
        hidden_states = hidden_states.reshape(shape)
        return hidden_states


class FastResidualUnit(nn.Module):
    def __init__(self, dim: int, dilation: int = 1):
        super().__init__()
        pad = ((7 - 1) * dilation) // 2
        self.snake1 = Snake1d(dim)
        self.conv1 = weight_norm(nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad))
        self.snake2 = Snake1d(dim)
        self.conv2 = weight_norm(nn.Conv1d(dim, dim, kernel_size=1))

    def forward(self, x):
        h = self.conv1(self.snake1(x))
        h = self.conv2(self.snake2(h))
        pad = (x.shape[-1] - h.shape[-1]) // 2
        if pad > 0:
            x = x[..., pad:-pad]
        return x + h


class FastDecoderBlock(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, stride: int = 1):
        super().__init__()
        self.snake1 = Snake1d(in_dim)
        self.conv_t = weight_norm(nn.ConvTranspose1d(
            in_dim, out_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2),
        ))
        self.res1 = FastResidualUnit(out_dim, dilation=1)
        self.res2 = FastResidualUnit(out_dim, dilation=3)

    def forward(self, x):
        x = self.snake1(x)
        x = self.conv_t(x)
        x = self.res1(x)
        x = self.res2(x)
        return x


class FastOobleckDecoder(nn.Module):
    def __init__(self, channels=128, input_channels=64, audio_channels=2,
                 upsampling_ratios=None, channel_multiples=None):
        super().__init__()
        if upsampling_ratios is None:
            upsampling_ratios = [10, 6, 4, 4, 2]
        if channel_multiples is None:
            channel_multiples = [1, 2, 4, 8, 8]
        strides = upsampling_ratios
        cm = [1] + channel_multiples
        self.conv1 = weight_norm(nn.Conv1d(input_channels, channels * cm[-1], kernel_size=7, padding=3))
        blocks = []
        for i, stride in enumerate(strides):
            in_dim = channels * cm[len(strides) - i]
            out_dim = channels * cm[len(strides) - i - 1]
            blocks.append(FastDecoderBlock(in_dim, out_dim, stride=stride))
        self.blocks = nn.ModuleList(blocks)
        self.final_snake = Snake1d(channels)
        self.conv2 = weight_norm(nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False))

    def forward(self, latents):
        x = self.conv1(latents)
        for block in self.blocks:
            x = block(x)
        x = self.final_snake(x)
        x = self.conv2(x)
        return x


# =========================================================================
# Metrics
# =========================================================================

def compute_snr(ref: torch.Tensor, gen: torch.Tensor) -> float:
    """SNR in dB between reference and generated audio."""
    min_len = min(ref.shape[-1], gen.shape[-1])
    ref = ref[..., :min_len]
    gen = gen[..., :min_len]
    noise = ref - gen
    signal_power = (ref ** 2).mean()
    noise_power = (noise ** 2).mean()
    if noise_power < 1e-10:
        return 100.0
    return 10 * torch.log10(signal_power / noise_power).item()


def compute_hf_energy_ratio(ref: torch.Tensor, gen: torch.Tensor, sr=48000, cutoff_hz=8000) -> dict:
    """Compare high-frequency energy between reference and generated audio.

    Returns dict with:
      - ref_hf_ratio: fraction of teacher energy above cutoff
      - gen_hf_ratio: fraction of student energy above cutoff
      - hf_energy_match: how close student HF energy is to teacher (1.0 = perfect)
      - spectral_rolloff_ref: frequency below which 85% of energy lives (teacher)
      - spectral_rolloff_gen: same for student
    """
    n_fft = 4096
    hop = 1024
    min_len = min(ref.shape[-1], gen.shape[-1])
    ref_flat = ref[..., :min_len].reshape(-1, min_len)
    gen_flat = gen[..., :min_len].reshape(-1, min_len)

    window = torch.hann_window(n_fft, device=ref.device)
    ref_stft = torch.stft(ref_flat, n_fft, hop_length=hop, window=window, return_complex=True, normalized=True)
    gen_stft = torch.stft(gen_flat, n_fft, hop_length=hop, window=window, return_complex=True, normalized=True)

    ref_power = ref_stft.abs().pow(2).mean(dim=(0, 2))  # [freq_bins]
    gen_power = gen_stft.abs().pow(2).mean(dim=(0, 2))

    freq_bins = n_fft // 2 + 1
    freqs = torch.linspace(0, sr / 2, freq_bins, device=ref.device)
    cutoff_bin = int(cutoff_hz * n_fft / sr)

    ref_hf = ref_power[cutoff_bin:].sum().item()
    ref_total = ref_power.sum().item()
    gen_hf = gen_power[cutoff_bin:].sum().item()
    gen_total = gen_power.sum().item()

    ref_hf_ratio = ref_hf / max(ref_total, 1e-10)
    gen_hf_ratio = gen_hf / max(gen_total, 1e-10)
    hf_match = min(gen_hf_ratio, ref_hf_ratio) / max(gen_hf_ratio, ref_hf_ratio, 1e-10)

    # Spectral rolloff (85%)
    ref_cumsum = ref_power.cumsum(0) / max(ref_total, 1e-10)
    gen_cumsum = gen_power.cumsum(0) / max(gen_total, 1e-10)
    ref_rolloff = freqs[(ref_cumsum >= 0.85).nonzero(as_tuple=True)[0][0]].item() if (ref_cumsum >= 0.85).any() else sr / 2
    gen_rolloff = freqs[(gen_cumsum >= 0.85).nonzero(as_tuple=True)[0][0]].item() if (gen_cumsum >= 0.85).any() else sr / 2

    return {
        "ref_hf_ratio": round(ref_hf_ratio, 4),
        "gen_hf_ratio": round(gen_hf_ratio, 4),
        "hf_energy_match": round(hf_match, 4),
        "spectral_rolloff_ref_hz": round(ref_rolloff),
        "spectral_rolloff_gen_hz": round(gen_rolloff),
    }


def compute_stft_distance(ref: torch.Tensor, gen: torch.Tensor) -> float:
    """Average log-magnitude STFT distance across multiple resolutions."""
    eps = 1e-5
    min_len = min(ref.shape[-1], gen.shape[-1])
    ref_flat = ref[..., :min_len].reshape(-1, min_len)
    gen_flat = gen[..., :min_len].reshape(-1, min_len)

    distances = []
    for n_fft in [256, 512, 1024, 2048]:
        hop = n_fft // 4
        window = torch.hann_window(n_fft, device=ref.device)
        ref_stft = torch.stft(ref_flat, n_fft, hop_length=hop, window=window, return_complex=True, normalized=True)
        gen_stft = torch.stft(gen_flat, n_fft, hop_length=hop, window=window, return_complex=True, normalized=True)
        dist = F.l1_loss(torch.log(gen_stft.abs() + eps), torch.log(ref_stft.abs() + eps))
        distances.append(dist.item())
    return sum(distances) / len(distances)


def compute_mel_distance(ref: torch.Tensor, gen: torch.Tensor, sr=48000) -> float:
    """Mel spectrogram L1 distance (1024-point)."""
    eps = 1e-5
    n_fft = 1024
    hop = 256
    n_mels = 80
    min_len = min(ref.shape[-1], gen.shape[-1])
    ref_flat = ref[..., :min_len].reshape(-1, min_len)
    gen_flat = gen[..., :min_len].reshape(-1, min_len)

    window = torch.hann_window(n_fft, device=ref.device)

    def mel_fb(device):
        f_min, f_max = 0.0, sr / 2.0
        freq_bins = n_fft // 2 + 1
        def hz2mel(f): return 2595.0 * math.log10(1.0 + f / 700.0)
        def mel2hz(m): return 700.0 * (10.0 ** (m / 2595.0) - 1.0)
        mel_points = torch.linspace(hz2mel(f_min), hz2mel(f_max), n_mels + 2, device=device)
        hz_points = mel2hz(mel_points)
        bins = (hz_points * n_fft / sr).long().clamp(0, freq_bins - 1)
        fb = torch.zeros(n_mels, freq_bins, device=device)
        for i in range(n_mels):
            l, c, r = bins[i], bins[i+1], bins[i+2]
            if c > l: fb[i, l:c] = torch.linspace(0, 1, c - l, device=device)
            if r > c: fb[i, c:r] = torch.linspace(1, 0, r - c, device=device)
        return fb

    fb = mel_fb(ref.device)
    ref_stft = torch.stft(ref_flat, n_fft, hop_length=hop, window=window, return_complex=True, normalized=True)
    gen_stft = torch.stft(gen_flat, n_fft, hop_length=hop, window=window, return_complex=True, normalized=True)
    ref_mel = torch.matmul(fb, ref_stft.abs().pow(2)).clamp(min=eps).log()
    gen_mel = torch.matmul(fb, gen_stft.abs().pow(2)).clamp(min=eps).log()
    return F.l1_loss(gen_mel, ref_mel).item()


# =========================================================================
# Main
# =========================================================================

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt", type=str, required=True, help="Path to student checkpoint")
    parser.add_argument("--audio", type=str, required=True, help="Path to WAV file or directory of MP3s")
    parser.add_argument("--num-clips", type=int, default=5, help="Number of clips to evaluate")
    parser.add_argument("--clip-duration", type=float, default=10.0, help="Duration per clip in seconds")
    parser.add_argument("--output-dir", type=str, default=None, help="Directory for audio outputs")
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--speed-trials", type=int, default=10, help="Number of speed measurement trials")
    args = parser.parse_args()

    device = args.device
    ckpt_path = Path(args.ckpt)
    ckpt_name = ckpt_path.stem

    if args.output_dir:
        out_dir = Path(args.output_dir)
    else:
        out_dir = ckpt_path.parent / f"eval_{ckpt_name}"
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"Loading checkpoint: {ckpt_path}")
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)

    config = ckpt.get("config", {})
    student = FastOobleckDecoder(
        channels=config.get("channels", 128),
        input_channels=config.get("input_channels", 64),
        audio_channels=config.get("audio_channels", 2),
        upsampling_ratios=config.get("upsampling_ratios", [10, 6, 4, 4, 2]),
        channel_multiples=config.get("channel_multiples", [1, 2, 4, 8, 8]),
    ).to(device)
    student.load_state_dict(ckpt["student_state_dict"])
    student.eval()
    step = ckpt.get("step", "unknown")
    print(f"Student loaded (step {step}), {sum(p.numel() for p in student.parameters()) / 1e6:.1f}M params")

    # Load teacher VAE (need full VAE for encoding)
    print("Loading teacher VAE...")
    from diffusers import AutoencoderOobleck
    vae = AutoencoderOobleck.from_pretrained("ACE-Step/Ace-Step1.5", subfolder="vae")
    vae = vae.to(device, dtype=torch.float32)
    vae.eval()
    teacher_decoder = vae.decoder
    print(f"Teacher loaded, {sum(p.numel() for p in teacher_decoder.parameters()) / 1e6:.1f}M params")

    # Load real audio
    sr = 48000
    hop = 1920
    clip_samples = int(args.clip_duration * sr)

    audio_path = Path(args.audio)
    if audio_path.is_dir():
        # Directory of MP3s: pick random tracks, convert to 48kHz stereo
        import glob
        import subprocess
        import tempfile
        import random

        mp3_files = sorted(glob.glob(str(audio_path / "**" / "*.mp3"), recursive=True))
        if not mp3_files:
            raise RuntimeError(f"No MP3 files found in {audio_path}")
        random.seed(42)  # reproducible selection
        random.shuffle(mp3_files)

        print(f"\nFound {len(mp3_files)} MP3 files in {audio_path}")
        print(f"Selecting {args.num_clips} random tracks, {args.clip_duration}s each")
        print(f"Output: {out_dir}\n")

        # Load one clip per track
        audio_clips = []
        track_names = []
        idx = 0
        while len(audio_clips) < args.num_clips and idx < len(mp3_files):
            mp3 = mp3_files[idx]
            idx += 1
            try:
                tmp = tempfile.mktemp(suffix=".wav")
                result = subprocess.run(
                    ["ffmpeg", "-y", "-i", mp3,
                     "-ar", str(sr), "-ac", "2", "-f", "wav", tmp],
                    capture_output=True, timeout=30,
                )
                if result.returncode != 0:
                    continue
                data, fsr = sf.read(tmp, dtype="float32")
                os.unlink(tmp)
                waveform = torch.tensor(data, dtype=torch.float32).T  # [2, samples]
                if waveform.shape[-1] < clip_samples:
                    continue
                # Random clip from the track
                start = random.randint(0, waveform.shape[-1] - clip_samples)
                clip = waveform[:, start:start + clip_samples]
                peak = clip.abs().max()
                if peak > 1e-6:
                    clip = clip / peak
                audio_clips.append(clip)
                track_names.append(Path(mp3).stem)
            except Exception:
                continue

        if len(audio_clips) < args.num_clips:
            print(f"WARNING: only loaded {len(audio_clips)} of {args.num_clips} requested clips")

    else:
        # Single WAV file: take evenly spaced clips
        print(f"\nLoading audio: {args.audio}")
        audio_data, audio_sr = sf.read(str(audio_path), dtype="float32")
        assert audio_sr == sr, f"Expected {sr}Hz, got {audio_sr}Hz"
        waveform = torch.tensor(audio_data, dtype=torch.float32).T
        print(f"Audio: {waveform.shape[1]/sr:.1f}s, {waveform.shape[0]} channels")
        total_samples = waveform.shape[1]

        audio_clips = []
        track_names = []
        for i in range(args.num_clips):
            start = int(i * (total_samples - clip_samples) / max(args.num_clips - 1, 1))
            clip = waveform[:, start:start + clip_samples]
            peak = clip.abs().max()
            if peak > 1e-6:
                clip = clip / peak
            audio_clips.append(clip)
            track_names.append(f"clip{i+1}")

    print(f"Evaluating {len(audio_clips)} clips of {args.clip_duration}s each")
    print(f"Output: {out_dir}\n")

    # Student vs teacher (distillation quality)
    all_snr = []
    all_stft = []
    all_mel = []
    all_hf = []
    # Teacher vs original (VAE reconstruction quality)
    all_snr_t_vs_orig = []
    all_stft_t_vs_orig = []
    all_mel_t_vs_orig = []
    # Student vs original (end-to-end practical quality)
    all_snr_s_vs_orig = []
    all_stft_s_vs_orig = []
    all_mel_s_vs_orig = []

    with torch.no_grad():
        for i, clip in enumerate(audio_clips):
            clip_gpu = clip.unsqueeze(0).to(device)
            enc_out = vae.encode(clip_gpu)
            z = enc_out.latent_dist.sample()

            teacher_audio = teacher_decoder(z)
            student_audio = student(z)

            # Trim all to same length
            min_len = min(teacher_audio.shape[-1], student_audio.shape[-1], clip_gpu.shape[-1])
            teacher_audio = teacher_audio[..., :min_len]
            student_audio = student_audio[..., :min_len]
            original_audio = clip_gpu[..., :min_len]

            # Student vs teacher (distillation fidelity)
            snr = compute_snr(teacher_audio, student_audio)
            stft_dist = compute_stft_distance(teacher_audio, student_audio)
            mel_dist = compute_mel_distance(teacher_audio, student_audio)
            hf = compute_hf_energy_ratio(teacher_audio, student_audio)
            all_snr.append(snr)
            all_stft.append(stft_dist)
            all_mel.append(mel_dist)
            all_hf.append(hf)

            # Teacher vs original (VAE ceiling)
            snr_to = compute_snr(original_audio, teacher_audio)
            stft_to = compute_stft_distance(original_audio, teacher_audio)
            mel_to = compute_mel_distance(original_audio, teacher_audio)
            all_snr_t_vs_orig.append(snr_to)
            all_stft_t_vs_orig.append(stft_to)
            all_mel_t_vs_orig.append(mel_to)

            # Student vs original (what the user actually hears)
            snr_so = compute_snr(original_audio, student_audio)
            stft_so = compute_stft_distance(original_audio, student_audio)
            mel_so = compute_mel_distance(original_audio, student_audio)
            all_snr_s_vs_orig.append(snr_so)
            all_stft_s_vs_orig.append(stft_so)
            all_mel_s_vs_orig.append(mel_so)

            name = track_names[i]
            print(f"  [{name}]")
            print(f"    student vs teacher:   SNR={snr:.1f} dB  STFT={stft_dist:.4f}  Mel={mel_dist:.4f}  HF_match={hf['hf_energy_match']:.3f}")
            print(f"    teacher vs original:  SNR={snr_to:.1f} dB  STFT={stft_to:.4f}  Mel={mel_to:.4f}")
            print(f"    student vs original:  SNR={snr_so:.1f} dB  STFT={stft_so:.4f}  Mel={mel_so:.4f}")

            # Save audio
            t_np = teacher_audio[0].cpu().numpy().T
            s_np = student_audio[0].cpu().numpy().T
            o_np = original_audio[0].cpu().numpy().T
            sf.write(out_dir / f"{name}_original.wav", o_np, sr)
            sf.write(out_dir / f"{name}_teacher.wav", t_np, sr)
            sf.write(out_dir / f"{name}_student.wav", s_np, sr)

    # Speed benchmark (separate, no contention with metrics computation)
    print(f"\nSpeed benchmark ({args.speed_trials} trials, {args.clip_duration}s clip)...")
    with torch.no_grad():
        bench_clip = audio_clips[0].unsqueeze(0).to(device)
        z_bench = vae.encode(bench_clip).latent_dist.sample()

    # Warmup
    for _ in range(3):
        _ = teacher_decoder(z_bench)
        _ = student(z_bench)
    torch.cuda.synchronize()

    teacher_times = []
    student_times = []
    for _ in range(args.speed_trials):
        torch.cuda.synchronize()
        t0 = time.time()
        _ = teacher_decoder(z_bench)
        torch.cuda.synchronize()
        teacher_times.append(time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        _ = student(z_bench)
        torch.cuda.synchronize()
        student_times.append(time.time() - t0)

    # Aggregate HF metrics
    avg_hf_match = np.mean([h["hf_energy_match"] for h in all_hf])
    avg_ref_hf = np.mean([h["ref_hf_ratio"] for h in all_hf])
    avg_gen_hf = np.mean([h["gen_hf_ratio"] for h in all_hf])
    avg_rolloff_ref = np.mean([h["spectral_rolloff_ref_hz"] for h in all_hf])
    avg_rolloff_gen = np.mean([h["spectral_rolloff_gen_hz"] for h in all_hf])

    # Summary
    print(f"\n{'='*60}")
    print(f"Checkpoint: {ckpt_name} (step {step})")
    print(f"{'='*60}")
    print(f"\n--- Student vs Teacher (distillation fidelity) ---")
    print(f"SNR:          {np.mean(all_snr):.1f} dB  (std {np.std(all_snr):.1f})")
    print(f"STFT dist:    {np.mean(all_stft):.4f}")
    print(f"Mel dist:     {np.mean(all_mel):.4f}")
    print(f"HF match:     {avg_hf_match:.3f}  (teacher HF={avg_ref_hf:.4f}, student HF={avg_gen_hf:.4f})")
    print(f"Rolloff:      teacher={avg_rolloff_ref:.0f}Hz, student={avg_rolloff_gen:.0f}Hz")
    print(f"\n--- Teacher vs Original (VAE reconstruction ceiling) ---")
    print(f"SNR:          {np.mean(all_snr_t_vs_orig):.1f} dB  (std {np.std(all_snr_t_vs_orig):.1f})")
    print(f"STFT dist:    {np.mean(all_stft_t_vs_orig):.4f}")
    print(f"Mel dist:     {np.mean(all_mel_t_vs_orig):.4f}")
    print(f"\n--- Student vs Original (end-to-end practical quality) ---")
    print(f"SNR:          {np.mean(all_snr_s_vs_orig):.1f} dB  (std {np.std(all_snr_s_vs_orig):.1f})")
    print(f"STFT dist:    {np.mean(all_stft_s_vs_orig):.4f}")
    print(f"Mel dist:     {np.mean(all_mel_s_vs_orig):.4f}")
    print(f"\n--- Speed ---")
    print(f"Teacher:      {np.mean(teacher_times)*1000:.0f}ms avg (std {np.std(teacher_times)*1000:.0f}ms)")
    print(f"Student:      {np.mean(student_times)*1000:.0f}ms avg (std {np.std(student_times)*1000:.0f}ms)")
    print(f"Speedup:      {np.mean(teacher_times)/np.mean(student_times):.2f}x")
    print(f"\nAudio saved to: {out_dir}")

    # Save metrics JSON
    results = {
        "checkpoint": str(ckpt_path),
        "step": step,
        "audio_source": str(args.audio),
        "num_clips": args.num_clips,
        "clip_duration_s": args.clip_duration,
        "student_vs_teacher": {
            "snr_mean": round(np.mean(all_snr), 2),
            "snr_std": round(np.std(all_snr), 2),
            "stft_dist": round(np.mean(all_stft), 4),
            "mel_dist": round(np.mean(all_mel), 4),
            "hf_energy_match": round(avg_hf_match, 4),
            "hf_ratio_teacher": round(avg_ref_hf, 4),
            "hf_ratio_student": round(avg_gen_hf, 4),
            "spectral_rolloff_teacher_hz": round(avg_rolloff_ref),
            "spectral_rolloff_student_hz": round(avg_rolloff_gen),
        },
        "teacher_vs_original": {
            "snr_mean": round(np.mean(all_snr_t_vs_orig), 2),
            "snr_std": round(np.std(all_snr_t_vs_orig), 2),
            "stft_dist": round(np.mean(all_stft_t_vs_orig), 4),
            "mel_dist": round(np.mean(all_mel_t_vs_orig), 4),
        },
        "student_vs_original": {
            "snr_mean": round(np.mean(all_snr_s_vs_orig), 2),
            "snr_std": round(np.std(all_snr_s_vs_orig), 2),
            "stft_dist": round(np.mean(all_stft_s_vs_orig), 4),
            "mel_dist": round(np.mean(all_mel_s_vs_orig), 4),
        },
        "speed": {
            "teacher_ms_avg": round(np.mean(teacher_times) * 1000, 1),
            "student_ms_avg": round(np.mean(student_times) * 1000, 1),
            "speedup": round(np.mean(teacher_times) / np.mean(student_times), 2),
        },
    }
    with open(out_dir / "metrics.json", "w") as f:
        json.dump(results, f, indent=2)


if __name__ == "__main__":
    main()