"""Run V3 face inference on an arbitrary audio file (e.g. a song).

Unlike `models/v3_face/infer.py` (which requires a pre-baked .npz with
ground-truth target arrays for L1 metrics) and `infer_e2e.py` (which expects
per-turn segmented audio + scenario lookup in the JSONL splits), this script
runs end-to-end on any single audio file with a user-chosen emotion label.

  python scripts/infer_song.py \\
      --audio "data/Coffee Goes Cold - AI Music.mp3" \\
      --sid song_coffee \\
      --emotion struggle

Output:
  data/viewer/<sid>_song_dataset.json     # blendshapes + meta
  data/viewer/<sid>_song_dataset.mp3      # symlink to source audio

To inspect the result, drop into the same viewer dir and either build a
mini player or temporarily hardcode the sid in blendshape-player-natural.html.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import librosa
import numpy as np
import torch

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

from scripts.compiler.constants import ARKIT_52_NAMES  # noqa: E402
from scripts.compiler.data_pipeline import EMOTION_LABELS, FPS, mel_features  # noqa: E402
from models.v3_face.infer import (  # noqa: E402
    crisp_mouth, smooth_brows, smooth_eye_squint, smooth_eye_wide,
    smooth_eye_look, inject_blinks, load_model,
)
from models.v3_face.infer_e2e import build_cond_smoothed  # noqa: E402

DEFAULT_CKPT = PROJECT_ROOT / "models" / "v3_face" / "checkpoints" / "best_expression_v18m.pt"
DEFAULT_VIEWER_DIR = PROJECT_ROOT / "data" / "viewer"
DEFAULT_EMOTION_LABELS_JSON = PROJECT_ROOT / "data" / "emotion" / "emotion_labels.json"


def default_vad_for(emotion: str) -> np.ndarray:
    """Look up the VAD anchor for a label from emotion_labels.json."""
    with open(DEFAULT_EMOTION_LABELS_JSON) as f:
        labels = json.load(f)["labels"]
    for L in labels:
        if L["name"] == emotion:
            return np.array(L["default_vad"], dtype=np.float32)
    raise ValueError(f"emotion {emotion!r} not found in {DEFAULT_EMOTION_LABELS_JSON}")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--audio", type=Path, required=True, help="Source audio (mp3/wav).")
    ap.add_argument("--sid", required=True, help="Output scenario ID (used in filenames).")
    ap.add_argument("--emotion", required=True, choices=EMOTION_LABELS,
                    help="Single emotion label applied across the whole clip.")
    ap.add_argument("--vad", nargs=3, type=float, default=None, metavar=("V", "A", "D"),
                    help="Override VAD vector (3 floats in [-1,1]). "
                         "Default: anchor from emotion_labels.json. "
                         "Useful for 'sad but passionate' — bump arousal manually.")
    ap.add_argument("--ckpt", type=Path, default=DEFAULT_CKPT)
    ap.add_argument("--viewer_dir", type=Path, default=DEFAULT_VIEWER_DIR)
    ap.add_argument("--device", default="cuda:0")
    # Post-processing (mirrors the locked POST_PROC settings used by serve_live).
    ap.add_argument("--no-crisp", action="store_true")
    ap.add_argument("--crisp-threshold", type=float, default=0.3)
    ap.add_argument("--crisp-scale", type=float, default=1.0)
    ap.add_argument("--crisp-sigma", type=float, default=1.0)
    ap.add_argument("--crisp-mouthclose-sigma", type=float, default=1.0)
    ap.add_argument("--no-brow-smooth", action="store_true")
    ap.add_argument("--brow-min-cutoff", type=float, default=1.5)
    ap.add_argument("--brow-beta", type=float, default=0.01)
    ap.add_argument("--brow-d-cutoff", type=float, default=1.0)
    ap.add_argument("--no-blinks", action="store_true")
    ap.add_argument("--blink-interval", type=float, default=4.0)
    ap.add_argument("--blink-expressive-cap", type=float, default=0.5)
    args = ap.parse_args()

    if not args.audio.exists():
        ap.error(f"audio not found: {args.audio}")
    args.viewer_dir.mkdir(parents=True, exist_ok=True)

    # Resolve VAD vector.
    vad = np.array(args.vad, dtype=np.float32) if args.vad is not None else default_vad_for(args.emotion)
    print(f"emotion={args.emotion}  vad={vad.tolist()}")

    device = torch.device(args.device)
    print(f"loading checkpoint {args.ckpt}")
    model, _cfg = load_model(args.ckpt, device)

    print(f"loading audio {args.audio}")
    wav, sr = librosa.load(str(args.audio), sr=16000, mono=True)
    dur_s = len(wav) / sr
    print(f"  duration={dur_s:.2f}s")
    if dur_s < 0.1:
        ap.error("audio too short")

    mel = mel_features(wav, sr=sr, fps=FPS)
    T = mel.shape[0]
    print(f"  mel frames={T}  ({T/FPS:.2f}s @ {FPS}fps)")

    # Treat the whole clip as one "turn" with the chosen emotion + VAD.
    cond = build_cond_smoothed(
        [args.emotion],
        np.expand_dims(vad, axis=0),
        [T],
        vad_smooth_sigma=0.0,
    )

    audio_t = torch.from_numpy(mel.astype(np.float32)).unsqueeze(0).to(device)
    cond_t = torch.from_numpy(cond.astype(np.float32)).unsqueeze(0).to(device)

    print("forward pass…")
    with torch.no_grad():
        pred = model(audio_t, cond_t).squeeze(0).cpu().numpy().astype(np.float32)
    print(f"  pred frames={pred.shape[0]}  channels={pred.shape[1]}")

    if not args.no_crisp:
        pred = crisp_mouth(pred,
                           threshold=args.crisp_threshold,
                           scale=args.crisp_scale,
                           pre_smooth_sigma=args.crisp_sigma,
                           mouth_close_sigma=args.crisp_mouthclose_sigma)
    if not args.no_brow_smooth:
        pred = smooth_brows(pred,
                            min_cutoff=args.brow_min_cutoff,
                            beta=args.brow_beta,
                            d_cutoff=args.brow_d_cutoff)
    if not args.no_blinks:
        pred = inject_blinks(pred,
                             scenario_id=args.sid,
                             mean_interval_s=args.blink_interval,
                             expressive_cap=args.blink_expressive_cap)

    new_base = f"{args.sid}_song_dataset"
    viewer_json = {
        "scenario_id": new_base,
        "fps": FPS,
        "num_frames": int(pred.shape[0]),
        "names": ARKIT_52_NAMES,
        "turns": [{
            "turn_idx": 0,
            "emotion": args.emotion,
            "vad": vad.tolist(),
            "text": f"[song: {args.audio.name}]",
            "speaker": "song",
        }],
        "blendshapes": np.round(pred, 4).tolist(),
    }
    out_json = args.viewer_dir / f"{new_base}.json"
    out_json.write_text(json.dumps(viewer_json, ensure_ascii=False))

    # Symlink the source mp3 next to the JSON so any player can find it.
    out_mp3 = args.viewer_dir / f"{new_base}.mp3"
    if out_mp3.exists() or out_mp3.is_symlink():
        out_mp3.unlink()
    out_mp3.symlink_to(args.audio.resolve())

    print(f"\nwrote {out_json}")
    print(f"      {out_mp3} -> {args.audio.resolve()}")


if __name__ == "__main__":
    main()