coremlit 0.1.1

//! CoreML wrapper for `pyannote_segmentation.mlmodelc` (spec §4) and the
//! powerset→multilabel decode that turns its per-frame log-probabilities
//! into dia's `segmentations` tensor layout.
//!
//! Ports the model-facing half of dia's `segment` stage — `SegmentModel`
//! (`diarization/src/segment/model.rs`) — over `coremlit` instead of `ort`,
//! plus the powerset argmax dia's owned audio-in pipeline performs inline
//! (`diarization/src/offline/owned.rs:476-499`; dia's own
//! `powerset_to_speakers_hard`, `diarization/src/segment/powerset.rs:
//! 68-87`). Ground truth: `tests/model_io.rs`'s
//! `pyannote_segmentation_io_matches_spec` introspection test
//! (`pyannote_segmentation.mlmodelc`: `audio [1, 1, 160_000]` f32 in,
//! `segments [1, 589, 7]` f32 out) and the design spec §4/§5.
//!
//! # dia contract match
//!
//! - **Padding**: dia's `SegmentModel::infer` (the ort analog,
//!   `diarization/src/segment/model.rs:280-357`) requires its caller to
//!   already have exactly `WINDOW_SAMPLES` samples —
//!   `debug_assert_eq!(samples.len(), WINDOW_SAMPLES as usize)`
//!   (`model.rs:281`) — and does not pad internally; zero-padding happens
//!   one layer up, in `Segmenter::emit_window`
//!   (`diarization/src/segment/segmenter.rs:250-266`) or, for the owned
//!   audio-in pipeline, in `OwnedDiarizationPipeline`'s own chunk loop
//!   (`diarization/src/offline/owned.rs:453-475`, `padded_chunk`).
//!   [`SegmentModel::infer`] mirrors that same boundary: it REJECTS
//!   (`InferError::InputLength`) rather than pads, so a caller (a future
//!   `Extractor`, mirroring `owned.rs`'s loop) is responsible for
//!   zero-padding short chunks before calling — exactly dia's contract,
//!   made a checked `Result` instead of a debug-only assertion.
//! - **Non-finite inputs and outputs**: dia's `infer` rejects non-finite
//!   *input* samples before running inference (`Error::NonFiniteInput`,
//!   `model.rs:283-290`) AND non-finite output logits after extracting them
//!   (`Error::NonFiniteOutput`, `model.rs:346-355`). [`SegmentModel::infer`]
//!   now mirrors BOTH: `check_finite_input` scans `samples` before the CoreML
//!   call ([`InferError::NonFiniteInput`]) and `check_finite` scans the
//!   extracted logits after ([`InferError::NonFiniteOutput`]). Relying on the
//!   output scan alone was a known gap: a NaN input is exactly the `ort`
//!   CoreML-EP corruption mode this crate exists to replace, and although it
//!   typically propagates to a non-finite logit, a kernel could in principle
//!   absorb it into finite garbage the output scan would miss — so the input
//!   scan closes that boundary directly, matching dia's own input-side guard
//!   and the identical scan the embed module already performs on its inputs.
//! - **Layout re-validation**: dia's `infer` re-validates its output's
//!   shape/layout on every call (`model.rs:313-338`) because, per its own
//!   doc comment, "Load-time dimension verification ... is reserved for a
//!   future revision once a stable ort metadata API is available"
//!   (`model.rs:134-138`). `crate::Model::description()` already IS
//!   that stable metadata API, so [`SegmentModel::from_file_with`]
//!   validates the declared contract (presence, dtype, base shape) once
//!   at construction (`ModelError::ContractMismatch`) instead of
//!   re-deriving it from scratch on every call — a deliberate, documented
//!   improvement dia's own comment anticipates. A stable *declared*
//!   contract at load time is not, however, a guarantee that every
//!   individual CoreML prediction actually returns output matching it —
//!   the runtime is a trust boundary independent of whether its metadata
//!   API is introspectable — so [`SegmentModel::infer`] also re-validates
//!   the predict-time `segments` tensor's shape on every call
//!   (`InferError::OutputShape`), matching dia's own per-call re-check.
//! - **dtype**: [`multilabel`] returns `Vec<f64>` because dia's
//!   `OfflineInput::segmentations` field is `&'a [f64]`
//!   (`diarization/src/offline/algo.rs:179`); dia itself produces those
//!   values by computing the hard 0/1 mask in `f32` and casting `as f64`
//!   at the point it writes into its own segmentations buffer
//!   (`diarization/src/offline/owned.rs:496`). `0.0_f32 as f64` and
//!   `1.0_f32 as f64` are exact (both values are exactly representable in
//!   both formats), so building the lookup table directly in `f64` here is
//!   bit-identical to dia's f32-then-cast.
//!
//! # `segments` is `log_softmax(·)`, NOT raw logits
//!
//! This module's docs asserted "raw powerset logits" until someone read
//! the graph. They are not. `pyannote_segmentation.mlmodelc/model.mil`
//! ends (immediately before its `-> (segments)` return):
//!
//! ```text
//! var_247_cast_fp16 = reduce_log_sum_exp(axes = [-1], keep_dims = true,
//!                                        x = linear_2_cast_fp16)
//! var_249_cast_fp16 = sub(x = linear_2_cast_fp16, y = var_247_cast_fp16)
//! segments          = cast(dtype = fp32, x = var_249_cast_fp16)
//! ```
//!
//! `segments` holds per-frame **log-probabilities**: `z − logsumexp(z)`, the
//! fused form of `log_softmax`. The fp32 `Segmentation.mlmodelc` variant
//! computes the same quantity and names its output `log_probs`, which is the
//! honest name.
//!
//! The "raw logits" mistake was not cosmetic: believing there was no `log` in
//! the graph is exactly why nobody looked at the `log`'s guard epsilon for a
//! full review cycle. See "fp16 safety" below.
//!
//! # Why argmaxing it without a softmax is still correct
//!
//! dia's audio-in pipeline runs `softmax_row(&row)` before argmaxing
//! (`diarization/src/offline/owned.rs:484-494`). Its ONNX segmentation model
//! emits log-probabilities — the same quantity this CoreML model does, NOT the
//! raw logits `z` this section once claimed: `tests/generate_goldens.rs` checks
//! the dia-ort output directly (every element `<= 0`, each 7-class row's
//! `sum(exp(row)) == 1.0`). So dia's `softmax_row` renormalizes those log-probs
//! back to probabilities (`softmax(log(softmax(z))) == softmax(z)`) rather than
//! converting raw logits; [`multilabel`] skips that step and argmaxes the
//! log-probs directly.
//!
//! In exact arithmetic `log(softmax(z))_i = z_i - logsumexp(z)`, and
//! `logsumexp(z)` is a single scalar shared by every class in the row. So
//! this model's output is each frame's logit vector shifted by a **per-row
//! constant** — and subtracting a constant from every element of a row
//! changes neither which element is largest nor which elements are exactly
//! equal:
//!
//! - `z_i > z_j  <=>  z_i - c > z_j - c`
//! - `z_i == z_j <=> z_i - c == z_j - c`
//!
//! Hence `argmax(log(softmax(z))) == argmax(softmax(z)) == argmax(z)`, ties
//! included — **in exact arithmetic**. That is the level at which
//! [`multilabel`]'s direct `hard_argmax` and dia's softmax-THEN-argmax decode
//! agree; they are NOT order-for-order identical over f32. dia's extra `softmax`
//! runs `exp`, which can round two log-probs one ULP apart to the SAME
//! probability, manufacturing a tie the lowest-index rule then breaks toward the
//! LOWER index — so the two decodes can pick different classes on a near-tie
//! (`tests/parity_seg.rs`'s `near_tie_softmax_can_flip_the_argmax`, the standing
//! witness for why the universal wording was wrong). They agree on every
//! committed golden row today (`tests/parity_seg.rs`'s
//! `golden_direct_and_dia_decode_agree`); a future re-cut that made them
//! disagree is a loud finding, not a silent pass.
//!
//! [`multilabel`] therefore consumes only the ORDERING of each row. It
//! never reads a magnitude, which is what keeps the decode correct in
//! spite of the next section.
//!
//! # fp16 safety of the shipped graph (one mechanism removed, not a proof)
//!
//! The artifact this module loaded until 2026-07-26 reached the same
//! log-probabilities through `softmax` → `log(epsilon = 0x0p+0)`. Every
//! epsilon below fp16's smallest subnormal (`2^-24` ≈ 5.96e-8) is inert
//! wherever the graph actually computes in fp16 — i.e. on the default
//! [`crate::ComputeUnits::All`] — so a softmax output that underflowed to 0
//! reached an unguarded `log(0)` and saturated instead of being clamped.
//! Measured over 1033 chunks of a real 17-minute recording
//! (`09_mrbeast_dollar_date`), the minimum `segments` value was **−45440** on
//! `All` against **−32.31** on `CpuOnly` (issue #15).
//!
//! The shipped graph has **no `log` op at all**: `reduce_log_sum_exp` → `sub`
//! computes `z − logsumexp(z)` directly. **That removes the mechanism above,
//! and removing that mechanism is the whole claim.** No `log` means no guard
//! epsilon, so the measured "underflowed softmax reaches `log(0)`" path cannot
//! occur.
//!
//! It is NOT a proof that no value can saturate at any precision on any
//! placement, and the wording that said so was reasoning one step past its
//! evidence. `sub` is not immune on its own terms: two finite fp16 operands
//! near the format's ±65504 maximum subtract to ≈ ∓131008, outside fp16's
//! range — and nothing here bounds the intermediates the trunk feeds it.
//! What the repo's gates actually bound is narrower than the old claim
//! sounded: `tests/fp16_guards.rs` bounds guard CONSTANTS against the `2^-24`
//! floor, and `tests/speaker/model_io.rs` bounds the artifact's BYTES (sha256)
//! and the tail's STRUCTURE (no `log` op, `reduce_log_sum_exp` present). None
//! of the three bounds a runtime value.
//!
//! - *Established*: the decomposed-softmax `log(0)` mechanism is gone, and on
//!   the same 1033-chunk measurement the minimum `segments` value on `All` is
//!   **−31.80** rather than −45440.
//! - *Not established*: any bound on `segments` for other inputs, other
//!   placements, or other hosts. An observed range on one recording is not a
//!   range guarantee.
//! - *What would settle it*: explicit finite/range testing — drive the graph
//!   at each placement with inputs chosen to push the trunk toward the fp16
//!   extremes and assert a bound on the output. This crate has no such test.
//!
//! **What this repair did NOT fix, and what this graph is NOT responsible
//! for.** The 8-speaker clip on which the shipping CoreML path returns 5
//! speakers at 16.6 % DER (the clip-09 pin in `coremlit-parity`'s
//! `tests/speaker/parity_shipping_der.rs`) is not repaired by the swap: re-running that suite's four gated clips
//! with only the artifact changed reproduced every gated number, clip 09
//! included.
//!
//! It is also **not this stage's collapse**. Measured at the shipping
//! configuration by `coremlit-parity`'s `tests/speaker/backend_factorial.rs` — int8 embedder, both
//! CoreML models on [`crate::ComputeUnits::All`], dia's ONNX as the reference
//! for whichever stage is not swapped — this graph on its own OVERcounts by
//! one speaker (9 of 8, 1.3011 % DER), while the CoreML *embedding* path on
//! its own reproduces the 5-of-8 collapse exactly (16.5904 %). An earlier
//! cross-product run with the fp32 embedder on `CpuOnly` attributed the whole
//! thing here; it does not survive re-measurement at the configuration this
//! crate ships. Within that embedding path the responsibility splits further —
//! the int8 palettization costs 2 speakers and the `All` placement 1, while the
//! embedding CONVERSION on its own is frame-perfect against dia-ort. See
//! `tests/speaker/model_io.rs`'s "Clip 09" section for all three tables side by
//! side.
//!
//! Which PART of this graph produces its own overcount is a further step, and
//! it is NOT established: `segments` is the only tensor either graph exposes,
//! so a divergence in it cannot be attributed to this tail rather than to the
//! trunk that feeds it. `backend_factorial`'s `seg_divergence` doc carries the
//! one decomposition the observable outputs do support (a monotone per-row
//! shift can only collapse an ordering into a tie, never invert one) and what
//! would settle the rest.
//!
//! What callers get: [`multilabel`] consumes per-row ORDERING only, and that
//! was never at risk from the epsilon. Magnitudes are unsaturated in the
//! measurements above, but they are still computed in fp16 on the ANE/GPU
//! (`All` can return a value a few thousandths above 0 where one class
//! dominates), so read them as fp16-precise log-probabilities — not as an
//! fp32-accurate log-likelihood, and not as a quantity this crate bounds.
//!
//! # Tie handling
//!
//! [`multilabel`]'s argmax uses strict `>` against a running max seeded
//! from class 0, exactly mirroring dia's `powerset_to_speakers_hard`
//! (`diarization/src/segment/powerset.rs:69-76`): on an exact tie, the
//! class with the LOWEST index wins, because a later equal value never
//! satisfies `p > max`.

use std::path::Path;

use crate::{ComputeUnits, DataType, Model, MultiArray};

use crate::model::contract::{Checked, Dim, FeatureContract, LoadContract, StateContract};

use crate::audio::speaker::error::{
  InferError, InputLength, ModelError, OutputShape, contract_violation,
};

/// Sample count of one segmentation-model chunk (10 s at 16 kHz). Matches
/// dia's `WINDOW_SAMPLES` (`diarization/src/segment/options.rs:18`) and the
/// introspected `pyannote_segmentation.mlmodelc` `audio` input's fixed
/// shape `[1, 1, 160_000]` (design spec §4; pinned by
/// `tests/model_io.rs::pyannote_segmentation_io_matches_spec`).
pub const SEG_CHUNK_SAMPLES: usize = 160_000;

/// Maximum simultaneous speakers the powerset encoding represents. Matches
/// dia's `MAX_SPEAKER_SLOTS` (`diarization/src/segment/options.rs:43`) —
/// `usize` here (not dia's `u8`) because every use in this crate is a
/// slice length or index, not a compact wire value.
pub const SEG_NUM_SLOTS: usize = 3;

/// Powerset class count: silence, A, B, C, A+B, A+C, B+C. Matches dia's
/// `POWERSET_CLASSES` (`diarization/src/segment/options.rs:40`) and the
/// introspected `segments` output's trailing dimension (design spec §4;
/// `[1, 589, 7]`).
pub const POWERSET_CLASSES: usize = 7;

/// Declared feature names on `pyannote_segmentation.mlmodelc`
/// (pinned by `tests/model_io.rs::pyannote_segmentation_io_matches_spec`).
mod names {
  pub const AUDIO: &str = "audio";
  pub const SEGMENTS: &str = "segments";
}

/// Default [`SegmentModelOptions::compute`]. `ComputeUnits::All` lets
/// CoreML schedule across ANE/GPU/CPU — the whole point of this crate
/// (design spec §1's ~20x segmentation uplift target). Model-gated tests
/// in this module instead load with `ComputeUnits::CpuOnly` for
/// determinism (no ANE compile-latency variance across runs), matching
/// `tests/model_io.rs`'s introspection convention (every load there also
/// uses `ComputeUnits::CpuOnly`) — production code keeps this default.
pub const DEFAULT_SEGMENT_COMPUTE: ComputeUnits = ComputeUnits::All;

#[cfg(feature = "serde")]
fn default_segment_compute() -> ComputeUnits {
  DEFAULT_SEGMENT_COMPUTE
}

/// Construction options for [`SegmentModel`] (rust-options-pattern).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct SegmentModelOptions {
  #[cfg_attr(feature = "serde", serde(default = "default_segment_compute"))]
  compute: ComputeUnits,
}

impl Default for SegmentModelOptions {
  fn default() -> Self {
    Self::new()
  }
}

impl SegmentModelOptions {
  /// Options matching the crate's default: [`DEFAULT_SEGMENT_COMPUTE`]
  /// (`ComputeUnits::All`).
  pub const fn new() -> Self {
    Self {
      compute: DEFAULT_SEGMENT_COMPUTE,
    }
  }

  /// Which hardware CoreML may schedule the segmentation model on.
  #[inline(always)]
  pub const fn compute(&self) -> ComputeUnits {
    self.compute
  }
  /// Builder form of [`Self::set_compute`].
  #[must_use]
  #[inline(always)]
  pub const fn with_compute(mut self, compute: ComputeUnits) -> Self {
    self.set_compute(compute);
    self
  }
  /// Sets [`Self::compute`] in place.
  #[inline(always)]
  pub const fn set_compute(&mut self, compute: ComputeUnits) -> &mut Self {
    self.compute = compute;
    self
  }
}

/// CoreML wrapper over `pyannote_segmentation.mlmodelc`: one
/// [`SEG_CHUNK_SAMPLES`]-sample chunk in, flattened `[num_frames *
/// POWERSET_CLASSES]` powerset **log-probabilities** out (the graph's tail
/// is `reduce_log_sum_exp` → `sub`; see the module doc) — layout-identical to dia's
/// `SegmentModel::infer` (`diarization/src/segment/model.rs:280-357`; see
/// the module doc's "dia contract match" section).
#[derive(Debug)]
pub struct SegmentModel {
  /// A `Checked`, never a bare [`Model`]: `segment_contract` below is the
  /// only way one is built, so removing the check from [`Self::from_file_with`]
  /// does not compile.
  model: Checked,
  num_frames: usize,
}

impl SegmentModel {
  /// Loads the model with [`SegmentModelOptions::new`] (`ComputeUnits::All`).
  ///
  /// # Errors
  /// As [`Self::from_file_with`].
  pub fn from_file(path: impl AsRef<Path>) -> Result<Self, ModelError> {
    Self::from_file_with(path, SegmentModelOptions::new())
  }

  /// Loads the model with custom options, checking it against this door's load
  /// contract.
  ///
  /// The model is held as a crate-internal `Checked` wrapper whose ONLY
  /// constructor runs that check. `segment_contract` IS this door's
  /// statement of what it needs:
  ///
  /// ```text
  /// input   audio     f32  [1, 1, 160_000]  every axis Exactly
  /// output  segments  f32  [1, F, 7]        F read back, at least 1
  /// state   none
  /// ```
  ///
  /// `F` is [`Self::num_frames`] — the ARTIFACT's, never hardcoded (589 on the
  /// staged `pyannote_segmentation.mlmodelc`).
  ///
  /// # What stating `F` as a contract axis newly refuses
  ///
  /// The check this replaced bound `F` from `segments.shape()[1]` with **no
  /// constraint check at all**, and [`crate::FeatureInfo::shape`] reports the
  /// DEFAULT shape of a flexible feature. A `RangeDims` `segments` head
  /// declared over 1..=4096 frames but defaulting to 589 therefore passed
  /// every clause, bound `F = 589`, and made every [`Self::infer`] allocate and
  /// range-check its output at that length — issue #137's defect (ii), the one
  /// that does not fail but computes against the wrong geometry.
  /// `Dim::AnyFixed` requires the whole feature to be
  /// [`crate::ShapeConstraint::Fixed`], so such a graph is refused at LOAD; its
  /// own clause also keeps the degenerate zero-frame case out, which the
  /// hand-written `>= 1` beside the old check did: an axis pinned at zero
  /// admits exactly one size, so only `Dim::AnyFixed`'s zero refusal sees it.
  ///
  /// The contract is also COMPLETE over what can make a conformant prediction
  /// fail rather than only over the features this door names: a graph carrying
  /// `audio` plus another REQUIRED input, or declaring a state buffer, is
  /// refused too — neither of which the shape checks this replaced could see.
  ///
  /// # Errors
  /// [`ModelError::Load`] if CoreML rejects the model;
  /// [`ModelError::ContractMismatch`] if a named feature is absent or its
  /// element type, rank, flexibility or a single axis is not the contract's;
  /// [`ModelError::UnsatisfiableInput`] if the graph requires an input this
  /// door never sends; [`ModelError::UnsatisfiableState`] if it declares a
  /// state buffer.
  pub fn from_file_with(
    path: impl AsRef<Path>,
    options: SegmentModelOptions,
  ) -> Result<Self, ModelError> {
    let model = Model::load(path, options.compute())?;
    let model = Checked::new(model, &segment_contract()).map_err(contract_violation)?;
    // Read back AFTER the check — see `Dim::AnyFixed`. The contract names
    // `segments` at rank 3, so both lookups are established by the check that
    // just passed.
    let num_frames = model
      .description()
      .output(names::SEGMENTS)
      .and_then(|segments| segments.shape().get(1).copied())
      .expect("the contract names `segments` at rank 3 and the check passed");

    Ok(Self { model, num_frames })
  }

  /// Output frame count for one chunk — the introspected `segments`
  /// shape's middle dimension (589 for `pyannote_segmentation.mlmodelc`,
  /// pinned by `tests/model_io.rs::pyannote_segmentation_io_matches_spec`;
  /// read dynamically at construction, not hardcoded).
  #[inline(always)]
  pub const fn num_frames(&self) -> usize {
    self.num_frames
  }

  /// Runs one segmentation chunk, returning flattened `[num_frames *
  /// POWERSET_CLASSES]` powerset **log-probabilities**, row-major
  /// `[frame][class]` — layout-identical to dia's `SegmentModel::infer`
  /// (`diarization/src/segment/model.rs:280-357`).
  ///
  /// These are `log(softmax(z))`, not the raw logits `z`: the shipped
  /// graph's tail is the FUSED `reduce_log_sum_exp` → `sub` (module doc,
  /// "`segments` is `log_softmax(·)`"). Consume the per-row ORDERING — the
  /// property [`multilabel`] needs, and the one a per-row constant shift
  /// preserves — rather than the magnitudes, which are computed in fp16 on
  /// ANE/GPU placements and are not bounded by anything this crate tests
  /// (module doc, "fp16 safety of the shipped graph").
  ///
  /// # Errors
  /// [`InferError::InputLength`] unless `samples.len() == SEG_CHUNK_SAMPLES`
  /// — see the module doc's "dia contract match" section for why this
  /// rejects rather than pads. [`InferError::NonFiniteInput`] if any input
  /// sample is NaN or infinite, scanned before the CoreML call (module doc,
  /// "Non-finite inputs and outputs"). [`InferError::Prediction`] /
  /// [`InferError::Tensor`] on a CoreML or tensor-construction failure —
  /// including a prediction whose runtime output set omits `segments`
  /// entirely ([`crate::PredictionError::MissingOutput`]; the
  /// construction-time contract pins the *declared* outputs, but the
  /// runtime provider's name set is CoreML's to produce per call).
  /// [`InferError::OutputShape`] if the predict-time `segments` tensor's
  /// shape diverges from `[1, num_frames, POWERSET_CLASSES]`. The
  /// construction-time contract validated once in [`Self::from_file_with`]
  /// is a claim about the model's declared shape, not a guarantee about
  /// every individual prediction; the CoreML runtime is a trust boundary
  /// re-checked here on every call, exactly as dia's `infer` re-validates
  /// output layout on every call
  /// (`diarization/src/segment/model.rs:313-338`). This also covers what
  /// [`crate::MultiArray::copy_into`] cannot: it validates only total
  /// element count, so an axes-swapped runtime output would otherwise be
  /// silently transposed into `logits` instead of erroring.
  /// [`InferError::NonFiniteOutput`] if any output logit is NaN or
  /// infinite — the exact `ort` CoreML-EP corruption mode (design spec §1)
  /// this crate exists to replace.
  pub fn infer(&self, samples: &[f32]) -> Result<Vec<f32>, InferError> {
    check_input_length(samples.len())?;
    check_finite_input(samples)?;

    let audio = MultiArray::from_slice(&[1, 1, SEG_CHUNK_SAMPLES], samples)?;
    let mut outputs = self.model.predict_with(&[(names::AUDIO, &audio)])?;
    let segments = outputs
      .take(names::SEGMENTS)
      .ok_or_else(|| crate::PredictionError::MissingOutput(names::SEGMENTS.to_string()))?;
    // The construction-time contract pins the model's DECLARED shape; the
    // CoreML runtime producing a specific prediction's tensor is a
    // separate trust boundary, re-checked here on every call exactly as
    // dia's `infer` re-validates output layout on every call
    // (`diarization/src/segment/model.rs:313-338`). `copy_into` below only
    // validates total element count, so this must run first.
    check_output_shape(segments.shape(), self.num_frames)?;

    let mut logits = vec![0.0f32; self.num_frames * POWERSET_CLASSES];
    segments.copy_into::<f32>(&mut logits)?;
    check_finite(&logits)?;

    Ok(logits)
  }
}

/// Validates `infer`'s input-length contract in isolation — hermetically
/// testable without a loaded model (`infer` itself needs `&self`, i.e. a
/// real loaded [`SegmentModel`], to reach the CoreML call this guards).
fn check_input_length(got: usize) -> Result<(), InferError> {
  if got != SEG_CHUNK_SAMPLES {
    return Err(InferError::InputLength(InputLength::new(
      got,
      SEG_CHUNK_SAMPLES,
    )));
  }
  Ok(())
}

/// Validates a predict-time `segments` tensor's shape against the
/// construction-time contract (`[1, num_frames, POWERSET_CLASSES]`) in
/// isolation — hermetically testable without a loaded model (`infer`
/// itself needs `&self`, i.e. a real loaded [`SegmentModel`], to reach the
/// CoreML call this guards). Catches exactly what
/// [`crate::MultiArray::copy_into`] cannot: it validates only total
/// element count, so an axes-swapped `[1, POWERSET_CLASSES, num_frames]`
/// tensor — the same element count as the expected
/// `[1, num_frames, POWERSET_CLASSES]` — would otherwise pass `copy_into`
/// silently and transpose frames and classes into `logits`.
fn check_output_shape(shape: &[usize], num_frames: usize) -> Result<(), InferError> {
  let shape_ok =
    shape.len() == 3 && shape[0] == 1 && shape[1] == num_frames && shape[2] == POWERSET_CLASSES;
  if !shape_ok {
    return Err(InferError::OutputShape(OutputShape::new(
      shape.to_vec(),
      vec![1, num_frames, POWERSET_CLASSES],
    )));
  }
  Ok(())
}

/// Scans `logits` for the first non-finite value — "the exact `ort`
/// CoreML-EP corruption mode this crate exists to replace" (see the module
/// doc). Extracted from [`SegmentModel::infer`] so it is hermetically
/// testable without a loaded model.
fn check_finite(logits: &[f32]) -> Result<(), InferError> {
  if let Some(index) = logits.iter().position(|v| !v.is_finite()) {
    return Err(InferError::NonFiniteOutput(index));
  }
  Ok(())
}

/// Scans `samples` for the first non-finite value, BEFORE the CoreML call —
/// hermetically testable without a loaded model. Mirrors the embed module's
/// identical input-side scan and dia's own `Error::NonFiniteInput` guard
/// (`diarization/src/segment/model.rs:283-290`): a NaN sample would otherwise
/// reach CoreML, and while it usually propagates to a non-finite logit the
/// output scan catches, a kernel could absorb it into finite garbage — see
/// the module doc's "Non-finite inputs and outputs" section.
fn check_finite_input(samples: &[f32]) -> Result<(), InferError> {
  if let Some(index) = samples.iter().position(|v| !v.is_finite()) {
    return Err(InferError::NonFiniteInput(index));
  }
  Ok(())
}

/// Powerset class index → hard 3-slot speaker mask. Order and encoding
/// match dia's `TABLE` exactly (`diarization/src/segment/powerset.rs:
/// 77-85`): index 0 is silence, 1..=3 are single speakers A/B/C, 4..=6 are
/// the three two-speaker overlaps A+B/A+C/B+C. `f64` here (dia's table is
/// `f32`) matches dia's `OfflineInput::segmentations` dtype — see the
/// module doc.
const POWERSET_TABLE: [[f64; SEG_NUM_SLOTS]; POWERSET_CLASSES] = [
  [0.0, 0.0, 0.0], // silence
  [1.0, 0.0, 0.0], // A
  [0.0, 1.0, 0.0], // B
  [0.0, 0.0, 1.0], // C
  [1.0, 1.0, 0.0], // A+B
  [1.0, 0.0, 1.0], // A+C
  [0.0, 1.0, 1.0], // B+C
];

/// Hard argmax over one frame's [`POWERSET_CLASSES`] logits, ties broken
/// toward the lowest class index — see the module doc's "Tie handling"
/// section.
///
/// The plain `>` comparison is dia's exact semantics
/// (`diarization/src/segment/powerset.rs:72`), including for non-finite
/// input (which [`SegmentModel::infer`] has already rejected on the
/// production path, exactly as dia's own callers do): `NaN > x` and
/// `x > NaN` are both false, so a NaN never displaces the running max, and
/// a NaN *seed* (class 0) sticks there. A NaN-total-ordering comparison
/// (`total_cmp`) would NOT be equivalent — it also orders `-0.0 < +0.0`,
/// breaking dia's lowest-index rule on that exact-tie case.
fn hard_argmax(row: &[f32; POWERSET_CLASSES]) -> usize {
  let mut argmax = 0usize;
  let mut max = row[0];
  for (k, &v) in row.iter().enumerate().skip(1) {
    if v > max {
      max = v;
      argmax = k;
    }
  }
  argmax
}

/// Powerset argmax → hard multilabel over a whole `[num_frames *
/// POWERSET_CLASSES]` logits buffer (e.g. [`SegmentModel::infer`]'s return
/// value): `[num_frames * SEG_NUM_SLOTS]` flattened `f64` 0.0/1.0,
/// frame-major (`frame * SEG_NUM_SLOTS + slot`) — the same layout dia's
/// `segmentations` buffer uses for one chunk
/// (`diarization/src/offline/owned.rs:496`:
/// `segs[(c * FRAMES_PER_WINDOW + f) * SLOTS_PER_CHUNK + s]`, `c = 0`). See
/// the module doc's "Why argmaxing it without a softmax is still correct"
/// and "Tie handling" sections for why this needs no softmax of its own
/// (the input is already `log(softmax(·))`, i.e. the logits shifted by a
/// per-row constant) and how ties resolve.
///
/// `logits` is expected to be finite, as [`SegmentModel::infer`] guarantees
/// for its own output — the same precondition dia's callers establish
/// before its powerset decode (`Segmenter::push_inference` rejects
/// `NonFiniteScores`; the owned pipeline's `infer` rejects
/// `NonFiniteOutput`). Non-finite values are not rejected here, matching
/// dia: a NaN simply never wins the argmax (see the private `hard_argmax`
/// helper's doc for the exact comparison semantics).
///
/// # Panics
/// Panics if `logits.len() != num_frames * POWERSET_CLASSES` — in every
/// build profile, mirroring dia's own fail-fast (dia's inline decode
/// indexes `logits[f * POWERSET_CLASSES + k]` directly,
/// `diarization/src/offline/owned.rs:482`, which panics on a short buffer
/// in release too; silently truncating instead would hand downstream
/// consumers a misaligned `segmentations` buffer).
pub fn multilabel(logits: &[f32], num_frames: usize) -> Vec<f64> {
  assert_eq!(
    logits.len(),
    num_frames * POWERSET_CLASSES,
    "logits.len() must equal num_frames * POWERSET_CLASSES"
  );
  let mut out = Vec::with_capacity(num_frames * SEG_NUM_SLOTS);
  for row in logits.as_chunks::<POWERSET_CLASSES>().0 {
    out.extend_from_slice(&POWERSET_TABLE[hard_argmax(row)]);
  }
  out
}

/// The load contract this door states: `audio [1, 1, 160_000]` f32 in,
/// `segments [1, F >= 1, 7]` f32 out, no state.
///
/// Data rather than a sequence of checks, and the ONLY thing
/// [`SegmentModel::from_file_with`] does beyond [`Model::load`] — see that
/// method for what the axis clauses newly refuse.
fn segment_contract() -> LoadContract {
  LoadContract::new(
    vec![FeatureContract::new(
      names::AUDIO,
      DataType::F32,
      vec![
        Dim::Exactly(1),
        Dim::Exactly(1),
        Dim::Exactly(SEG_CHUNK_SAMPLES),
      ],
    )],
    vec![FeatureContract::new(
      names::SEGMENTS,
      DataType::F32,
      vec![
        Dim::Exactly(1),
        Dim::AnyFixed,
        Dim::Exactly(POWERSET_CLASSES),
      ],
    )],
    StateContract::None,
  )
}

#[cfg(test)]
mod tests;