"""Standalone V3 face inference for the avatar E2E pipeline.

Given an assistant utterance (16k mono wav) + an emotion tag, produce the 52
ARKit blendshape timeline the avatar viewer consumes — WITHOUT V3's TTS or
KlueTeacher steps. In the E2E, OpenAI Realtime supplies the audio and the LLM
supplies the emotion tag, so we feed mel + emotion/VAD straight into the V3
face model. This is the audio-driven subset of
`face/animasync-face-v3/models/v3_face/serve_live.py`.

Pipeline:
    wav ─► mel_features                         (T, 80)   [raw log-mel, no norm — matches V3 npz "audio"]
    emotion ─► emotion_vad_anchors[emo][lvl].vad
            ─► build_conditioning               (T, 19)   [16 one-hot + 3 VAD]
    V3FaceModel(mel, cond)                       (T, 52)
    ─► crisp_mouth ─► smooth_brows ─► inject_blinks ─► clip[0,1]

Runs in the repo `.venv` (same interpreter as the body model — imports verified).

    python face_v3_infer.py --wav <path> --emotion joy --out out.json
"""
import argparse
import json
import os
import sys

import numpy as np
import librosa
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
# Localised like the body scripts (no pluto /home/ubuntu or /data/recover paths).
FACE = os.environ.get("FACE_V3_DIR") or os.path.abspath(
    os.path.join(HERE, "..", "face", "animasync-face-v3"))
sys.path.insert(0, FACE)

from scripts.compiler.constants import ARKIT_52_NAMES                          # noqa: E402
from scripts.compiler.data_pipeline import (                                    # noqa: E402
    EMOTION_LABELS, mel_features, build_conditioning,
)
from models.v3_face.infer import (                                             # noqa: E402
    load_model, crisp_mouth, smooth_brows, inject_blinks,
)

DEFAULT_CKPT = os.path.join(FACE, "models", "v3_face", "checkpoints", "best_expression_v14.pt")
ANCHORS = os.path.join(FACE, "data", "emotion", "emotion_vad_anchors.json")
FPS = 30
SR = 16000

# V3's shipped live post-processing (serve_live.py POST_PROC).
POST = dict(crisp_threshold=0.3, crisp_scale=1.0, crisp_sigma=1.3, crisp_mouthclose_sigma=1.0,
            brow_min_cutoff=2.0, brow_beta=0.01, brow_d_cutoff=1.0,
            blink_interval=6.0, blink_expressive_cap=0.5)


def emotion_to_vad(emotion, level=4):
    """Look up the affective VAD anchor (valence/arousal/dominance) for an emotion.

    Same selection gen_v3_face.py uses: the anchor entry whose intensity `level`
    is closest to the requested one.
    """
    anchors = json.load(open(ANCHORS))["anchors"]
    lv = anchors.get(emotion, anchors["neutral"])
    pk = min(lv, key=lambda a: abs(a["level"] - level))
    return np.asarray(pk["vad"], np.float32)


def load_face_model(ckpt=DEFAULT_CKPT, device=None):
    """Load once, reuse across turns (the E2E server keeps this warm)."""
    device = device or torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model, cfg = load_model(ckpt, device)
    return model, device


def render_face(wav_path, emotion, level=4, model_dev=None, ckpt=DEFAULT_CKPT,
                scenario_id="e2e"):
    """wav + emotion -> (T, 52) blendshapes in ARKit_52 order, post-processed."""
    wav, _ = librosa.load(wav_path, sr=SR, mono=True)
    return render_face_samples(wav, emotion, level=level, model_dev=model_dev,
                               ckpt=ckpt, scenario_id=scenario_id)


def render_face_samples(wav, emotion, level=4, model_dev=None, ckpt=DEFAULT_CKPT,
                        scenario_id="e2e"):
    """In-memory mono 16k PCM + emotion -> post-processed ARKit-52 frames."""
    if model_dev is None:
        model_dev = load_face_model(ckpt)
    model, device = model_dev

    wav = np.asarray(wav, dtype=np.float32).reshape(-1)
    if wav.size < 1024:
        wav = np.pad(wav, (0, 1024 - wav.size))
    mel = mel_features(wav, sr=SR, fps=FPS)                      # (T, 80) raw log-mel
    T = mel.shape[0]
    vad = emotion_to_vad(emotion, level)
    cond = build_conditioning(emotion, vad.tolist(), T)          # (T, 19)

    with torch.no_grad():
        audio = torch.from_numpy(mel).unsqueeze(0).to(device)
        c = torch.from_numpy(cond).unsqueeze(0).to(device)
        bs = model(audio, c).squeeze(0).cpu().numpy().astype(np.float32)   # (T, 52)

    bs = np.clip(bs, 0.0, 1.0)
    bs = crisp_mouth(bs, threshold=POST["crisp_threshold"], scale=POST["crisp_scale"],
                     pre_smooth_sigma=POST["crisp_sigma"],
                     mouth_close_sigma=POST["crisp_mouthclose_sigma"])
    bs = smooth_brows(bs, min_cutoff=POST["brow_min_cutoff"], beta=POST["brow_beta"],
                      d_cutoff=POST["brow_d_cutoff"], fps=FPS)
    bs = inject_blinks(bs, scenario_id=scenario_id, mean_interval_s=POST["blink_interval"],
                       expressive_cap=POST["blink_expressive_cap"], fps=FPS)
    return np.clip(bs, 0.0, 1.0).astype(np.float32), T


class FaceStream:
    """Causal V3 wrapper that returns frames aligned to each body chunk."""

    def __init__(self, model_dev, scenario_id, level=4, blend=0):
        self.model_dev = model_dev
        self.scenario_id = scenario_id
        self.level = level
        self.blend = blend
        self.pcm = np.zeros(0, dtype=np.float32)
        self._last = None
        self.frames_emitted = 0

    def push_pcm(self, pcm_chunk, emotion, n_frames, frame_start=None):
        """Append PCM and return the exact absolute frame range emitted by body.

        ``librosa``'s centered STFT produces one more mel frame than the
        duration-derived 30fps body timeline (12,800 samples -> 25 mel frames,
        while body emits frames 0..23). Selecting ``bs[-n_frames:]`` therefore
        led face by one frame. Track the body frame cursor explicitly and slice
        ``[frame_start:frame_start + n_frames]`` instead.

        A request may also contain more PCM than body consumed into complete
        steps. Limit inference to the sample boundary represented by the body
        frame end so that pending future audio cannot leak into this packet's
        post-processing.
        """
        n_frames = int(n_frames)
        frame_start = self.frames_emitted if frame_start is None else int(frame_start)
        if n_frames < 0:
            raise ValueError("n_frames must be non-negative")
        if frame_start != self.frames_emitted:
            raise ValueError(
                f"face frame discontinuity: expected {self.frames_emitted}, got {frame_start}")

        pcm = np.asarray(pcm_chunk, dtype=np.float32).reshape(-1)
        if pcm.size:
            self.pcm = np.concatenate([self.pcm, pcm])
        if n_frames <= 0:
            return np.zeros((0, len(ARKIT_52_NAMES)), dtype=np.float32)

        # V3 is a causal TCN. Reusing the warm model over the accumulated
        # utterance is an incremental inference boundary without changing the
        # trained network. Do not include samples beyond body's emitted cursor.
        frame_end = frame_start + n_frames
        represented_samples = int(round(frame_end * SR / FPS))
        inference_pcm = self.pcm[:min(len(self.pcm), represented_samples)]
        bs, _ = render_face_samples(inference_pcm, emotion, level=self.level,
                                    model_dev=self.model_dev,
                                    scenario_id=self.scenario_id)
        if len(bs) < frame_end:
            bs = np.pad(bs, ((0, frame_end - len(bs)), (0, 0)), mode="edge")
        out = bs[frame_start:frame_end].copy()

        if self._last is not None and self.blend > 0:
            k = min(self.blend, len(out))
            for i in range(k):
                a = (i + 1) / (k + 1)
                out[i] = (1.0 - a) * self._last + a * out[i]
        self._last = out[-1].copy()
        self.frames_emitted = frame_end
        return np.clip(out, 0.0, 1.0).astype(np.float32)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--wav", required=True)
    ap.add_argument("--emotion", default="neutral",
                    help=f"one of: {', '.join(EMOTION_LABELS)}")
    ap.add_argument("--level", type=int, default=4, help="VAD intensity 1-5")
    ap.add_argument("--ckpt", default=DEFAULT_CKPT)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    if a.emotion not in EMOTION_LABELS:
        sys.exit(f"unknown emotion {a.emotion!r} — must be one of {EMOTION_LABELS}")

    sid = os.path.splitext(os.path.basename(a.wav))[0]
    bs, T = render_face(a.wav, a.emotion, a.level, ckpt=a.ckpt, scenario_id=sid)

    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    json.dump({"scenario_id": os.path.splitext(os.path.basename(a.out))[0],
               "fps": FPS, "num_frames": int(T), "names": ARKIT_52_NAMES,
               "blendshapes": [[round(float(v), 4) for v in bs[t]] for t in range(T)]},
              open(a.out, "w"))
    print(f"face: {a.emotion:10s} {T:4d}f  "
          f"jawOpenMax={bs[:, 24].max():.2f}  browInnerMax={bs[:, 2].max():.2f}  "
          f"smileMax={bs[:, 44].max():.2f} -> {a.out}")


if __name__ == "__main__":
    main()