code/scripts/run_comprehensive_eval.py

#!/usr/bin/env python3
"""Comprehensive model comparison evaluation.

Generates audio for multiple models × reference voices × prompts,
scores with aux losses, builds HTML report with embedded MP3.
"""

import argparse
import base64
import hashlib
import io
import json
import logging
import math
import os
import random
import shutil
import subprocess
import sys
import tempfile
import time
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import numpy as np

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

# ── Paths ──────────────────────────────────────────────────────────────────
DRAMABOX_DIR = "/home/deployer/laion/DramaBox"
INFERENCE_SCRIPT = os.path.join(DRAMABOX_DIR, "src", "inference.py")
CHECKPOINT = os.path.join(DRAMABOX_DIR, "models", "ltx-2.3-22b-dev-audio-only-v13-merged.safetensors")
FULL_CHECKPOINT = os.path.join(DRAMABOX_DIR, "models", "ltx-2.3-22b-dev.safetensors")
GEMMA_ROOT = "/home/deployer/.cache/dramabox/models--unsloth--gemma-3-12b-it-bnb-4bit/snapshots/826e729dbaeea4ecb143738eed2bcf3539ebf7bf"
PYTHON = "/home/deployer/miniconda3/envs/ml-general/bin/python"
VAP_DIR = "/home/deployer/laion/Voice-Acting-Pipeline"

# ── Models ─────────────────────────────────────────────────────────────────
MODELS = {
    "vanilla": {
        "name": "Vanilla DramaBox",
        "lora": None,
        "desc": "Base DramaBox model without any fine-tuning (ResembleAI/Dramabox)"
    },
    "laionbox_v01": {
        "name": "LaionBox v0.1-wip",
        "lora": "/home/deployer/.cache/huggingface/hub/models--laion--laionbox-v0.1-wip/snapshots/66176d2a653a013a7b71c1ccb7a7a4d4cf514b0d/lora_epoch5.safetensors",
        "desc": "Previous best LoRA (5-epoch diff reward training on DramaBox+Emolia data)"
    },
    "nat_best_flow": {
        "name": "Nat-Only Best-Flow (step 160)",
        "lora": os.path.join(VAP_DIR, "finetune_output/nat_only_2ep/lora_step_00160.safetensors"),
        "desc": "Naturalness-only training (VoiceCLAP-7B), best flow loss=0.528 at step 160"
    },
    "nat_best_nat": {
        "name": "Nat-Only Best-Naturalness (step 190)",
        "lora": os.path.join(VAP_DIR, "finetune_output/nat_only_2ep/lora_step_00190.safetensors"),
        "desc": "Naturalness-only training (VoiceCLAP-7B), best naturalness=0.111 at step 190"
    },
}


def fetch_json_url(url):
    """Fetch JSON from a URL."""
    import urllib.request
    with urllib.request.urlopen(url) as resp:
        return json.loads(resp.read().decode())


def select_prompts():
    """Select 5 prompts from each JSON + 5 podcast-style emotion prompts."""
    prompts = []

    # 1. Extreme Physical challenges
    logging.info("Fetching extreme physical challenges...")
    data = fetch_json_url("https://raw.githubusercontent.com/LAION-AI/Voice-Acting-Pipeline/refs/heads/main/data/acting_challenges_extreme_physical.json")
    random.seed(42)
    selected = random.sample(data, 5)
    for entry in selected:
        prompts.append({
            "id": f"extreme_{entry['id']}",
            "source": "acting_challenges_extreme_physical",
            "title": entry["title"],
            "prompt": entry["instruction"],
        })

    # 2. Existing Inspired challenges
    logging.info("Fetching existing inspired challenges...")
    data = fetch_json_url("https://raw.githubusercontent.com/LAION-AI/Voice-Acting-Pipeline/refs/heads/main/data/acting_challenges_existing_inspired.json")
    random.seed(43)
    selected = random.sample(data, 5)
    for entry in selected:
        prompts.append({
            "id": f"inspired_{entry['id']}",
            "source": "acting_challenges_existing_inspired",
            "title": entry["title"],
            "prompt": entry["instruction"],
        })

    # 3. CC2-C Archetype
    logging.info("Fetching CC2-C archetypes...")
    data = fetch_json_url("https://raw.githubusercontent.com/LAION-AI/Voice-Acting-Pipeline/refs/heads/main/data/dramabox_cc2c_archetype.json")
    random.seed(44)
    selected = random.sample(data, 5)
    for entry in selected:
        prompts.append({
            "id": f"archetype_{entry['id']}",
            "source": "dramabox_cc2c_archetype",
            "title": entry.get("sample_info", {}).get("archetype", f"Archetype {entry['id']}"),
            "prompt": entry["dramabox_prompt"],
        })

    # 4. Podcast-style emotion prompts (30 words direct speech each)
    podcast_prompts = [
        {
            "id": "podcast_joy",
            "source": "podcast_emotion",
            "title": "Podcast - Joy",
            "prompt": 'A joyful podcast host speaks with infectious warmth and excitement: "Oh my goodness, I cannot believe we actually made it happen! This is the most incredible day of my entire life, and I want to share every beautiful moment with all of you listening right now!"'
        },
        {
            "id": "podcast_anger",
            "source": "podcast_emotion",
            "title": "Podcast - Anger",
            "prompt": 'An angry investigative journalist speaks with barely contained fury: "They lied to us, every single one of them. They looked us straight in the eye, smiled with those perfect teeth, and told us everything was absolutely fine while people were dying in the streets!"'
        },
        {
            "id": "podcast_fear",
            "source": "podcast_emotion",
            "title": "Podcast - Fear",
            "prompt": 'A terrified true-crime narrator whispers with trembling voice: "The door creaked open slowly behind me. I could hear breathing that wasn\'t mine, heavy and ragged, getting closer with each heartbeat. My hands were shaking so badly I could barely hold the phone."'
        },
        {
            "id": "podcast_sadness",
            "source": "podcast_emotion",
            "title": "Podcast - Sadness",
            "prompt": 'A grieving storyteller speaks softly with deep sorrow: "She was only thirty-two years old when the diagnosis came. We sat together in that cold hospital room, holding hands, both knowing that our time together had suddenly become so incredibly precious and painfully finite."'
        },
        {
            "id": "podcast_gratitude",
            "source": "podcast_emotion",
            "title": "Podcast - Gratitude",
            "prompt": 'A deeply grateful motivational speaker addresses listeners with sincere warmth: "To every single person who believed in me when I couldn\'t believe in myself, who stayed when it would have been easier to leave, from the absolute bottom of my heart, thank you so much."'
        },
    ]
    prompts.extend(podcast_prompts)

    return prompts


def get_ref_audios(refs_dir):
    """Get list of reference audio files."""
    refs = []
    for f in sorted(os.listdir(refs_dir)):
        if f.endswith((".wav", ".mp3", ".flac")):
            refs.append(os.path.join(refs_dir, f))
    return refs


def build_task_list(models, prompts, refs, seeds, output_dir):
    """Build list of generation tasks."""
    tasks = []
    for model_key, model_info in models.items():
        for prompt_info in prompts:
            for seed in seeds:
                # Voice clone tasks (with reference)
                for ref_path in refs:
                    ref_name = Path(ref_path).stem
                    out_name = f"{model_key}__{prompt_info['id']}__{ref_name}__seed{seed}.wav"
                    tasks.append({
                        "model_key": model_key,
                        "model_name": model_info["name"],
                        "lora": model_info["lora"],
                        "prompt": prompt_info["prompt"],
                        "prompt_id": prompt_info["id"],
                        "prompt_title": prompt_info["title"],
                        "prompt_source": prompt_info["source"],
                        "ref": ref_path,
                        "ref_name": ref_name,
                        "seed": seed,
                        "output": os.path.join(output_dir, "wavs", out_name),
                        "mode": "voice_clone",
                    })
                # Unconditional task (no reference)
                out_name = f"{model_key}__{prompt_info['id']}__uncond__seed{seed}.wav"
                tasks.append({
                    "model_key": model_key,
                    "model_name": model_info["name"],
                    "lora": model_info["lora"],
                    "prompt": prompt_info["prompt"],
                    "prompt_id": prompt_info["id"],
                    "prompt_title": prompt_info["title"],
                    "prompt_source": prompt_info["source"],
                    "ref": None,
                    "ref_name": "uncond",
                    "seed": seed,
                    "output": os.path.join(output_dir, "wavs", out_name),
                    "mode": "unconditional",
                })
    return tasks


def run_inference_task(task, gpu_id):
    """Run a single inference task on a specific GPU."""
    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    # Remove bad LD_LIBRARY_PATH
    env.pop("LD_LIBRARY_PATH", None)

    cmd = [
        PYTHON, INFERENCE_SCRIPT,
        "--prompt", task["prompt"],
        "--output", task["output"],
        "--checkpoint", CHECKPOINT,
        "--full-checkpoint", FULL_CHECKPOINT,
        "--gemma-root", GEMMA_ROOT,
        "--seed", str(task["seed"]),
    ]

    if task["lora"]:
        cmd.extend(["--lora", task["lora"], "--lora-rank", "128"])

    if task["ref"]:
        cmd.extend(["--voice-sample", task["ref"]])
    else:
        cmd.append("--no-ref")

    t0 = time.time()
    try:
        result = subprocess.run(
            cmd, capture_output=True, text=True, timeout=300, env=env,
            cwd=DRAMABOX_DIR
        )
        elapsed = time.time() - t0
        success = result.returncode == 0 and os.path.exists(task["output"])
        if not success:
            logging.warning(f"Failed: {task['output']} (rc={result.returncode})")
            if result.stderr:
                logging.warning(f"  stderr: {result.stderr[-200:]}")
        return {**task, "success": success, "elapsed": elapsed, "gpu": gpu_id}
    except subprocess.TimeoutExpired:
        return {**task, "success": False, "elapsed": 300, "gpu": gpu_id}
    except Exception as e:
        return {**task, "success": False, "elapsed": time.time() - t0, "gpu": gpu_id, "error": str(e)}


def run_generation_phase(tasks, num_gpus=8):
    """Run all generation tasks across multiple GPUs."""
    logging.info(f"Generating {len(tasks)} audio samples across {num_gpus} GPUs...")

    results = []
    completed = 0

    with ProcessPoolExecutor(max_workers=num_gpus) as executor:
        futures = {}
        for i, task in enumerate(tasks):
            gpu_id = i % num_gpus
            fut = executor.submit(run_inference_task, task, gpu_id)
            futures[fut] = task

        for fut in as_completed(futures):
            result = fut.result()
            results.append(result)
            completed += 1
            if completed % 20 == 0 or completed == len(tasks):
                ok = sum(1 for r in results if r.get("success"))
                logging.info(f"  [{completed}/{len(tasks)}] {ok} successful")

    ok = sum(1 for r in results if r.get("success"))
    logging.info(f"Generation complete: {ok}/{len(tasks)} successful")
    return results


def score_all_audio(results, classifiers_dir, device="cuda:0"):
    """Score all generated audio with CLAP-small, CLAP-large (7B), centroid, quality MLP, and WavLM-SV."""
    import torch
    import torchaudio
    import torch.nn.functional as F
    import torch.nn as nn
    import soundfile as sf

    logging.info("Loading scoring models...")

    # ── CLAP-small ──
    from transformers import AutoModel, AutoTokenizer, AutoFeatureExtractor
    clap_model = AutoModel.from_pretrained("laion/voiceclap-small", trust_remote_code=True).to(device).eval()
    clap_tokenizer = AutoTokenizer.from_pretrained("laion/voiceclap-small", trust_remote_code=True)

    # ── CLAP-large (VoiceCLAP 7B with INT4) ──
    logging.info("Loading VoiceCLAP-7B (INT4)...")
    from sentence_transformers import SentenceTransformer
    from transformers import BitsAndBytesConfig
    bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
    clap_large = SentenceTransformer(
        "gijs/voiceclap-lco-7b-lora",
        model_kwargs={"quantization_config": bnb_config, "torch_dtype": torch.bfloat16, "trust_remote_code": True},
        trust_remote_code=True,
    )
    logging.info(f"VoiceCLAP-7B loaded, embedding dim={clap_large.get_sentence_embedding_dimension()}")

    # ── Centroids (from CLAP-small embeddings) ──
    emb_data = torch.load(os.path.join(classifiers_dir, "clap_embeddings.pt"),
                          map_location="cpu", weights_only=False)
    dramabox_embs = emb_data["dramabox_embeddings"]
    emilia_embs = emb_data["emilia_embeddings"]
    n_train = int(len(dramabox_embs) * 0.8)
    synth_centroid = F.normalize(dramabox_embs[:n_train].float().mean(0, keepdim=True), p=2, dim=-1).to(device)
    real_centroid = F.normalize(emilia_embs[:n_train].float().mean(0, keepdim=True), p=2, dim=-1).to(device)

    # ── Quality MLP ──
    ckpt = torch.load(os.path.join(classifiers_dir, "quality_classifier.pt"),
                      map_location="cpu", weights_only=False)

    class BinaryMLP(nn.Module):
        def __init__(self, d_in, h1, h2):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(d_in, h1), nn.ReLU(), nn.Dropout(0.3),
                nn.Linear(h1, h2), nn.ReLU(), nn.Dropout(0.3),
                nn.Linear(h2, 1))
        def forward(self, x):
            return self.net(x)

    quality_mlp = BinaryMLP(ckpt["input_dim"], ckpt["hidden1"], ckpt["hidden2"])
    quality_mlp.load_state_dict(ckpt["model_state_dict"])
    quality_mlp.eval().to(device)

    # ── WavLM-SV ──
    from transformers import WavLMForXVector
    wavlm = WavLMForXVector.from_pretrained("microsoft/wavlm-base-plus-sv").to(device).eval()
    wavlm_fe = AutoFeatureExtractor.from_pretrained("microsoft/wavlm-base-plus-sv")

    # ── Pre-encode positive/negative text for both CLAP models ──
    pos_text = "Realistic, genuine, spontaneous, authentic, sensual, natural voice with all imperfections and organic microdistractions a natural situation brings with it"
    neg_text = "distorted, unnatural, robotic, distortion"

    # CLAP-small text embeddings
    with torch.no_grad():
        pos_tok = clap_tokenizer(pos_text, return_tensors="pt", padding=True, truncation=True, max_length=77).to(device)
        neg_tok = clap_tokenizer(neg_text, return_tensors="pt", padding=True, truncation=True, max_length=77).to(device)
        pos_emb = clap_model.encode_text(pos_tok["input_ids"], pos_tok.get("attention_mask"))
        neg_emb = clap_model.encode_text(neg_tok["input_ids"], neg_tok.get("attention_mask"))
        pos_emb = pos_emb / pos_emb.norm(dim=-1, keepdim=True)
        neg_emb = neg_emb / neg_emb.norm(dim=-1, keepdim=True)

    # CLAP-large text embeddings
    with torch.no_grad():
        pos_emb_large = clap_large.encode([pos_text], convert_to_tensor=True)
        neg_emb_large = clap_large.encode([neg_text], convert_to_tensor=True)
        pos_emb_large = F.normalize(pos_emb_large, p=2, dim=-1)
        neg_emb_large = F.normalize(neg_emb_large, p=2, dim=-1)

    logging.info("All scoring models loaded. Starting scoring...")

    scored = []
    for i, result in enumerate(results):
        if not result.get("success"):
            scored.append({**result, "scores": None})
            continue

        wav_path = result["output"]
        try:
            waveform, sr = torchaudio.load(wav_path)
            if waveform.shape[0] > 1:
                waveform = waveform.mean(dim=0, keepdim=True)

            with torch.no_grad():
                # Resample to 16kHz for CLAP-small and WavLM
                if sr != 16000:
                    waveform_16k = torchaudio.functional.resample(waveform, sr, 16000)
                else:
                    waveform_16k = waveform

                # ── CLAP-small scoring ──
                audio_emb = clap_model.encode_waveform(waveform_16k.to(device), sample_rate=16000)
                audio_emb = audio_emb / audio_emb.norm(dim=-1, keepdim=True)

                nat_score_small = (audio_emb @ pos_emb.T).item() - (audio_emb @ neg_emb.T).item()
                cent_score = (audio_emb @ real_centroid.T).item() - (audio_emb @ synth_centroid.T).item()
                quality_score = torch.sigmoid(quality_mlp(audio_emb)).item()

                # ── CLAP-large (7B) scoring via temp WAV file ──
                tmp_path = f"/dev/shm/eval_clap_large_{os.getpid()}.wav"
                try:
                    # Save as 16kHz mono WAV for CLAP-large
                    sf.write(tmp_path, waveform_16k.squeeze(0).numpy(), 16000)
                    audio_emb_large = clap_large.encode([{"audio": tmp_path}], convert_to_tensor=True)
                    audio_emb_large = F.normalize(audio_emb_large, p=2, dim=-1)
                    nat_score_large = (audio_emb_large @ pos_emb_large.T).item() - (audio_emb_large @ neg_emb_large.T).item()
                finally:
                    if os.path.exists(tmp_path):
                        os.remove(tmp_path)

                # ── Speaker similarity (WavLM-SV) ──
                spk_score = None
                if result["ref"]:
                    ref_wav, ref_sr = torchaudio.load(result["ref"])
                    if ref_wav.shape[0] > 1:
                        ref_wav = ref_wav.mean(dim=0, keepdim=True)
                    if ref_sr != 16000:
                        ref_wav_16k = torchaudio.functional.resample(ref_wav, ref_sr, 16000)
                    else:
                        ref_wav_16k = ref_wav

                    gen_wav_16k = waveform_16k

                    # Truncate to 10s max
                    max_len = 16000 * 10
                    ref_wav_16k = ref_wav_16k[:, :max_len]
                    gen_wav_16k = gen_wav_16k[:, :max_len]

                    ref_in = wavlm_fe(ref_wav_16k.squeeze(0), sampling_rate=16000, return_tensors="pt", padding=True)
                    gen_in = wavlm_fe(gen_wav_16k.squeeze(0), sampling_rate=16000, return_tensors="pt", padding=True)

                    ref_emb_sv = wavlm(**{k: v.to(device) for k, v in ref_in.items()}).embeddings
                    gen_emb_sv = wavlm(**{k: v.to(device) for k, v in gen_in.items()}).embeddings

                    ref_emb_sv = ref_emb_sv / ref_emb_sv.norm(dim=-1, keepdim=True)
                    gen_emb_sv = gen_emb_sv / gen_emb_sv.norm(dim=-1, keepdim=True)
                    spk_score = (ref_emb_sv @ gen_emb_sv.T).item()

                scores = {
                    "naturalness_small": round(nat_score_small, 4),
                    "naturalness_large": round(nat_score_large, 4),
                    "centroid": round(cent_score, 4),
                    "quality": round(quality_score, 4),
                    "speaker_sim": round(spk_score, 4) if spk_score is not None else None,
                }

            scored.append({**result, "scores": scores})

        except Exception as e:
            logging.warning(f"Scoring failed for {wav_path}: {e}")
            import traceback
            traceback.print_exc()
            scored.append({**result, "scores": None})

        if (i + 1) % 20 == 0:
            logging.info(f"  Scored {i+1}/{len(results)}")

    logging.info(f"Scoring complete: {sum(1 for s in scored if s.get('scores'))} scored")
    return scored


def wav_to_mp3_base64(wav_path, bitrate="96k"):
    """Convert WAV to base64-encoded 96kbps mono MP3."""
    try:
        result = subprocess.run(
            ["ffmpeg", "-y", "-i", wav_path, "-ac", "1", "-ab", bitrate,
             "-f", "mp3", "pipe:1"],
            capture_output=True, timeout=30
        )
        if result.returncode == 0:
            return base64.b64encode(result.stdout).decode("ascii")
    except Exception:
        pass
    return None


def build_html_report(scored_results, models, prompts, refs, output_path):
    """Build comprehensive HTML report with embedded MP3 audio."""
    logging.info("Building HTML report...")

    # Compute mean scores per model
    ALL_METRICS = ["naturalness_small", "naturalness_large", "centroid", "quality", "speaker_sim"]

    model_scores = defaultdict(lambda: defaultdict(list))
    for r in scored_results:
        if r.get("scores"):
            mk = r["model_key"]
            for k, v in r["scores"].items():
                if v is not None:
                    model_scores[mk][k].append(v)
            # Backward compat: old scored_results may have "naturalness" instead of "naturalness_small"
            if "naturalness" in r["scores"] and "naturalness_small" not in r["scores"]:
                if r["scores"]["naturalness"] is not None:
                    model_scores[mk]["naturalness_small"].append(r["scores"]["naturalness"])

    model_means = {}
    for mk in models:
        means = {}
        for metric in ALL_METRICS:
            vals = model_scores[mk][metric]
            means[metric] = round(sum(vals) / len(vals), 4) if vals else None
        model_means[mk] = means

    # Build HTML
    html_parts = []
    html_parts.append("""<!DOCTYPE html>
<html><head>
<meta charset="utf-8">
<title>DramaBox Model Comparison - Comprehensive Evaluation</title>
<style>
body { font-family: 'Segoe UI', system-ui, sans-serif; margin: 20px; background: #0d1117; color: #c9d1d9; }
h1 { color: #58a6ff; border-bottom: 2px solid #30363d; padding-bottom: 10px; }
h2 { color: #79c0ff; margin-top: 30px; }
h3 { color: #d2a8ff; }
.score-table { border-collapse: collapse; margin: 15px 0; width: 100%; }
.score-table th, .score-table td { border: 1px solid #30363d; padding: 8px 12px; text-align: center; }
.score-table th { background: #161b22; color: #58a6ff; }
.score-table tr:nth-child(even) { background: #161b22; }
.score-table tr:hover { background: #1c2128; }
.model-card { background: #161b22; border: 1px solid #30363d; border-radius: 8px; padding: 15px; margin: 10px 0; }
.metric-good { color: #3fb950; font-weight: bold; }
.metric-bad { color: #f85149; }
.metric-neutral { color: #d29922; }
.prompt-section { background: #161b22; border-left: 3px solid #58a6ff; padding: 10px 15px; margin: 15px 0; border-radius: 0 8px 8px 0; }
.audio-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(280px, 1fr)); gap: 12px; margin: 10px 0; }
.audio-card { background: #0d1117; border: 1px solid #30363d; border-radius: 6px; padding: 10px; }
.audio-card audio { width: 100%; margin: 5px 0; }
.audio-card .model-name { font-weight: bold; color: #58a6ff; font-size: 0.9em; }
.audio-card .scores { font-size: 0.8em; color: #8b949e; }
.audio-card .scores span { margin-right: 8px; }
.explanation { background: #161b22; border: 1px solid #30363d; border-radius: 8px; padding: 20px; margin: 20px 0; line-height: 1.6; }
.best-score { background: #0d2818; border: 1px solid #238636; }
.summary-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)); gap: 15px; margin: 20px 0; }
.summary-card { background: #161b22; border: 1px solid #30363d; border-radius: 8px; padding: 15px; text-align: center; }
.summary-card .model-name { font-size: 1.1em; font-weight: bold; color: #58a6ff; margin-bottom: 10px; }
.summary-card .metric { margin: 5px 0; }
.summary-card .metric .label { color: #8b949e; font-size: 0.85em; }
.summary-card .metric .value { font-size: 1.3em; font-weight: bold; }
</style>
</head><body>
""")

    html_parts.append("<h1>DramaBox Model Comparison - Comprehensive Evaluation</h1>")
    html_parts.append(f"<p>Generated: {time.strftime('%Y-%m-%d %H:%M UTC')}</p>")

    # Score explanations
    html_parts.append("""
<div class="explanation">
<h2>Score Explanations</h2>
<p>Each generated audio sample is evaluated on five metrics that measure different aspects of speech quality:</p>
<ul>
<li><strong>Naturalness-Small (CLAP-Small Net Score)</strong>: Measures naturalness using VoiceCLAP-small (lightweight model).
Computes: cos(audio, "natural realistic voice") - cos(audio, "robotic distorted").
Range: typically -0.5 to +0.5. <span class="metric-good">Higher is better</span>.</li>

<li><strong>Naturalness-Large (CLAP-7B Net Score)</strong>: Measures naturalness using VoiceCLAP-7B (gijs/voiceclap-lco-7b-lora).
Same formula as above but with a much larger 7B-parameter CLAP model producing 3584-dim embeddings.
More accurate signal than CLAP-small. <span class="metric-good">Higher is better</span>.</li>

<li><strong>Centroid Score</strong>: Measures similarity to real vs synthetic speech distributions.
Computes: cos(audio_embedding, real_speech_centroid) - cos(audio_embedding, synthetic_centroid).
Range: typically -0.5 to +0.5. <span class="metric-good">Higher is better</span> (closer to real speech distribution).</li>

<li><strong>Quality P(real)</strong>: A binary classifier (MLP) predicting probability that the audio is real speech.
Range: 0.0 to 1.0. <span class="metric-good">Higher is better</span>.
Score > 0.5 means "more likely real than synthetic".</li>

<li><strong>Speaker Similarity</strong>: Measures voice identity preservation during voice cloning.
Uses WavLM speaker verification embeddings: cos(reference_voice, generated_voice).
Range: 0.0 to 1.0. <span class="metric-good">Higher is better</span>.
Only computed for voice-cloned samples. Score > 0.85 = good, > 0.90 = excellent.</li>
</ul>
</div>
""")

    # Model summary cards
    html_parts.append("<h2>Model Summary (Mean Scores)</h2>")
    html_parts.append('<div class="summary-grid">')

    # Find best per metric for highlighting
    best_per_metric = {}
    for metric in ALL_METRICS:
        best_mk = max(models.keys(), key=lambda mk: model_means[mk].get(metric) or -999)
        best_per_metric[metric] = best_mk

    METRIC_LABELS = [
        ("naturalness_small", "Nat (CLAP-Small)"),
        ("naturalness_large", "Nat (CLAP-7B)"),
        ("centroid", "Centroid"),
        ("quality", "Quality P(real)"),
        ("speaker_sim", "Speaker Sim"),
    ]

    for mk, minfo in models.items():
        means = model_means[mk]
        html_parts.append(f'<div class="summary-card">')
        html_parts.append(f'<div class="model-name">{minfo["name"]}</div>')
        html_parts.append(f'<div style="font-size:0.8em;color:#8b949e;margin-bottom:8px">{minfo["desc"]}</div>')
        for metric, label in METRIC_LABELS:
            val = means.get(metric)
            cls = "metric-good" if best_per_metric.get(metric) == mk else ""
            val_str = f"{val:.4f}" if val is not None else "N/A"
            star = " ★" if best_per_metric.get(metric) == mk else ""
            html_parts.append(f'<div class="metric"><div class="label">{label}</div>'
                            f'<div class="value {cls}">{val_str}{star}</div></div>')
        html_parts.append('</div>')
    html_parts.append('</div>')

    # Detailed score table
    html_parts.append("<h2>Detailed Score Table</h2>")
    html_parts.append('<table class="score-table">')
    header = '<tr><th>Model</th>'
    for metric, label in METRIC_LABELS:
        header += f'<th>{label} ↑</th>'
    header += '</tr>'
    html_parts.append(header)
    for mk, minfo in models.items():
        means = model_means[mk]
        row = f'<tr><td><strong>{minfo["name"]}</strong></td>'
        for metric, _ in METRIC_LABELS:
            val = means.get(metric)
            is_best = best_per_metric.get(metric) == mk
            cls = "best-score" if is_best else ""
            val_str = f"{val:.4f}" if val is not None else "N/A"
            row += f'<td class="{cls}">{val_str}</td>'
        row += '</tr>'
        html_parts.append(row)
    html_parts.append('</table>')

    # Per-prompt results with audio
    html_parts.append("<h2>Audio Samples by Prompt</h2>")

    # Group results by prompt
    by_prompt = defaultdict(list)
    for r in scored_results:
        by_prompt[r["prompt_id"]].append(r)

    for prompt_info in prompts:
        pid = prompt_info["id"]
        results_for_prompt = by_prompt.get(pid, [])

        html_parts.append(f'<div class="prompt-section">')
        html_parts.append(f'<h3>{prompt_info["title"]} <small>({prompt_info["source"]})</small></h3>')
        prompt_text = prompt_info["prompt"]
        if len(prompt_text) > 300:
            prompt_text = prompt_text[:300] + "..."
        html_parts.append(f'<p style="color:#8b949e;font-size:0.9em">{prompt_text}</p>')

        # Group by ref
        by_ref = defaultdict(list)
        for r in results_for_prompt:
            by_ref[r["ref_name"]].append(r)

        for ref_name in sorted(by_ref.keys()):
            ref_results = by_ref[ref_name]
            ref_label = "Unconditional" if ref_name == "uncond" else f"Ref: {ref_name}"
            html_parts.append(f'<h4 style="color:#d2a8ff;margin:10px 0 5px">{ref_label}</h4>')

            # If this is a voice clone, show reference audio
            if ref_name != "uncond" and ref_results:
                ref_path = ref_results[0].get("ref")
                if ref_path and os.path.exists(ref_path):
                    ref_b64 = wav_to_mp3_base64(ref_path)
                    if ref_b64:
                        html_parts.append(f'<div style="margin-bottom:8px"><small style="color:#8b949e">Reference:</small> '
                                        f'<audio controls preload="none"><source src="data:audio/mpeg;base64,{ref_b64}" type="audio/mpeg"></audio></div>')

            html_parts.append('<div class="audio-grid">')
            for r in sorted(ref_results, key=lambda x: list(models.keys()).index(x["model_key"])):
                html_parts.append('<div class="audio-card">')
                html_parts.append(f'<div class="model-name">{r["model_name"]}</div>')

                if r.get("success") and os.path.exists(r["output"]):
                    b64 = wav_to_mp3_base64(r["output"])
                    if b64:
                        html_parts.append(f'<audio controls preload="none"><source src="data:audio/mpeg;base64,{b64}" type="audio/mpeg"></audio>')
                    else:
                        html_parts.append('<div style="color:#f85149">MP3 conversion failed</div>')
                else:
                    html_parts.append('<div style="color:#f85149">Generation failed</div>')

                if r.get("scores"):
                    s = r["scores"]
                    html_parts.append('<div class="scores">')
                    # Naturalness-small (backward compat: key may be "naturalness" or "naturalness_small")
                    nat_s = s.get("naturalness_small", s.get("naturalness"))
                    if nat_s is not None:
                        nat_cls = "metric-good" if nat_s > 0.3 else ("metric-bad" if nat_s < 0 else "metric-neutral")
                        html_parts.append(f'<span class="{nat_cls}">NatS:{nat_s:.3f}</span>')
                    # Naturalness-large
                    nat_l = s.get("naturalness_large")
                    if nat_l is not None:
                        nat_l_cls = "metric-good" if nat_l > 0.3 else ("metric-bad" if nat_l < 0 else "metric-neutral")
                        html_parts.append(f'<span class="{nat_l_cls}">NatL:{nat_l:.3f}</span>')
                    # Centroid
                    if s.get("centroid") is not None:
                        cent_cls = "metric-good" if s["centroid"] > 0 else "metric-bad"
                        html_parts.append(f'<span class="{cent_cls}">Cent:{s["centroid"]:.3f}</span>')
                    # Quality
                    if s.get("quality") is not None:
                        q_cls = "metric-good" if s["quality"] > 0.7 else ("metric-bad" if s["quality"] < 0.3 else "metric-neutral")
                        html_parts.append(f'<span class="{q_cls}">Q:{s["quality"]:.3f}</span>')
                    # Speaker similarity
                    if s.get("speaker_sim") is not None:
                        spk_cls = "metric-good" if s["speaker_sim"] > 0.85 else ("metric-bad" if s["speaker_sim"] < 0.7 else "metric-neutral")
                        html_parts.append(f'<span class="{spk_cls}">Spk:{s["speaker_sim"]:.3f}</span>')
                    html_parts.append('</div>')

                html_parts.append('</div>')
            html_parts.append('</div>')  # audio-grid

        html_parts.append('</div>')  # prompt-section

    html_parts.append("</body></html>")

    html_content = "\n".join(html_parts)
    with open(output_path, "w") as f:
        f.write(html_content)

    logging.info(f"HTML report: {output_path} ({len(html_content)/(1024*1024):.1f} MB)")
    return model_means


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default=os.path.join(VAP_DIR, "comprehensive_eval"))
    parser.add_argument("--refs-dir", default="/home/deployer/laion/test-refs")
    parser.add_argument("--num-gpus", type=int, default=8)
    parser.add_argument("--seeds", type=int, nargs="+", default=[42])
    parser.add_argument("--classifiers-dir", default=os.path.join(VAP_DIR, "classifiers"))
    parser.add_argument("--skip-generation", action="store_true")
    parser.add_argument("--skip-scoring", action="store_true")
    args = parser.parse_args()

    args.output_dir = os.path.abspath(args.output_dir)
    os.makedirs(os.path.join(args.output_dir, "wavs"), exist_ok=True)

    # 1. Select prompts
    prompts = select_prompts()
    logging.info(f"Selected {len(prompts)} prompts")

    # Save prompts
    with open(os.path.join(args.output_dir, "prompts.json"), "w") as f:
        json.dump(prompts, f, indent=2)

    # 2. Get reference audios
    refs = get_ref_audios(args.refs_dir)
    logging.info(f"Found {len(refs)} reference audios")

    # 3. Build task list
    tasks = build_task_list(MODELS, prompts, refs, args.seeds, args.output_dir)
    logging.info(f"Total tasks: {len(tasks)} ({len(MODELS)} models × {len(prompts)} prompts × ({len(refs)} refs + 1 uncond) × {len(args.seeds)} seeds)")

    # 4. Generate (skip tasks where WAV already exists)
    if not args.skip_generation:
        tasks_to_gen = [t for t in tasks if not os.path.exists(t["output"])]
        tasks_existing = [t for t in tasks if os.path.exists(t["output"])]
        logging.info(f"Skipping {len(tasks_existing)} existing WAVs, generating {len(tasks_to_gen)} new samples")
        if tasks_to_gen:
            gen_results = run_generation_phase(tasks_to_gen, args.num_gpus)
        else:
            gen_results = []
        # Mark existing files as successful
        existing_results = [{**t, "success": True, "elapsed": 0, "gpu": -1} for t in tasks_existing]
        results = existing_results + gen_results
        with open(os.path.join(args.output_dir, "gen_results.json"), "w") as f:
            json.dump(results, f, indent=2, default=str)
    else:
        with open(os.path.join(args.output_dir, "gen_results.json")) as f:
            results = json.load(f)

    # 5. Score
    if not args.skip_scoring:
        scored = score_all_audio(results, args.classifiers_dir)
        with open(os.path.join(args.output_dir, "scored_results.json"), "w") as f:
            json.dump(scored, f, indent=2, default=str)
    else:
        with open(os.path.join(args.output_dir, "scored_results.json")) as f:
            scored = json.load(f)

    # 6. Build HTML
    model_means = build_html_report(
        scored, MODELS, prompts, refs,
        os.path.join(args.output_dir, "eval_report.html")
    )

    # Print summary
    print("\n" + "="*70)
    print("MODEL COMPARISON SUMMARY")
    print("="*70)
    for mk, minfo in MODELS.items():
        means = model_means[mk]
        print(f"\n{minfo['name']}:")
        for metric in ["naturalness_small", "naturalness_large", "centroid", "quality", "speaker_sim"]:
            val = means.get(metric)
            print(f"  {metric:20s}: {val:.4f}" if val else f"  {metric:20s}: N/A")
    print("="*70)


if __name__ == "__main__":
    main()