#!/usr/bin/env python3
"""
pick_voice_ref.py
Pick the best reference clips for voice generation from the dataset made by prepare_dataset.py.

Chatterbox copies the voice, tone and energy of ONE reference clip, using about the first
10 seconds of it. So the reference matters a lot: this script ranks the kept dataset clips
by how strongly they match Brian's voice, how long they are, and how clean they are, then
copies the top picks into a folder so you can audition them.

Usage:
  python pick_voice_ref.py ..\\Brian-voice\\dataset_test --name brian
  python pick_voice_ref.py ..\\Lily-voice\\dataset --name lily
"""

import argparse
import csv
import re
import shutil
import sys
from pathlib import Path

import numpy as np
import soundfile as sf


def analyze(path):
    """Quick quality checks on a clip: loudness, silence, clipping."""
    x, sr = sf.read(str(path), dtype="float32", always_2d=True)
    x = x.mean(axis=1)
    frame = int(sr * 0.05)
    n = len(x) // frame
    if n == 0:
        return None
    rms = np.sqrt(np.mean(x[: n * frame].reshape(n, frame) ** 2, axis=1) + 1e-12)
    db = 20 * np.log10(rms + 1e-12)
    speech = db > (db.max() - 35)
    return {
        "seconds": len(x) / sr,
        "speech_ratio": float(speech.mean()),          # how much of the clip is talking
        "clipped": float(np.mean(np.abs(x) >= 0.999)),
    }


def main():
    p = argparse.ArgumentParser(description="Rank dataset clips as voice references.")
    p.add_argument("dataset", type=Path, help="Folder made by prepare_dataset.py (has report.csv)")
    p.add_argument("--name", default="brian",
                   help="Voice name for the default output folder (e.g. lily -> lily_voice_refs)")
    p.add_argument("--out", type=Path, help="Output folder (default: <name>_voice_refs)")
    p.add_argument("--top", type=int, default=6, help="How many candidates to copy")
    p.add_argument("--min-sec", type=float, default=7.0,
                   help="Shortest usable clip (Turbo needs more than 5 s)")
    p.add_argument("--max-sec", type=float, default=16.0)
    args = p.parse_args()

    voice_slug = re.sub(r"[^a-zA-Z0-9]+", "_", args.name.strip()).strip("_").lower()
    if not voice_slug:
        p.error("--name must contain at least one letter or number")
    out = args.out if args.out is not None else Path(f"{voice_slug}_voice_refs")

    report = args.dataset / "report.csv"
    wavs = args.dataset / "wavs"
    if not report.exists():
        sys.exit(f"Can't find {report}. Point this at the folder prepare_dataset.py created.")

    rows = []
    with open(report, newline="", encoding="utf-8") as f:
        for r in csv.DictReader(f):
            if r.get("status") != "kept":
                continue
            wav = wavs / f"{r['clip_id']}.wav"
            if not wav.exists():
                continue
            q = analyze(wav)
            if q is None or not (args.min_sec <= q["seconds"] <= args.max_sec):
                continue
            if q["clipped"] > 0.0005:
                continue
            sim = float(r["similarity"])
            # Strong voice match matters most; prefer clips that are mostly speech and ~10-14 s.
            length_bonus = 1.0 - min(abs(q["seconds"] - 12.0) / 6.0, 1.0)
            score = 3.0 * sim + 1.0 * q["speech_ratio"] + 0.5 * length_bonus
            rows.append({**r, **q, "wav": wav, "score": score})

    if not rows:
        sys.exit(f"No kept clips between {args.min_sec:.0f} and {args.max_sec:.0f} seconds. "
                 "Process more videos first, or lower --min-sec (not below 5.5).")

    rows.sort(key=lambda r: -r["score"])
    out.mkdir(parents=True, exist_ok=True)
    print(f"{len(rows)} usable clips found. Top {min(args.top, len(rows))}:\n")
    with open(out / "refs.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["file", "seconds", "similarity", "score", "source_clip", "text"])
        for i, r in enumerate(rows[: args.top], 1):
            name = f"ref_{i:02d}.wav"
            shutil.copy2(r["wav"], out / name)
            w.writerow([name, f"{r['seconds']:.1f}", r["similarity"], f"{r['score']:.2f}",
                        r["clip_id"], r["text"]])
            print(f"  {name}  {r['seconds']:4.1f}s  match {float(r['similarity']):.2f}  "
                  f"\"{r['text'][:70]}{'...' if len(r['text']) > 70 else ''}\"")
    print(f"\nCopied to {out.resolve()}")
    print("Listen to them and pick one whose energy fits the line you're generating:")
    print("calm explaining, excited, or talking to camera. That clip's tone carries over.")


if __name__ == "__main__":
    main()