a2f-mlx 0.1.0

//! a2f-mlx — Audio2Face-3D v2.3.1 (James) as a hand-written MLX graph (Rust).
//!
//! Mirrors the validated Python port (~/dev/a2f-mlx/a2f_mlx.py), which itself
//! mirrors soniqo/speech-swift's Audio2Face3DMLXRuntime.swift. Weights come
//! from `tools/export_mlx_weights.py` — extracted from NVIDIA's official
//! network.onnx (keep-NVIDIA provenance; no third-party repack).
//!
//! Graph:
//!   audio branch : wav2vec2-style conv frontend + 4-layer transformer
//!   freq branch  : 25 hann windows → 32-lag autocorrelation → 5× conv2d
//!   join         : 4× conv2d stride(2,1) + emotion broadcast → dual MLP heads
//!   out          : 169 = 140 skin | 10 tongue | 15 jaw | 4 eye

use std::collections::HashMap;
use std::path::Path;

use anyhow::{bail, Context, Result};
use mlx_rs::{
    ops::{self, indexing::IndexOp},
    Array,
};
use safetensors::SafeTensors;

use a2f_core::{
    hann640, Emotion, EXPLICIT_EMO, IMPLICIT_EMO, N_EMO, OUT_COUNT, SKIN, TONGUE,
};

pub const HEADS: usize = 12;
const SQRT2: f32 = std::f32::consts::SQRT_2;
const EPS: f32 = 1e-5;

/// Tensor-name families ("variants") present in the official ONNX. For this
/// bundle the live one is `"2230"` — decided by ONNX parity, not by name.
pub const VARIANT_2230: Variant = Variant {
    emo_proj: "2230",
    fe_scale: "2235",
    fe_bias: "2236",
    feat_proj: "2237",
    pos_conv: "2240",
    audio_map: "2289",
    layers: [
        [2241, 2244, 2247, 2250, 2251, 2252],
        [2253, 2256, 2259, 2262, 2263, 2264],
        [2265, 2268, 2271, 2274, 2275, 2276],
        [2277, 2280, 2283, 2286, 2287, 2288],
    ],
};

#[derive(Clone, Copy)]
pub struct Variant {
    pub emo_proj: &'static str,
    pub fe_scale: &'static str,
    pub fe_bias: &'static str,
    pub feat_proj: &'static str,
    pub pos_conv: &'static str,
    pub audio_map: &'static str,
    /// [q, k, v, o, ff1, ff2] per transformer layer (MatMul initializer names).
    pub layers: [[u32; 6]; 4],
}

fn p(layer: usize, rest: &str) -> String {
    format!("wave2vec_model.encoder.transformer.layers.{layer}.{rest}")
}

/// Load safetensors into MLX arrays with layouts baked for the graph:
/// conv1d torch [O,I,K] → [O,K,I]; conv2d torch [O,I,kH,kW] → [O,kH,kW,I];
/// torch-style matmul [out,in] → [in,out]. Named as-is otherwise.
fn load_weights(bytes: &[u8], v: &Variant) -> Result<HashMap<String, Array>> {
    let st = SafeTensors::deserialize(bytes)?;
    let mut m = HashMap::new();
    let get = |m: &mut HashMap<String, Array>, name: &str, data: &[u8], shape: &[usize]| {
        let f: Vec<f32> = data
            .chunks_exact(4)
            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect();
        let sh: Vec<i32> = shape.iter().map(|&d| d as i32).collect();
        m.insert(name.to_string(), Array::from_slice(&f, &sh));
    };
    for (name, view) in st.iter() {
        match view.dtype() {
            safetensors::Dtype::F32 => {
                get(&mut m, name, view.data(), view.shape())
            }
            other => bail!("tensor `{name}` has dtype {other:?}, need F32 — re-export"),
        }
    }
    // bake layouts (transpose once at load, not per call)
    let conv1d = [
        "wave2vec_model.feature_extractor.conv_layers.0.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.1.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.2.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.3.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.4.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.5.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.6.conv.weight",
        v.pos_conv,
    ];
    for name in conv1d {
        let a = m.remove(name).unwrap();
        m.insert(name.into(), a.transpose_axes(&[0, 2, 1])?);
    }
    let conv2d = ["freq1.weight", "freq2.weight", "freq3.weight", "freq4.weight", "freq5.weight",
        "time1.weight", "time2.weight", "time3.weight", "time4.weight"];
    for name in conv2d {
        let a = m.remove(name).unwrap();
        m.insert(name.into(), a.transpose_axes(&[0, 2, 3, 1])?);
    }
    let outin = ["time5.weight", "out_mapping.weight", "out_mapping1.weight",
        "out_mapping_tongue.weight", "out_mapping1_tongue.weight"];
    for name in outin {
        let a = m.remove(name).unwrap();
        m.insert(name.into(), a.transpose()?);
    }
    for layer in v.layers.iter() {
        for id in layer {
            let name = id.to_string();
            let a = m.remove(&name).unwrap();
            // MatMul initializers are [in, out] already — keep.
            m.insert(name, a);
        }
    }
    Ok(m)
}

pub struct A2FMlx {
    w: HashMap<String, Array>,
    v: Variant,
    input_strength: f32,
    default_emotion: Emotion,
    hann: Vec<f32>,
    combine: Vec<f32>, // normalized combine weights
}

impl A2FMlx {
    pub fn load(bundle_dir: &Path, input_strength: Option<f32>) -> Result<Self> {
        Self::load_with_variant(bundle_dir, VARIANT_2230, input_strength)
    }

    pub fn load_with_variant(
        bundle_dir: &Path,
        variant: Variant,
        input_strength: Option<f32>,
    ) -> Result<Self> {
        let bytes = std::fs::read(bundle_dir.join("audio2face3d.safetensors"))
            .with_context(|| format!("read {}", bundle_dir.join("audio2face3d.safetensors").display()))?;
        let w = load_weights(&bytes, &variant)?;

        let mc: serde_json::Value = serde_json::from_str(
            &std::fs::read_to_string(bundle_dir.join("model_config.json"))
                .unwrap_or_else(|_| "{}".into()),
        )?;
        let input_strength = input_strength.unwrap_or(
            mc.get("config")
                .and_then(|c| c.get("input_strength"))
                .and_then(|v| v.as_f64())
                .unwrap_or(1.3) as f32,
        );
        let default_emotion = Emotion::from_raw(
            &std::fs::read(bundle_dir.join("default_emotion.f32"))
                .context("default_emotion.f32 missing from bundle")?,
        )?;

        // combine_weights / Σ (python: cw = cw / cw.sum())
        let cw_arr = w.get("combine_weights").context("combine_weights")?;
        let cw: Vec<f32> = cw_arr.as_slice::<f32>().to_vec();
        let sum: f32 = cw.iter().sum();
        let combine = cw.into_iter().map(|c| c / sum).collect();

        Ok(Self {
            w,
            v: variant,
            input_strength,
            default_emotion,
            hann: hann640(),
            combine,
        })
    }

    pub fn default_emotion(&self) -> &Emotion {
        &self.default_emotion
    }

    fn t(&self, name: &str) -> Result<Array> {
        self.w
            .get(name)
            .cloned()
            .with_context(|| format!("missing tensor `{name}`"))
    }

    fn bias1(&self, name: &str) -> Result<Array> {
        self.t(name)
    }

    fn w3(&self, id: u32) -> Result<Array> {
        self.t(&id.to_string())
    }

    /// x: NCL [B,C,T]; weight pre-baked [O,K,I/g]; returns NCL [B,O,T].
    fn conv1d_ncl(&self, x: &Array, weight: &Array, bias: Option<&Array>,
                  stride: i32, padding: i32, groups: i32) -> Result<Array> {
        let y = ops::conv1d(x.transpose_axes(&[0, 2, 1])?, weight, stride, padding, 1, groups)?;
        let y = y.transpose_axes(&[0, 2, 1])?; // back to NCL
        match bias {
            Some(b) => {
                let b = b.reshape(&[1, -1, 1])?;
                Ok(y + b)
            }
            None => Ok(y),
        }
    }

    /// x: NCHW; weight pre-baked [O,kH,kW,I]; returns NCHW.
    fn conv2d_nchw(&self, x: &Array, weight: &Array, bias: &Array,
                   stride: (i32, i32), padding: (i32, i32)) -> Result<Array> {
        let y = ops::conv2d(x.transpose_axes(&[0, 2, 3, 1])?, weight, stride, padding, (1, 1), 1)?;
        let y = y.transpose_axes(&[0, 3, 1, 2])? + bias.reshape(&[1, -1, 1, 1])?;
        Ok(y)
    }

    fn gelu(&self, x: &Array) -> Result<Array> {
        let t = x / Array::from_f32(SQRT2);
        let t = ops::erf(&t)?;
        Ok(x * (t + Array::from_f32(1.0)) * Array::from_f32(0.5))
    }

    fn layer_norm(&self, x: &Array, wname: &str, bname: &str) -> Result<Array> {
        let mean = x.mean_axis(-1, Some(true))?;
        let centered = x - &mean;
        let var = (&centered * &centered).mean_axis(-1, Some(true))?;
        let norm = &centered / (&var + Array::from_f32(EPS)).sqrt()?;
        Ok(&norm * self.t(wname)? + self.t(bname)?)
    }

    /// wav2vec2-style audio branch. Input [B,1,T]; returns [B,T',C] combined.
    fn audio_branch(&self, x: &Array) -> Result<Array> {
        let fe = |i: u32| {
            format!("wave2vec_model.feature_extractor.conv_layers.{i}.conv.weight")
        };
        let mut x = self.conv1d_ncl(x, &self.t(&fe(0))?, None, 5, 0, 1)?;
        // instance norm with learned scale/bias, then gelu
        let mean = x.mean_axis(-1, Some(true))?;
        let centered = &x - &mean;
        let var = (&centered * &centered).mean_axis(-1, Some(true))?;
        x = &centered / (&var + Array::from_f32(EPS)).sqrt()?
            * self.t(self.v.fe_scale)?.reshape(&[1, -1, 1])?
            + self.t(self.v.fe_bias)?.reshape(&[1, -1, 1])?;
        x = self.gelu(&x)?;
        for i in 1..7 {
            x = self.conv1d_ncl(&x, &self.t(&fe(i))?, None, 2, 0, 1)?;
            x = self.gelu(&x)?;
        }

        // feature projection: NLC + LN + linear
        let mut x = x.transpose_axes(&[0, 2, 1])?; // [B,T,C]
        x = self.layer_norm(
            &x,
            "wave2vec_model.encoder.feature_projection.layer_norm.weight",
            "wave2vec_model.encoder.feature_projection.layer_norm.bias",
        )?;
        x = x.matmul(&self.t(self.v.feat_proj)?)?
            + self.bias1("wave2vec_model.encoder.feature_projection.projection.bias")?;

        // positional conv (groups=16, pad 64, truncate) + gelu, residual
        let t_len = x.shape()[1];
        let pos = self.conv1d_ncl(
            &x.transpose_axes(&[0, 2, 1])?,
            &self.t(self.v.pos_conv)?,
            Some(&self.bias1("wave2vec_model.encoder.transformer.pos_conv_embed.conv.bias")?),
            1, 64, 16,
        )?;
        let pos = pos.index((.., .., ..(t_len as i32)));
        x = x + self.gelu(&pos)?.transpose_axes(&[0, 2, 1])?;

        x = self.layer_norm(
            &x,
            "wave2vec_model.encoder.transformer.layer_norm.weight",
            "wave2vec_model.encoder.transformer.layer_norm.bias",
        )?;

        let shape = x.shape();
        let (b, t_len, hidden) = (shape[0], shape[1], shape[2]);
        let head_dim = hidden / HEADS as i32;
        let mut outputs = Vec::with_capacity(4);
        for (li, ids) in self.v.layers.iter().enumerate() {
            let pre = p(li, "attention");
            let q = (x.matmul(&self.w3(ids[0])?)? + self.bias1(&format!("{pre}.q_proj.bias"))?)
                .reshape(&[b, t_len, HEADS as i32, head_dim])?
                .transpose_axes(&[0, 2, 1, 3])?;
            let k = (x.matmul(&self.w3(ids[1])?)? + self.bias1(&format!("{pre}.k_proj.bias"))?)
                .reshape(&[b, t_len, HEADS as i32, head_dim])?
                .transpose_axes(&[0, 2, 3, 1])?;
            let v = (x.matmul(&self.w3(ids[2])?)? + self.bias1(&format!("{pre}.v_proj.bias"))?)
                .reshape(&[b, t_len, HEADS as i32, head_dim])?
                .transpose_axes(&[0, 2, 1, 3])?;
            let mut scores = (&q * Array::from_f32(1.0 / (head_dim as f32).sqrt())).matmul(&k)?;
            let smax = scores.max_axes(&[-1], Some(true))?;
            scores = scores - smax;
            let ctx = ops::softmax_axes(&scores, &[-1], None)?
                .matmul(&v)?
                .transpose_axes(&[0, 2, 1, 3])?
                .reshape(&[b, t_len, hidden])?;
            x = x + ctx.matmul(&self.w3(ids[3])?)?
                + self.bias1(&format!("{pre}.out_proj.bias"))?;
            x = self.layer_norm(&x, &p(li, "layer_norm.weight"), &p(li, "layer_norm.bias"))?;
            let ff = x.matmul(&self.w3(ids[4])?)?
                + self.bias1(&p(li, "feed_forward.intermediate_dense.bias"))?;
            let ff = self.gelu(&ff)?;
            x = x + ff.matmul(&self.w3(ids[5])?)?
                + self.bias1(&p(li, "feed_forward.output_dense.bias"))?;
            x = self.layer_norm(
                &x,
                &p(li, "final_layer_norm.weight"),
                &p(li, "final_layer_norm.bias"),
            )?;
            outputs.push(x.clone());
        }
        // weighted combine (weights normalized at load)
        let mut acc = &outputs[0] * Array::from_f32(self.combine[0]);
        for (o, c) in outputs.iter().zip(self.combine.iter()).skip(1) {
            acc = acc + o * Array::from_f32(*c);
        }
        Ok(acc)
    }

    /// Autocorrelation frequency branch. Input [B,1,T] (raw audio); [B,c,h,w] NCHW.
    fn freq_branch(&self, x: &Array) -> Result<Array> {
        let scaled = x * self.t("in_scale_autocorr")?;
        let hann = Array::from_slice(&self.hann, &[1, 1, 1, 640]);
        let mut windows = Vec::with_capacity(25);
        for i in 0..25 {
            let w = scaled.index((
                ..,
                ..,
                (i * 320) as i32..((i * 320 + 640) as i32),
            ));
            windows.push(w.reshape(&[w.shape()[0], 1, 1, 640])?);
        }
        let frames = ops::concatenate(&windows, 2)?;
        let mean = frames.mean_axis(-1, Some(true))?;
        let centered = (frames - mean) * hann;
        let mut corrs = Vec::with_capacity(32);
        for lag in 0..32 {
            let prod = if lag == 0 {
                &centered * &centered
            } else {
                let lhs = centered.index((.., .., .., ..(640 - lag)));
                let rhs = centered.index((.., .., .., (lag as i32)..));
                &lhs * &rhs
            };
            corrs.push(prod.sum_axes(&[-1], Some(true))?);
        }
        let mut out = ops::concatenate(&corrs, -1)? * Array::from_f32(1.0 / 640.0);
        let spec = [
            ("freq1", (1, 2), (0, 1)),
            ("freq2", (1, 2), (0, 1)),
            ("freq3", (1, 2), (0, 1)),
            ("freq4", (1, 2), (0, 1)),
            ("freq5", (1, 1), (0, 0)),
        ];
        for (name, stride, padding) in spec {
            let y = self.conv2d_nchw(
                &out,
                &self.t(&format!("{name}.weight"))?,
                &self.t(&format!("{name}.bias"))?,
                stride,
                padding,
            )?;
            out = ops::maximum(&y, Array::from_f32(0.0))? * Array::from_f32(SQRT2);
        }
        Ok(out)
    }

    /// window: BUFFER_LEN f32 samples; emotion: N_EMO f32. Returns 169 coefs.
    pub fn coefficients(&self, window: &[f32], emotion: &Emotion) -> Result<Vec<f32>> {
        use a2f_core::BUFFER_LEN;
        assert_eq!(window.len(), BUFFER_LEN);
        let scaled: Vec<f32> = window.iter().map(|s| s * self.input_strength).collect();
        let audio = Array::from_slice(&scaled, &[1, 1, BUFFER_LEN as i32]);
        let emo = Array::from_slice(&emotion.0, &[1, 1, N_EMO as i32]);
        let out = self.forward(&audio, &emo)?;
        Ok(out.as_slice::<f32>().to_vec())
    }

    /// audio [1,1,T]; emotion [1,1,26] → [1,1,169]
    pub fn forward(&self, audio: &Array, emotion: &Array) -> Result<Array> {
        let implicit = emotion.index((.., .., ..(IMPLICIT_EMO as i32)));
        let explicit = emotion.index((.., .., (IMPLICIT_EMO as i32)..));
        let projected = explicit.matmul(&self.t(self.v.emo_proj)?)?
            + self.bias1("emo_linear1.bias")?;

        let freq = self.freq_branch(audio)?;
        let aud = self.audio_branch(audio)?;

        let head = ops::maximum(
            &(aud.matmul(&self.t(self.v.audio_map)?)?
                + self.bias1("audio_feature_map.bias")?),
            Array::from_f32(0.0),
        )?;
        let hs = head.shape();
        // [B,1,C,T] → [B,1,C,T,1]
        let aud5 = head
            .reshape(&[hs[0], 1, hs[1], hs[2], 1])?
            .transpose_axes(&[0, 1, 3, 2, 4])?; // [B,1,T,C,1]
        let fs = freq.shape();
        let freq5 = freq.reshape(&[fs[0], 1, fs[1], fs[2], fs[3]])?;

        let joined = ops::concatenate(&[aud5, freq5], 2)?;
        let js = joined.shape();
        let mut x = joined.reshape(&[js[0] * js[1], js[2], js[3], js[4]])?;

        for name in ["time1", "time2", "time3", "time4"] {
            let y = self.conv2d_nchw(
                &x,
                &self.t(&format!("{name}.weight"))?,
                &self.t(&format!("{name}.bias"))?,
                (2, 1),
                (1, 0),
            )?;
            x = ops::maximum(&y, Array::from_f32(0.0))? * Array::from_f32(SQRT2);
            let xs = x.shape();
            let (b, _, h, w) = (xs[0], xs[1], xs[2], xs[3]);
            let imp = ops::broadcast_to(
                &implicit.reshape(&[b, IMPLICIT_EMO as i32, 1, 1])?,
                &[b, IMPLICIT_EMO as i32, h, w],
            )?;
            let prj = ops::broadcast_to(
                &projected.reshape(&[b, 8, 1, 1])?,
                &[b, 8, h, w],
            )?;
            x = ops::concatenate(&[x, imp, prj], 1)?;
        }

        let xs = x.shape();
        let flat = x.reshape(&[xs[0], xs[1] * xs[2] * xs[3]])?;
        let time5 = (flat.matmul(&self.t("time5.weight")?)?
            + self.bias1("time5.bias")?)
            .reshape(&[xs[0], 1, 256])?;
        let time5 = ops::maximum(&time5, Array::from_f32(0.0))? * Array::from_f32(SQRT2);
        let head_in = ops::concatenate(&[time5, implicit, projected], 2)?;
        let his = head_in.shape();
        let head_flat = head_in.reshape(&[his[0] * his[1], his[2]])?;

        let face = (head_flat.matmul(&self.t("out_mapping.weight")?)?
            + self.bias1("out_mapping.bias")?)
            .reshape(&[his[0] * his[1], a2f_core::FACE_COUNT as i32])?;
        let face = (face.matmul(&self.t("out_mapping1.weight")?)?
            + self.bias1("out_mapping1.bias")?)
            .reshape(&[his[0], his[1], a2f_core::FACE_COUNT as i32])?;
        let tongue = (head_flat.matmul(&self.t("out_mapping_tongue.weight")?)?
            + self.bias1("out_mapping_tongue.bias")?)
            .reshape(&[his[0] * his[1], TONGUE as i32])?;
        let tongue = (tongue.matmul(&self.t("out_mapping1_tongue.weight")?)?
            + self.bias1("out_mapping1_tongue.bias")?)
            .reshape(&[his[0], his[1], TONGUE as i32])?;

        let skin = face.index((.., .., ..(SKIN as i32)));
        let jaw_eye = face.index((.., .., (SKIN as i32)..));
        Ok(ops::concatenate(&[skin, tongue, jaw_eye], 2)?)
    }

    pub fn input_strength(&self) -> f32 {
        self.input_strength
    }

    pub fn emotion_kind(&self) -> (usize, usize) {
        (IMPLICIT_EMO, EXPLICIT_EMO)
    }

    pub fn out_count(&self) -> usize {
        OUT_COUNT
    }
}

impl A2FMlx {
    /// Parity helper: run one fixture window/emotion pair, return coefficients.
    pub fn run_fixture(&self, window: &[f32], emotion: &[f32]) -> Result<Vec<f32>> {
        self.coefficients(window, &Emotion(emotion.to_vec()))
    }
}