mozu-audio 0.3.0

//! Haas recognition using macOS Core ML. Models remain loaded between independent utterances.
//! Create and use an engine on a single worker thread; no unsafe Send/Sync guarantees are added.

pub mod audio;
pub mod batch;
mod decoder;
pub mod models;
pub mod recording;
pub mod result;
pub mod service;
pub mod streaming;
mod tensor;
pub mod vad;
pub mod vocabulary;
pub mod window;
#[cfg(test)]
use decoder::label_probability;
use decoder::{DecoderModels, decode};
#[cfg(test)]
use tensor::flat_offset;
use tensor::{array_shape, multi_array, read_array, read_floating_tensor};
use vocabulary::{CompiledVocabulary, Vocabulary};
use window::TimedToken;

use objc2::{
    AnyThread,
    rc::{Retained, autoreleasepool},
    runtime::{AnyObject, ProtocolObject},
};
use objc2_core_ml::{
    MLComputeUnits, MLDictionaryFeatureProvider, MLFeatureDescription, MLFeatureProvider,
    MLFeatureValue, MLModel, MLModelConfiguration, MLMultiArray, MLMultiArrayDataType,
};
use objc2_foundation::{NSDictionary, NSError, NSString, NSURL};
use rustfft::{FftPlanner, num_complex::Complex32};
use serde::{Deserialize, Serialize};
use std::{
    fs,
    path::{Path, PathBuf},
    time::Instant,
};

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("I/O failed: path={path}, error={source}")]
    Io {
        path: PathBuf,
        source: std::io::Error,
    },
    #[error("JSON failed: path={path}, error={source}")]
    Json {
        path: PathBuf,
        source: serde_json::Error,
    },
    #[error("WAV failed: path={path}, error={source}")]
    Wav { path: PathBuf, source: hound::Error },
    #[error("Core ML failed: operation={operation}, {message}")]
    CoreMl { operation: String, message: String },
    #[error("{0}")]
    Invalid(String),
    #[error("Model download failed: url={url}, {message}")]
    ModelDownload { url: String, message: String },
    #[error("Recognition cancelled")]
    Cancelled,
}
pub type Result<T> = std::result::Result<T, Error>;

/// Allowed Core ML execution devices. Core ML chooses placement within this set.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ComputeUnits {
    CpuOnly,
    CpuAndGpu,
    CpuAndNeuralEngine,
    All,
}
impl ComputeUnits {
    pub(crate) fn native(self) -> MLComputeUnits {
        match self {
            Self::CpuOnly => MLComputeUnits::CPUOnly,
            Self::CpuAndGpu => MLComputeUnits::CPUAndGPU,
            Self::CpuAndNeuralEngine => MLComputeUnits::CPUAndNeuralEngine,
            Self::All => MLComputeUnits::All,
        }
    }
}

/// Encoder and decoder execution policy, selected once when loading models.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ComputeConfig {
    pub encoder: ComputeUnits,
    pub decoder: ComputeUnits,
}
impl ComputeConfig {
    pub const fn standard() -> Self {
        Self {
            encoder: ComputeUnits::CpuAndNeuralEngine,
            decoder: ComputeUnits::CpuOnly,
        }
    }
}

fn read_bytes(path: &Path) -> Result<Vec<u8>> {
    fs::read(path).map_err(|source| Error::Io {
        path: path.into(),
        source,
    })
}
fn read_json<T: serde::de::DeserializeOwned>(path: &Path) -> Result<T> {
    serde_json::from_slice(&read_bytes(path)?).map_err(|source| Error::Json {
        path: path.into(),
        source,
    })
}
fn write_json<T: Serialize>(path: &Path, value: &T) -> Result<()> {
    let bytes = serde_json::to_vec_pretty(value).map_err(|source| Error::Json {
        path: path.into(),
        source,
    })?;
    fs::write(path, bytes).map_err(|source| Error::Io {
        path: path.into(),
        source,
    })
}
fn read_f32(path: &Path) -> Result<Vec<f32>> {
    let bytes = read_bytes(path)?;
    if bytes.len() % 4 != 0 {
        return Err(Error::Invalid(format!(
            "Invalid float32 file: path={}, bytes={}",
            path.display(),
            bytes.len()
        )));
    }
    let values: Vec<f32> = bytes
        .chunks_exact(4)
        .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
        .collect();
    if values.iter().any(|x| !x.is_finite()) {
        return Err(Error::Invalid(format!(
            "Non-finite asset: {}",
            path.display()
        )));
    }
    Ok(values)
}

/// Inclusive bounds for one mono 16 kHz model invocation, measured in samples.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct InputLimits {
    minimum_samples: usize,
    maximum_samples: usize,
}

impl InputLimits {
    pub fn new(minimum_samples: usize, maximum_samples: usize) -> Result<Self> {
        if minimum_samples == 0 || minimum_samples > maximum_samples {
            return Err(Error::Invalid(format!(
                "Invalid model input limits: minimum={minimum_samples}, maximum={maximum_samples}"
            )));
        }
        Ok(Self {
            minimum_samples,
            maximum_samples,
        })
    }

    pub fn minimum_samples(self) -> usize {
        self.minimum_samples
    }
    pub fn maximum_samples(self) -> usize {
        self.maximum_samples
    }

    fn validate_samples(self, samples: usize) -> Result<()> {
        if !(self.minimum_samples..=self.maximum_samples).contains(&samples) {
            return Err(Error::Invalid(format!(
                "Audio outside model input limits: samples={samples}, minimum={}, maximum={}",
                self.minimum_samples, self.maximum_samples
            )));
        }
        Ok(())
    }
}

fn validate_array_features(
    descriptions: &NSDictionary<NSString, MLFeatureDescription>,
    features: &[(&str, &[usize], MLMultiArrayDataType)],
    operation: &str,
) -> Result<()> {
    for (name, expected_shape, expected_type) in features {
        let description = descriptions.objectForKey(&NSString::from_str(name));
        // SAFETY: The description is a valid Core ML object; non-array features are rejected.
        let constraint = description
            .and_then(|value| unsafe { value.multiArrayConstraint() })
            .ok_or_else(|| {
                Error::Invalid(format!(
                    "Missing model array: operation={operation}, feature={name}"
                ))
            })?;
        // SAFETY: Shape and dtype belong to this retained immutable constraint.
        let actual_shape: Vec<usize> = unsafe { constraint.shape() }
            .iter()
            .map(|n| n.unsignedIntegerValue())
            .collect();
        let actual_type = unsafe { constraint.dataType() };
        if actual_shape != *expected_shape || actual_type != *expected_type {
            return Err(Error::Invalid(format!(
                "Model array contract mismatch: operation={operation}, feature={name}, expected_shape={expected_shape:?}, actual_shape={actual_shape:?}, expected_dtype={expected_type:?}, actual_dtype={actual_type:?}"
            )));
        }
    }
    Ok(())
}

fn validate_encoder_inputs(model: &MLModel, metadata: &Metadata, path: &Path) -> Result<()> {
    // SAFETY: The model owns its input descriptions.
    let descriptions = unsafe { model.modelDescription().inputDescriptionsByName() };
    let features_shape = [1, metadata.mel_bins, metadata.max_mel_frames];
    validate_array_features(
        &descriptions,
        &[
            ("features", &features_shape, MLMultiArrayDataType::Float32),
            ("length", &[1], MLMultiArrayDataType::Int32),
        ],
        &format!("encoder inputs model={}", path.display()),
    )
}

#[derive(Deserialize)]
struct Metadata {
    model: String,
    revision: String,
    sample_rate: u32,
    fft_size: usize,
    hop_size: usize,
    window_size: usize,
    mel_bins: usize,
    max_mel_frames: usize,
    min_mel_frames: usize,
    preemphasis: f32,
    log_guard: f32,
    normalization_epsilon: f32,
    embedding_rows: usize,
    embedding_width: usize,
    blank_id: usize,
}
#[derive(Deserialize)]
struct Piece {
    text: String,
    kind: u8,
}
#[derive(Deserialize)]
struct Tokenizer {
    pieces: Vec<Piece>,
    unknown_surface: String,
    strip_leading_space: bool,
}
#[derive(Deserialize)]
struct TokenFixture {
    ids: Vec<usize>,
    expected: String,
}
#[derive(Deserialize)]
struct Sample {
    id: String,
    wav: PathBuf,
    expected_features: PathBuf,
    expected_mel_frames: usize,
    expected_tokens: Vec<usize>,
    expected_text: String,
}
struct Assets {
    metadata: Metadata,
    input_limits: InputLimits,
    window: Vec<f32>,
    filterbank: Vec<f32>,
    embeddings: Vec<f32>,
    tokenizer: Tokenizer,
}
impl Assets {
    fn load(directory: &Path) -> Result<Self> {
        let metadata: Metadata = read_json(&directory.join("metadata.json"))?;
        if metadata.fft_size != 512
            || metadata.window_size != 400
            || metadata.hop_size != 160
            || metadata.mel_bins != 128
            || !matches!(metadata.max_mel_frames, 1500 | 3072)
            || metadata.sample_rate != 16000
            || metadata.embedding_rows != 3073
            || metadata.embedding_width != 640
            || metadata.min_mel_frames != 64
            || metadata.blank_id != 3072
        {
            return Err(Error::Invalid(
                "Asset metadata does not match the verified Haas tensor contract".into(),
            ));
        }
        let window = read_f32(&directory.join("window.f32"))?;
        let filterbank = read_f32(&directory.join("filterbank.f32"))?;
        let embeddings = read_f32(&directory.join("embeddings.f32"))?;
        let tokenizer: Tokenizer = read_json(&directory.join("tokenizer.json"))?;
        if window.len() != 400
            || filterbank.len() != 128 * 257
            || embeddings.len() != 3073 * 640
            || tokenizer.pieces.len() != 3072
        {
            return Err(Error::Invalid(format!(
                "Invalid asset dimensions: window={}, filterbank={}, embeddings={}, pieces={}",
                window.len(),
                filterbank.len(),
                embeddings.len(),
                tokenizer.pieces.len()
            )));
        }
        let input_limits = InputLimits::new(
            metadata.min_mel_frames * metadata.hop_size,
            metadata.max_mel_frames * metadata.hop_size,
        )?;
        Ok(Self {
            metadata,
            input_limits,
            window,
            filterbank,
            embeddings,
            tokenizer,
        })
    }
}

fn read_wav(path: &Path, expected_rate: u32, limits: InputLimits) -> Result<Vec<f32>> {
    if expected_rate != recording::SAMPLE_RATE {
        return Err(Error::Invalid(format!(
            "Unsupported WAV target rate: {expected_rate}"
        )));
    }
    let mut source = batch::WavSource::open(path)?;
    if source.total_samples() > limits.maximum_samples() as u64 {
        return Err(Error::Invalid(format!(
            "WAV exceeds model input window: path={}, samples={}, maximum={}; use batch::transcribe_wav",
            path.display(),
            source.total_samples(),
            limits.maximum_samples()
        )));
    }
    let mut audio = Vec::new();
    loop {
        let block = source.read_block(16000)?;
        if block.is_empty() {
            break;
        }
        audio.extend(block);
    }
    if audio.len() as u64 != source.total_samples() {
        return Err(Error::Invalid(
            "WAV ended before its declared sample count".into(),
        ));
    }
    Ok(audio)
}
struct Features {
    values: Vec<f32>,
    valid_frames: usize,
}
/// Reproduce NeMo's constant-centered STFT, original BF16-trained filterbank and
/// unbiased per-feature normalization. No Python or model inference is involved.
fn preprocess(audio: &[f32], assets: &Assets) -> Result<Features> {
    let m = &assets.metadata;
    let frames = audio.len() / m.hop_size;
    if frames < m.min_mel_frames || frames > m.max_mel_frames {
        return Err(Error::Invalid(format!(
            "Audio length outside fixed-window contract: samples={}, mel_frames={frames}, supported={}..{}",
            audio.len(),
            m.min_mel_frames,
            m.max_mel_frames
        )));
    }
    if audio.iter().any(|sample| !sample.is_finite()) {
        return Err(Error::Invalid("Audio contains non-finite samples".into()));
    }
    let emphasized: Vec<f32> = audio
        .iter()
        .enumerate()
        .map(|(i, &v)| {
            if i == 0 {
                v
            } else {
                v - m.preemphasis * audio[i - 1]
            }
        })
        .collect();
    let mut planner = FftPlanner::<f32>::new();
    let fft = planner.plan_fft_forward(m.fft_size);
    let bands: Vec<Vec<(usize, f32)>> = assets
        .filterbank
        .chunks_exact(257)
        .map(|row| {
            row.iter()
                .copied()
                .enumerate()
                .filter(|&(_, v)| v != 0.0)
                .collect()
        })
        .collect();
    let mut logarithms = vec![0.0_f32; m.mel_bins * frames];
    let mut buffer = vec![Complex32::new(0.0, 0.0); m.fft_size];
    for frame in 0..frames {
        buffer.fill(Complex32::new(0.0, 0.0));
        for (j, &window) in assets.window.iter().enumerate() {
            let sample_index = (frame * m.hop_size + j) as isize - (m.window_size / 2) as isize;
            if sample_index >= 0 && (sample_index as usize) < audio.len() {
                buffer[j + (m.fft_size - m.window_size) / 2].re =
                    emphasized[sample_index as usize] * window;
            }
        }
        fft.process(&mut buffer);
        let power: Vec<f32> = buffer[..257]
            .iter()
            .map(|v| {
                let magnitude = (v.re * v.re + v.im * v.im).sqrt();
                magnitude * magnitude
            })
            .collect();
        for (mel, band) in bands.iter().enumerate() {
            let energy = band
                .iter()
                .map(|&(bin, weight)| weight as f64 * power[bin] as f64)
                .sum::<f64>() as f32;
            logarithms[mel * frames + frame] = (energy + m.log_guard).ln();
        }
    }
    let mut values = vec![0.0_f32; m.mel_bins * m.max_mel_frames];
    for mel in 0..m.mel_bins {
        let row = &logarithms[mel * frames..(mel + 1) * frames];
        let mean = (row.iter().map(|&v| v as f64).sum::<f64>() / frames as f64) as f32;
        let variance = (row
            .iter()
            .map(|&v| {
                let d = v - mean;
                (d * d) as f64
            })
            .sum::<f64>()
            / (frames - 1) as f64) as f32;
        let std = variance.sqrt() + m.normalization_epsilon;
        for (t, &v) in row.iter().enumerate() {
            values[mel * m.max_mel_frames + t] = (v - mean) / std;
        }
    }
    if values.iter().any(|x| !x.is_finite()) {
        return Err(Error::Invalid(
            "Native preprocessing produced non-finite values".into(),
        ));
    }
    Ok(Features {
        values,
        valid_frames: frames,
    })
}

/// SentencePiece replaces each malformed byte separately, including an
/// incomplete trailing sequence. Rust's from_utf8_lossy groups some such bytes.
fn decode_bytes(bytes: &[u8]) -> String {
    let mut output = String::new();
    let mut index = 0;
    while index < bytes.len() {
        let length = match bytes[index] {
            0x00..=0x7f => 1,
            0xc2..=0xdf => 2,
            0xe0..=0xef => 3,
            0xf0..=0xf4 => 4,
            _ => 0,
        };
        if length > 0
            && index + length <= bytes.len()
            && let Ok(text) = std::str::from_utf8(&bytes[index..index + length])
        {
            output.push_str(text);
            index += length;
        } else {
            output.push('\u{fffd}');
            index += 1;
        }
    }
    output
}
impl Tokenizer {
    fn decode(&self, ids: &[usize]) -> Result<String> {
        let mut output = String::new();
        let mut bytes = Vec::new();
        let mut strip_prefix = self.strip_leading_space;
        for &id in ids {
            let piece = self.pieces.get(id).ok_or_else(|| {
                Error::Invalid(format!(
                    "Token ID outside tokenizer: id={id}, vocabulary={}",
                    self.pieces.len()
                ))
            })?;
            if piece.kind == 6 {
                let text = piece
                    .text
                    .strip_prefix("<0x")
                    .and_then(|p| p.strip_suffix('>'))
                    .ok_or_else(|| {
                        Error::Invalid(format!(
                            "Malformed byte piece: id={id}, piece={}",
                            piece.text
                        ))
                    })?;
                let byte = u8::from_str_radix(text, 16).map_err(|e| {
                    Error::Invalid(format!(
                        "Invalid byte piece: id={id}, text={text}, error={e}"
                    ))
                })?;
                bytes.push(byte);
                continue;
            }
            if !bytes.is_empty() {
                output.push_str(&decode_bytes(&bytes));
                bytes.clear();
                strip_prefix = false;
            }
            match piece.kind {
                2 => {
                    output.push_str(&self.unknown_surface);
                    strip_prefix = false;
                }
                3 => {}
                1 | 4 | 5 => {
                    let text = if strip_prefix {
                        piece.text.strip_prefix('▁').unwrap_or(&piece.text)
                    } else {
                        &piece.text
                    };
                    output.push_str(&text.replace('▁', " "));
                    strip_prefix = false;
                }
                other => {
                    return Err(Error::Invalid(format!(
                        "Unsupported SentencePiece kind: id={id}, kind={other}"
                    )));
                }
            }
        }
        if !bytes.is_empty() {
            output.push_str(&decode_bytes(&bytes));
        }
        Ok(output)
    }
}

fn core_message(error: &NSError) -> String {
    format!(
        "domain={}, code={}, description={}",
        error.domain(),
        error.code(),
        error.localizedDescription()
    )
}
fn core_error(operation: &str, error: Retained<NSError>) -> Error {
    Error::CoreMl {
        operation: operation.into(),
        message: core_message(&error),
    }
}

fn load_model(path: &Path, units: MLComputeUnits) -> Result<Retained<MLModel>> {
    if path.extension().and_then(|s| s.to_str()) != Some("mlmodelc") || !path.is_dir() {
        return Err(Error::Invalid(format!(
            "Expected precompiled .mlmodelc directory: {}",
            path.display()
        )));
    }
    let name = path
        .to_str()
        .ok_or_else(|| Error::Invalid(format!("Non-UTF8 model path: {}", path.display())))?;
    let url = NSURL::fileURLWithPath(&NSString::from_str(name));
    // SAFETY: Standard Core ML configuration, URL and model-loading APIs.
    unsafe {
        let config = MLModelConfiguration::new();
        config.setComputeUnits(units);
        MLModel::modelWithContentsOfURL_configuration_error(&url, &config)
            .map_err(|e| core_error(&format!("load {}", path.display()), e))
    }
}
fn feature_provider(
    features: &[(&str, &MLMultiArray)],
) -> Result<Retained<MLDictionaryFeatureProvider>> {
    let names: Vec<_> = features
        .iter()
        .map(|(name, _)| NSString::from_str(name))
        .collect();
    // SAFETY: Each value wraps a valid retained multi-array, not an arbitrary object.
    let values: Vec<_> = features
        .iter()
        .map(|(_, array)| unsafe { MLFeatureValue::featureValueWithMultiArray(array) })
        .collect();
    let keys: Vec<&NSString> = names.iter().map(|n| &**n).collect();
    // The SDK's dictionary initializer requires AnyObject at this API boundary.
    let objects: Vec<&AnyObject> = values.iter().map(|v| v.as_ref()).collect();
    let dictionary = NSDictionary::from_slices(&keys, &objects);
    // SAFETY: Every dictionary value is a retained MLFeatureValue.
    unsafe {
        MLDictionaryFeatureProvider::initWithDictionary_error(
            MLDictionaryFeatureProvider::alloc(),
            &dictionary,
        )
        .map_err(|e| core_error("construct feature provider", e))
    }
}

fn predict(
    model: &MLModel,
    features: &[(&str, &MLMultiArray)],
) -> Result<Retained<ProtocolObject<dyn MLFeatureProvider>>> {
    let provider = feature_provider(features)?;
    // SAFETY: The provider owns valid feature arrays for this synchronous prediction.
    unsafe {
        model
            .predictionFromFeatures_error(ProtocolObject::from_ref(&*provider))
            .map_err(|e| {
                core_error(
                    &format!(
                        "predict inputs={:?}",
                        features.iter().map(|(n, _)| *n).collect::<Vec<_>>()
                    ),
                    e,
                )
            })
    }
}

fn output_array(
    provider: &ProtocolObject<dyn MLFeatureProvider>,
    name: &str,
) -> Result<Retained<MLMultiArray>> {
    // SAFETY: This is a valid Core ML feature provider; output type is checked explicitly.
    unsafe { provider.featureValueForName(&NSString::from_str(name)) }
        .and_then(|v| unsafe { v.multiArrayValue() })
        .ok_or_else(|| Error::Invalid(format!("Missing multi-array model output: {name}")))
}
struct Encoded {
    values: Vec<f32>,
    storage_frames: usize,
    valid_frames: usize,
}
fn encode(model: &MLModel, features: &Features, metadata: &Metadata) -> Result<Encoded> {
    let input = multi_array(
        &[1, metadata.mel_bins, metadata.max_mel_frames],
        &features.values,
    )?;
    let length = multi_array(
        &[1],
        &[i32::try_from(features.valid_frames).map_err(|e| Error::Invalid(e.to_string()))?],
    )?;
    let result = predict(model, &[("features", &input), ("length", &length)])?;
    let encoded = output_array(&result, "encoded")?;
    let shape = array_shape(&encoded);
    if shape.len() != 3 || shape[0] != 1 || shape[1] != 1024 {
        return Err(Error::Invalid(format!(
            "Unexpected encoder output shape: {shape:?}"
        )));
    }
    let values = read_floating_tensor(&encoded)?;
    let length_output = output_array(&result, "encoded_length")?;
    let actual_length = read_array::<i32>(&length_output)?;
    if actual_length.len() != 1
        || actual_length[0] <= 0
        || actual_length[0] as usize > shape[2]
        || values.iter().any(|x| !x.is_finite())
    {
        return Err(Error::Invalid(format!(
            "Invalid encoder output: shape={shape:?}, length={actual_length:?}"
        )));
    }
    Ok(Encoded {
        values,
        storage_frames: shape[2],
        valid_frames: actual_length[0] as usize,
    })
}

/// Cooperative cancellation, checked between stages and on every TDT decoding step.
/// An already-running Core ML prediction completes before cancellation is observed.
#[derive(Clone, Default)]
pub struct Cancellation(std::sync::Arc<CancellationState>);
#[derive(Default)]
struct CancellationState {
    cancelled: std::sync::atomic::AtomicBool,
    parents: Vec<Cancellation>,
}
impl Cancellation {
    pub fn new() -> Self {
        Self::default()
    }
    pub fn cancel(&self) {
        self.0
            .cancelled
            .store(true, std::sync::atomic::Ordering::Release);
    }
    fn linked(first: &Self, second: &Self) -> Self {
        Self(std::sync::Arc::new(CancellationState {
            cancelled: std::sync::atomic::AtomicBool::new(false),
            parents: vec![first.clone(), second.clone()],
        }))
    }
    fn check(&self) -> Result<()> {
        if self.0.cancelled.load(std::sync::atomic::Ordering::Acquire) {
            return Err(Error::Cancelled);
        }
        for parent in &self.0.parents {
            parent.check()?;
        }
        Ok(())
    }
}

/// One independent utterance; inference states are never carried into the next call.
#[derive(Debug, Serialize)]
pub struct Recognition {
    pub text: String,
    pub tokens: Vec<usize>,
    pub token_timings: Vec<TimedToken>,
    pub audio_seconds: f64,
    pub preprocess_seconds: f64,
    pub encode_seconds: f64,
    pub decode_seconds: f64,
}

/// Owns local model assets and retained Core ML models. Dropping releases these handles.
/// Accepts mono 16 kHz input within the loaded model's input_limits().
pub struct HaasEngine {
    assets: Assets,
    encoder: Retained<MLModel>,
    decoder: DecoderModels,
    compute: ComputeConfig,
    vocabulary_id: uuid::Uuid,
}
impl HaasEngine {
    /// Load local compiled models without downloads or child processes.
    pub fn load(directory: &Path, encoder_path: &Path, predictor_path: &Path) -> Result<Self> {
        Self::load_with_compute(
            directory,
            encoder_path,
            predictor_path,
            ComputeConfig::standard(),
        )
    }

    pub fn load_with_compute(
        directory: &Path,
        encoder_path: &Path,
        predictor_path: &Path,
        compute: ComputeConfig,
    ) -> Result<Self> {
        Self::load_models(directory, encoder_path, compute, || {
            Ok(DecoderModels::Fused(load_model(
                predictor_path,
                compute.decoder.native(),
            )?))
        })
    }

    /// Load separate FP32 prediction and joint models. The joint may return logits or scalar decisions.
    pub fn load_split(
        directory: &Path,
        encoder_path: &Path,
        prediction_path: &Path,
        joint_path: &Path,
    ) -> Result<Self> {
        Self::load_split_with_compute(
            directory,
            encoder_path,
            prediction_path,
            joint_path,
            ComputeConfig::standard(),
        )
    }

    pub fn load_split_with_compute(
        directory: &Path,
        encoder_path: &Path,
        prediction_path: &Path,
        joint_path: &Path,
        compute: ComputeConfig,
    ) -> Result<Self> {
        Self::load_models(directory, encoder_path, compute, || {
            DecoderModels::split(
                load_model(prediction_path, compute.decoder.native())?,
                load_model(joint_path, compute.decoder.native())?,
                joint_path,
            )
        })
    }

    fn load_models(
        directory: &Path,
        encoder_path: &Path,
        compute: ComputeConfig,
        load_decoder: impl FnOnce() -> Result<DecoderModels>,
    ) -> Result<Self> {
        autoreleasepool(|_| {
            let assets = Assets::load(directory)?;
            let encoder = load_model(encoder_path, compute.encoder.native())?;
            validate_encoder_inputs(&encoder, &assets.metadata, encoder_path)?;
            Ok(Self {
                assets,
                encoder,
                decoder: load_decoder()?,
                compute,
                vocabulary_id: uuid::Uuid::new_v4(),
            })
        })
    }
    /// Sample bounds verified against the encoder's declared input shape.
    pub fn input_limits(&self) -> InputLimits {
        self.assets.input_limits
    }
    pub fn compute_config(&self) -> ComputeConfig {
        self.compute
    }
    pub fn compile_vocabulary(
        &self,
        vocabulary: Vocabulary,
        tokenizer_model: &Path,
    ) -> Result<CompiledVocabulary> {
        self.decoder.require_logits()?;
        CompiledVocabulary::compile(
            vocabulary,
            tokenizer_model,
            &self.assets.tokenizer,
            self.vocabulary_id,
        )
    }
    pub fn transcribe_samples_with_vocabulary(
        &self,
        samples: &[f32],
        sample_rate: u32,
        vocabulary: &CompiledVocabulary,
        cancellation: &Cancellation,
    ) -> Result<Recognition> {
        if sample_rate != 16000 {
            return Err(Error::Invalid(format!(
                "Expected mono16000Hz vocabulary input, got sample_rate={sample_rate}"
            )));
        }
        vocabulary.validate_engine(self.vocabulary_id)?;
        let (_, mut result) = self.recognize_using(samples, cancellation, |encoded| {
            let plain = decode(&self.decoder, encoded, &self.assets, cancellation)?;
            if !vocabulary.has_speech(samples) {
                return Ok(plain);
            }
            let boosted = decoder::decode_biased(
                &self.decoder,
                encoded,
                &self.assets,
                vocabulary,
                cancellation,
            )?;
            Ok(vocabulary.guard(&plain, &boosted))
        })?;
        result.text = vocabulary.format(&result.text);
        Ok(result)
    }

    pub fn transcribe_samples(
        &self,
        samples: &[f32],
        sample_rate: u32,
        cancellation: &Cancellation,
    ) -> Result<Recognition> {
        if sample_rate != 16000 {
            return Err(Error::Invalid(format!(
                "Expected mono 16000 Hz input, got sample_rate={sample_rate}"
            )));
        }
        self.recognize(samples, cancellation)
            .map(|(_, recognition)| recognition)
    }
    pub fn transcribe_wav(&self, path: &Path, cancellation: &Cancellation) -> Result<Recognition> {
        let samples = read_wav(path, self.assets.metadata.sample_rate, self.input_limits())?;
        self.transcribe_samples(&samples, self.assets.metadata.sample_rate, cancellation)
    }
    /// Return shared metadata and raw vocabulary pieces with measured ASR time.
    pub fn transcribe_result(
        &self,
        samples: &[f32],
        sample_rate: u32,
        cancellation: &Cancellation,
    ) -> Result<result::AsrResult> {
        let started = Instant::now();
        let recognition = self.transcribe_samples(samples, sample_rate, cancellation)?;
        let metrics = result::InferenceMetrics::measured(
            samples.len(),
            started.elapsed().as_secs_f64(),
            Some(result::StageTiming::from_recognition(&recognition)),
        );
        let transcript = window::Transcript {
            text: recognition.text,
            tokens: recognition.token_timings,
            audio_samples: samples.len() as u64,
        };
        let result = result::AsrResult::build(self, &transcript, metrics)?;
        cancellation.check()?;
        Ok(result)
    }
    pub fn transcribe_wav_result(
        &self,
        path: &Path,
        cancellation: &Cancellation,
    ) -> Result<result::AsrResult> {
        cancellation.check()?;
        let samples = read_wav(path, self.assets.metadata.sample_rate, self.input_limits())?;
        self.transcribe_result(&samples, self.assets.metadata.sample_rate, cancellation)
    }
    pub fn token_piece(&self, id: usize) -> Result<result::TokenPiece> {
        let piece = self.assets.tokenizer.pieces.get(id).ok_or_else(|| {
            Error::Invalid(format!(
                "Token ID outside tokenizer: id={id}, vocabulary={}",
                self.assets.tokenizer.pieces.len()
            ))
        })?;
        let kind = match piece.kind {
            1 => result::TokenKind::Normal,
            2 => result::TokenKind::Unknown,
            3 => result::TokenKind::Control,
            4 => result::TokenKind::UserDefined,
            5 => result::TokenKind::Unused,
            6 => result::TokenKind::Byte,
            kind => {
                return Err(Error::Invalid(format!(
                    "Unsupported token kind: id={id}, kind={kind}"
                )));
            }
        };
        Ok(result::TokenPiece {
            text: piece.text.clone(),
            kind,
        })
    }
    fn recognize(
        &self,
        samples: &[f32],
        cancellation: &Cancellation,
    ) -> Result<(Features, Recognition)> {
        self.recognize_using(samples, cancellation, |encoded| {
            decode(&self.decoder, encoded, &self.assets, cancellation)
        })
    }
    fn recognize_using<F>(
        &self,
        samples: &[f32],
        cancellation: &Cancellation,
        decode_tokens: F,
    ) -> Result<(Features, Recognition)>
    where
        F: FnOnce(&Encoded) -> Result<Vec<TimedToken>>,
    {
        autoreleasepool(|_| {
            cancellation.check()?;
            self.input_limits().validate_samples(samples.len())?;
            let start = Instant::now();
            let features = preprocess(samples, &self.assets)?;
            let preprocess_seconds = start.elapsed().as_secs_f64();
            cancellation.check()?;
            let start = Instant::now();
            let encoded = encode(&self.encoder, &features, &self.assets.metadata)?;
            let encode_seconds = start.elapsed().as_secs_f64();
            cancellation.check()?;
            let start = Instant::now();
            let token_timings = decode_tokens(&encoded)?;
            let decode_seconds = start.elapsed().as_secs_f64();
            let tokens: Vec<usize> = token_timings.iter().map(|t| t.token_id).collect();
            let text = self.assets.tokenizer.decode(&tokens)?;
            cancellation.check()?;
            Ok((
                features,
                Recognition {
                    text,
                    tokens,
                    token_timings,
                    audio_seconds: samples.len() as f64 / 16000.0,
                    preprocess_seconds,
                    encode_seconds,
                    decode_seconds,
                },
            ))
        })
    }
}

#[derive(Serialize)]
struct FeatureCheck {
    id: String,
    valid_frames: usize,
    expected_valid_frames: usize,
    max_absolute_error: f32,
    mean_absolute_error: f64,
}
fn check_features(sample: &Sample, features: &Features) -> Result<FeatureCheck> {
    let expected = read_f32(&sample.expected_features)?;
    if features.values.len() != expected.len() {
        return Err(Error::Invalid(format!(
            "Feature shape mismatch: clip={}, actual={}, expected={}",
            sample.id,
            features.values.len(),
            expected.len()
        )));
    }
    let differences: Vec<f32> = features
        .values
        .iter()
        .zip(&expected)
        .map(|(a, b)| (a - b).abs())
        .collect();
    Ok(FeatureCheck {
        id: sample.id.clone(),
        valid_frames: features.valid_frames,
        expected_valid_frames: sample.expected_mel_frames,
        max_absolute_error: differences.iter().copied().fold(0.0, f32::max),
        mean_absolute_error: differences.iter().map(|&v| v as f64).sum::<f64>()
            / differences.len() as f64,
    })
}
#[derive(Serialize)]
struct TokenCheck {
    index: usize,
    expected: String,
    actual: String,
    matches: bool,
}
#[derive(Serialize)]
struct Verification {
    tokenizer: Vec<TokenCheck>,
    features: Vec<FeatureCheck>,
    feature_tolerance: f32,
    passed: bool,
}
pub fn verify_assets(directory: &Path, output: &Path) -> Result<()> {
    let assets = Assets::load(directory)?;
    let fixtures: Vec<TokenFixture> = read_json(&directory.join("tokenizer-fixtures.json"))?;
    let tokenizer = fixtures
        .iter()
        .enumerate()
        .map(|(index, f)| {
            let actual = assets.tokenizer.decode(&f.ids)?;
            Ok(TokenCheck {
                index,
                expected: f.expected.clone(),
                matches: actual == f.expected,
                actual,
            })
        })
        .collect::<Result<Vec<_>>>()?;
    let samples: Vec<Sample> = read_json(&directory.join("samples.json"))?;
    let mut features = Vec::new();
    for sample in &samples {
        let audio = read_wav(
            &sample.wav,
            assets.metadata.sample_rate,
            assets.input_limits,
        )?;
        features.push(check_features(sample, &preprocess(&audio, &assets)?)?);
    }
    let tolerance = 0.001;
    let passed = tokenizer.iter().all(|c| c.matches)
        && features.iter().all(|c| {
            c.valid_frames == c.expected_valid_frames && c.max_absolute_error <= tolerance
        });
    write_json(
        output,
        &Verification {
            tokenizer,
            features,
            feature_tolerance: tolerance,
            passed,
        },
    )?;
    if !passed {
        return Err(Error::Invalid(format!(
            "Native asset verification failed; see {}",
            output.display()
        )));
    }
    eprintln!("Native tokenizer and preprocessing verification passed");
    Ok(())
}

#[derive(Serialize)]
struct Transcript {
    id: String,
    text: String,
    expected_text: String,
    tokens: Vec<usize>,
    tokens_match: bool,
    text_matches: bool,
    feature_check: FeatureCheck,
    preprocess_seconds: f64,
    encode_seconds: f64,
    decode_seconds: f64,
}
#[derive(Serialize)]
struct Report {
    runtime: String,
    model: String,
    revision: String,
    model_load_seconds: f64,
    samples: Vec<Transcript>,
    cancellation_checked: bool,
    passed: bool,
}
pub fn verify_transcriptions(
    directory: &Path,
    encoder_path: &Path,
    predictor_path: &Path,
    output: &Path,
) -> Result<()> {
    let assets = Assets::load(directory)?;
    let samples: Vec<Sample> = read_json(&directory.join("samples.json"))?;
    let start = Instant::now();
    let engine = HaasEngine::load(directory, encoder_path, predictor_path)?;
    let model_load_seconds = start.elapsed().as_secs_f64();
    let mut transcripts = Vec::new();
    let cancelled = Cancellation::new();
    cancelled.cancel();
    let cancellation_checked = matches!(
        engine.transcribe_samples(&vec![0.0; 16000], 16000, &cancelled),
        Err(Error::Cancelled)
    );
    // Forward then reverse order exposes stale decoder state and cancellation leakage.
    for sample in samples.iter().chain(samples.iter().rev()) {
        let row = autoreleasepool(|_| -> Result<Transcript> {
            let audio = read_wav(
                &sample.wav,
                assets.metadata.sample_rate,
                assets.input_limits,
            )?;
            let (features, recognition) = engine.recognize(&audio, &Cancellation::new())?;
            let feature_check = check_features(sample, &features)?;
            let Recognition {
                tokens,
                text,
                preprocess_seconds,
                encode_seconds,
                decode_seconds,
                ..
            } = recognition;
            Ok(Transcript {
                id: sample.id.clone(),
                expected_text: sample.expected_text.clone(),
                tokens_match: tokens == sample.expected_tokens,
                text_matches: text == sample.expected_text,
                text,
                tokens,
                feature_check,
                preprocess_seconds,
                encode_seconds,
                decode_seconds,
            })
        })?;
        eprintln!(
            "{}: tokens_match={}, text_matches={}",
            row.id, row.tokens_match, row.text_matches
        );
        transcripts.push(row);
    }
    let passed = cancellation_checked
        && transcripts.iter().all(|r| {
            r.tokens_match
                && r.text_matches
                && r.feature_check.max_absolute_error <= 0.001
                && r.feature_check.valid_frames == r.feature_check.expected_valid_frames
        });
    write_json(
        output,
        &Report {
            runtime: "Rust + macOS Core ML; no Python or Swift process".into(),
            model: assets.metadata.model.clone(),
            revision: assets.metadata.revision.clone(),
            model_load_seconds,
            samples: transcripts,
            cancellation_checked,
            passed,
        },
    )?;
    if !passed {
        return Err(Error::Invalid(format!(
            "Native transcription mismatch; see {}",
            output.display()
        )));
    }
    Ok(())
}
fn copy_compiled_directory(source: &Path, destination: &Path) -> Result<()> {
    fs::create_dir(destination).map_err(|source| Error::Io {
        path: destination.into(),
        source,
    })?;
    let entries = fs::read_dir(source).map_err(|error| Error::Io {
        path: source.into(),
        source: error,
    })?;
    for entry in entries {
        let entry = entry.map_err(|error| Error::Io {
            path: source.into(),
            source: error,
        })?;
        let origin = entry.path();
        let target = destination.join(entry.file_name());
        let kind = entry.file_type().map_err(|source| Error::Io {
            path: origin.clone(),
            source,
        })?;
        if kind.is_dir() {
            copy_compiled_directory(&origin, &target)?;
        } else if kind.is_file() {
            fs::copy(&origin, &target).map_err(|source| Error::Io {
                path: target,
                source,
            })?;
        } else {
            return Err(Error::Invalid(format!(
                "Unexpected file type in compiled model: {}",
                origin.display()
            )));
        }
    }
    Ok(())
}

/// Explicit build-time compilation; transcription accepts precompiled models only.
pub fn compile_model(source: &Path, destination: &Path) -> Result<()> {
    if source.extension().and_then(|s| s.to_str()) != Some("mlpackage")
        || !source.is_dir()
        || destination.extension().and_then(|s| s.to_str()) != Some("mlmodelc")
        || destination.exists()
    {
        return Err(Error::Invalid(format!(
            "Expected existing .mlpackage source and new .mlmodelc destination: source={}, destination={}",
            source.display(),
            destination.display()
        )));
    }
    let path = source
        .to_str()
        .ok_or_else(|| Error::Invalid(format!("Non-UTF8 model path: {}", source.display())))?;
    let url = NSURL::fileURLWithPath(&NSString::from_str(path));
    let (sender, receiver) = std::sync::mpsc::channel::<Result<PathBuf>>();
    let operation = format!("compile {}", source.display());
    let handler = block2::RcBlock::new(move |compiled: *mut NSURL, error: *mut NSError| {
        // SAFETY: Apple keeps callback arguments valid for the duration of the callback.
        let result = unsafe {
            if let Some(error) = error.as_ref() {
                Err(Error::CoreMl {
                    operation: operation.clone(),
                    message: core_message(error),
                })
            } else if let Some(compiled) = compiled.as_ref() {
                compiled
                    .path()
                    .map(|path| PathBuf::from(path.to_string()))
                    .ok_or_else(|| {
                        Error::Invalid("Core ML compiler returned a URL without a file path".into())
                    })
            } else {
                Err(Error::Invalid(
                    "Core ML compilation returned neither a model URL nor an error".into(),
                ))
            }
        };
        if let Err(error) = sender.send(result) {
            eprintln!("Core ML compilation callback receiver disconnected: {error}");
        }
    });
    // SAFETY: The callback is retained through completion; no Objective-C object crosses Rust threads.
    unsafe { MLModel::compileModelAtURL_completionHandler(&url, &handler) };
    let temporary = receiver.recv().map_err(|error| {
        Error::Invalid(format!("Core ML compilation channel failed: {error}"))
    })??;
    copy_compiled_directory(&temporary, destination)?;
    eprintln!("Compiled model: {}", destination.display());
    Ok(())
}

#[cfg(test)]
mod tests {
    #[test]
    fn label_posteriors_are_stable_and_exclude_duration_classes() {
        let probability = super::label_probability(&[10000.0, 10000.0], 0).unwrap();
        assert!((probability - 0.5).abs() < 0.000001);
        assert!(
            (super::label_probability(&[0.0, 1.0], 1).unwrap()
                - super::label_probability(&[100.0, 101.0], 1).unwrap())
            .abs()
                < 0.000001
        );
        assert!(super::label_probability(&[], 0).is_err());
        assert!(super::label_probability(&[f32::NAN], 0).is_err());
        assert!(super::label_probability(&[0.0], 1).is_err());
    }

    use super::*;
    #[test]
    fn coreml_half_precision_output_is_read_as_declared() {
        autoreleasepool(|_| {
            let values = [
                half::f16::from_f32(0.0),
                half::f16::from_f32(1.0),
                half::f16::from_f32(-2.0),
            ];
            let array = multi_array(&[1, 3], &values).expect("valid float16 array");
            assert_eq!(
                read_floating_tensor(&array).expect("read declared float16"),
                [0.0, 1.0, -2.0]
            );
        });
    }
    #[test]
    fn tokenizer_preserves_real_space_after_dummy_prefix() {
        let tokenizer = Tokenizer {
            pieces: vec![
                Piece {
                    text: "▁".into(),
                    kind: 1,
                },
                Piece {
                    text: "空".into(),
                    kind: 1,
                },
            ],
            unknown_surface: " ⁇ ".into(),
            strip_leading_space: true,
        };
        assert_eq!(
            tokenizer.decode(&[0, 0, 1]).expect("valid token fixture"),
            " 空"
        );
    }
    #[test]
    fn malformed_sentencepiece_bytes_are_replaced_individually() {
        assert_eq!(decode_bytes(&[0xe3, 0x81]), "��");
        assert_eq!(decode_bytes("𠮷😀".as_bytes()), "𠮷😀");
    }
    #[test]
    fn tensor_strides_preserve_logical_row_major_order() {
        let offsets: Vec<_> = (0..6)
            .map(|i| flat_offset(i, &[1, 3, 2], &[6, 1, 3]))
            .collect();
        assert_eq!(offsets, [0, 3, 1, 4, 2, 5]);
    }
}