qsm_core/unwrap/

romeo.rs

//! ROMEO weight calculation for phase unwrapping
//!
//! Calculates edge weights for region-growing phase unwrapping based on:
//! - Phase coherence
//! - Phase gradient coherence (multi-echo)
//! - Magnitude coherence
//! - Magnitude weights
//!
//! Reference:
//! Dymerska, B., Eckstein, K., Bachrata, B., Siow, B., Trattnig, S., Shmueli, K.,
//! Robinson, S.D. (2021). "Phase unwrapping with a rapid opensource minimum spanning
//! tree algorithm (ROMEO)." Magnetic Resonance in Medicine, 85(4):2294-2308.
//! https://doi.org/10.1002/mrm.28563
//!
//! Reference implementation: https://github.com/korbinian90/MriResearchTools.jl
//! (`romeo`/`unwrap` re-exported from <https://github.com/korbinian90/ROMEO.jl>).
//!
//! # Reference pipeline
//!
//! Facts from ROMEO.jl (commit 08776d4, 2026-09-02) and MriResearchTools.jl (771eadf), the
//! authors' code; file:line refer to those trees.
//!
//! - **Weights.** `:romeo` (phase coherence, phase gradient coherence, phase linearity,
//!   magnitude coherence) is the API's default (`src/weights.jl:56,75-76`), which
//!   [`RomeoParams::default`] matches; `:romeo4` sets the same four flags (`src/weights.jl:81-82`).
//!   The CLI's `-w romeo` default becomes `romeo3` (no phase linearity) whenever a magnitude is
//!   given and `romeo4` otherwise (`ext/RomeoApp/caller.jl:49-55`): [`RomeoParams::romeo_cli`].
//! - **Weight bins.** Each weight in `[0, 1]` goes to the bin `max(round((1 − w)·255), 1)`
//!   (`rescale`, `src/weights.jl:139-146`); the bins order the growth, so they are ROMEO.jl's here
//!   (stored as `256 − bin`). An edge is weighted when its first voxel is in the mask
//!   (`mask[I] && checkbounds(J)`, `src/weights.jl:99-100`), with the magnitude masked, so the
//!   growth can step one voxel out of the mask along +x, +y or +z (and through such voxels).
//! - **Multi-echo mode.** Temporal (template) unwrapping is the default: the template echo
//!   (`template=1`, i.e. index 0 here) is unwrapped spatially, with the echo after it (or before
//!   it, for a later template) as `phase2` for the gradient-coherence weight, and the others
//!   temporally outward from it (`src/unwrapping.jl:68-93`). Individual unwrapping is opt-in,
//!   and the CLI says it "might be necessary if there is large movement (timeseries) or
//!   phase-offset-correction is not applicable" (`argparse.jl:95-100`). Both match
//!   [`RomeoParams::default`] (`template: 0`, `individual: false`).
//! - **`temporal_uncertain_unwrapping`.** 0.5 in the API (`src/unwrapping.jl:52,70`), which
//!   [`RomeoParams::default`] and CLEAR-SWI's `romeo_combine` use; the CLI's default is 0 and it
//!   marks the option EXPERIMENTAL (`argparse.jl:165-172`); [`RomeoParams::romeo_cli`] has 0.
//! - **`correct_global`** is off by default in both the API and the CLI (`src/unwrapping.jl:43`,
//!   `argparse.jl:129-135`). (QSMxT's Python releases ran the CLI with
//!   `--individual-unwrapping --correct-global`, `qsmxt/interfaces/nipype_interface_romeo.py:66`;
//!   that is QSMxT's choice, not ROMEO's default.) Where it applies, as in ROMEO.jl's 3D `unwrap!`:
//!   to a single volume, to the template echo before the others are unwrapped temporally from
//!   it, and in individual mode to each echo and then across echoes
//!   ([`correct_multi_echo_wraps`]); see `unwrap_individual` for the two places (threads and even
//!   medians) where QSM.rs is deterministic and ROMEO.jl is not.
//! - **Undefined phase (NaN).** ROMEO.jl takes NaN phase (e.g. MCPC-3D-S's output where its
//!   offset is undefined, [`crate::utils::phase_offset_removal`]) as "no data": any edge touching
//!   it gets a NaN weight, which `rescale` makes invalid, so the growth never reaches it and it
//!   stays NaN; the global corrections skip non-finite voxels. QSM.rs does the same, including
//!   Julia's NaN-propagating `max(0, x)` in the gradient-coherence and phase-linearity weights
//!   (phase linearity is 0.5 next to a NaN, as in ROMEO.jl; before 2026-11 it was 0 here).
//! - **Mask.** The CLI unwraps inside `robustmask` of the template echo's magnitude by default
//!   (`argparse.jl:38-42`, `caller.jl:201-204`), a magnitude threshold with hole filling
//!   (MriResearchTools `src/masking.jl:7-29`), not a brain mask; with no magnitude it unwraps
//!   unmasked. CLEAR-SWI calls ROMEO with no mask at all
//!   (CLEARSWI.jl `src/phase_processing.jl:122`). The QSM.rs pipeline passes its brain mask.
//! - **Phase offsets.** When the CLI computes a B0 map from multi-echo data it turns MCPC-3D-S
//!   phase-offset correction on (monopolar) unless told otherwise (`caller.jl:80-91`), with
//!   smoothing sigma 7 mm converted to voxels (`argparse.jl:79-91`, `caller.jl:137-141`);
//!   bipolar correction is opt-in and needs ≥ 3 echoes (`argparse.jl:71-76`). MCPC-3D-S itself
//!   is in [`crate::utils::phase_offset_removal`]; its defaults are in the module docs of
//!   `utils::multi_echo`. `calculateB0_unwrapped` states that "the phase offsets have to be
//!   removed prior" (MriResearchTools `src/romeofunctions.jl:3-7`).
//! - **Seed.** ROMEO.jl seeds at the voxel whose three forward edges have the best summed bin
//!   (ties to the highest linear index) and corrects the seed's 2π count from `phase2` with
//!   multi-echo input, or wraps it into (−π, π] otherwise (`src/seed.jl`). So does every path here
//!   ([`crate::region_grow::grow_region_unwrap_full`]). Before, the default single-seed path
//!   seeded at the mask's centre of mass ([`find_seed_point`], no longer used here), which could
//!   change the global 2π offset and the growth order.
//! - **Resolution, field strength, orientation.** Nothing in ROMEO depends on them: the weights
//!   use phase differences between 6-connected neighbours and echo-time ratios only.
//!
//! # Agreement with the reference
//!
//! In double precision QSM.rs reproduces ROMEO.jl's `unwrap` exactly (the same number of turns at
//! every voxel) in template mode with the API's and the CLI's settings and in individual mode
//! with global correction, and MriResearchTools' `mcpc3ds` to 9e-16 rad
//! (`tests/romeo_romeojl_parity.rs`, a noisy three-echo phantom; ROMEO.jl 08776d4 and
//! MriResearchTools 771eadf in the QSM-CI `julia-qsm` image `sha256:b3c048ac…`). Given the
//! CLI's own MCPC-3D-S output on the QSM-CI in-vivo travelling-head scan (Prisma 2.89 T, 1 mm,
//! five echoes), QSM.rs's template unwrapping gives the CLI's unwrapped echoes exactly.
//!
//! `tests/romeo_validity_romeojl_parity.rs` (ROMEO.jl 1.6.0, MriResearchTools 3.8.0, same image, one
//! Julia thread) adds a phantom whose brain mask reaches past MCPC-3D-S's mask: the NaN validity,
//! individual mode with and without global correction, and template mode with and without it,
//! on the corrected and on the raw phase, all give ROMEO.jl's turns at every voxel and its B0 to
//! 1e-14 (measured 2026-11-04). **Known gap:** template mode's re-unwrapping of uncertain voxels
//! (`temporal_uncertain_unwrapping > 0`, the API default) is not yet ROMEO.jl's. ROMEO.jl rates the
//! voxels with `voxelquality` (Float32 weights before binning, no mask, the sum over six edges
//! divided by six, NaN for an edge touching a NaN) and grows from every voxel it marks visited,
//! those outside the mask included (`temporal_uncertain_unwrapping!`); `temporal_uncertain_rewrap`
//! uses [`voxel_quality_romeo`] (binned weights, in-mask neighbours only) and grows from the
//! certain voxels inside the mask. On that phantom 14 of 3,000 voxels of echoes 2 and 3 differ
//! after MCPC-3D-S and 850 of 3,392 on the raw phase; the fixture's template cases therefore set
//! it to 0 (the CLI's default). QSMxT's default (individual mode) does not use it.
//!
//! The `romeo` CLI itself reads the NIfTI phase and magnitude as Float32 and computes in single
//! precision; QSM.rs computes in f64 (kept: double precision throughout, Ashley 2026-10-28).
//! That alone leaves a residual on the in-vivo scan. The CLI's defaults (`romeo -B` with a brain
//! mask) against the pipeline with [`RomeoParams::romeo_cli`]: B0 NRMSE 2.9 % (corr 0.9995),
//! from 2π differences in ~1,250 voxels of the MCPC-3D-S HIP unwrap that propagate through the
//! offset smoothing (fed the CLI's own single-precision HIP, QSM.rs differs in 36 voxels; fed
//! its corrected phase, in none). Before the seed, weight bins, edge masking, `mcpc3ds`'s own
//! mask and smoothing were ROMEO.jl's, the same comparison gave NRMSE 25 %; with individual
//! unwrapping, global correction, `romeo4` and a 4 mm offset smoothing on both sides it went
//! from 6.4 % to 1.1 %.

use std::f64::consts::PI;

use crate::region_grow::{grow_region_unwrap_from_visited, grow_region_unwrap_full};
use crate::Grid;
#[cfg(feature = "parallel")]
use rayon::prelude::*;

/// Weight calculation scheme for ROMEO unwrapping.
///
/// Matches ROMEO.jl weight type options. See Dymerska et al. (2021) for details.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum RomeoWeightType {
    /// Phase coherence + phase gradient coherence + phase linearity + mag coherence
    /// (ROMEO.jl's default `:romeo`, `src/weights.jl:75-76`).
    Romeo,
    /// Phase coherence + mag coherence only (no multi-echo temporal info).
    Romeo2,
    /// Phase coherence + phase gradient coherence + mag coherence.
    Romeo3,
    /// Alias for Romeo.
    Romeo4,
    /// All components including magnitude weighting.
    Romeo6,
    /// Best-path reliability weights (Abdul-Rahman), used in place of the ROMEO
    /// weight components. Equivalent to setting [`RomeoParams::bestpath`].
    ///
    /// This is still ROMEO's region growing, just differently weighted. For the
    /// original best-path algorithm, see
    /// [`unwrap_bestpath`](crate::unwrap::unwrap_bestpath).
    BestPath,
}

impl RomeoWeightType {
    /// Map weight type to 6 component flags matching ROMEO.jl:
    /// \[phase_coherence, phase_gradient_coherence, phase_linearity,
    ///  mag_coherence, mag_weight, mag_weight2\]
    fn weight_flags(&self) -> [bool; 6] {
        match self {
            RomeoWeightType::Romeo  => [true, true, true, true, false, false],
            RomeoWeightType::Romeo2 => [true, false, false, true, false, false],
            RomeoWeightType::Romeo3 => [true, true, false, true, false, false],
            RomeoWeightType::Romeo4 => [true, true, true, true, false, false],
            RomeoWeightType::Romeo6 => [true, true, true, true, true, true],
            RomeoWeightType::BestPath => [false; 6], // uses different calculation
        }
    }
}

/// Parameters for ROMEO phase unwrapping.
///
/// Weight components can be toggled individually. Each multiplies into the
/// final edge weight as `0.1 + 0.9 * component_value`.
///
/// The default enables phase coherence, phase gradient coherence,
/// phase linearity, and magnitude coherence (equivalent to `:romeo` in ROMEO.jl).
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Debug)]
pub struct RomeoParams {
    // -- Weight component flags --

    /// Phase coherence: `1 - |wrap(Δφ)| / π`. Always recommended. (default: true)
    pub phase_coherence: bool,
    /// Phase gradient coherence: consistency between echo pairs. Needs multi-echo
    /// data (phase2 + TEs). Automatically disabled if unavailable. (default: true)
    pub phase_gradient_coherence: bool,
    /// Phase linearity: second-derivative smoothness along each edge. (default: true)
    pub phase_linearity: bool,
    /// Magnitude coherence: `(min/max)²` of neighbor magnitudes. (default: true)
    pub mag_coherence: bool,
    /// Magnitude weight: penalizes low-signal voxels relative to the 95th
    /// percentile. (default: false)
    pub mag_weight: bool,
    /// Magnitude weight 2: penalizes abnormally high signal (flow artifacts).
    /// (default: false)
    pub mag_weight2: bool,
    /// Use Best-path weights (Abdul-Rahman) instead of ROMEO weights.
    /// When true, the individual weight flags above are ignored. (default: false)
    pub bestpath: bool,

    // -- Multi-echo options --

    /// Template echo index for spatial unwrapping, 0-indexed (default: 0).
    /// Only used for multi-echo template-based unwrapping.
    pub template: usize,
    /// Unwrap each echo individually instead of template-based (default: false).
    pub individual: bool,
    /// Correct global 2π offsets between echoes (default: false).
    /// Recommended when using individual mode.
    pub correct_global: bool,
    /// Quality threshold for re-unwrapping uncertain voxels after temporal
    /// unwrapping. Range \[0, 1\], set to 0 to disable. (default: 0.5)
    pub temporal_uncertain_unwrapping: f64,
    /// Maximum number of seed regions (default: 1, max: 255).
    pub max_seeds: u8,
    /// Merge neighboring regions after unwrapping (default: false).
    pub merge_regions: bool,
    /// Correct each region's median to nearest 0 by adding n·2π (default: false).
    pub correct_regions: bool,
    /// Additional phase tolerance beyond π for neighbor differences.
    /// Range \[0, π\]. (default: 0.0)
    pub wrap_addition: f64,
}

impl Default for RomeoParams {
    fn default() -> Self {
        Self {
            // Default matches :romeo in ROMEO.jl
            phase_coherence: true,
            phase_gradient_coherence: true,
            phase_linearity: true,
            mag_coherence: true,
            mag_weight: false,
            mag_weight2: false,
            bestpath: false,
            template: 0,
            individual: false,
            correct_global: false,
            temporal_uncertain_unwrapping: 0.5,
            max_seeds: 1,
            merge_regions: false,
            correct_regions: false,
            wrap_addition: 0.0,
        }
    }
}

impl RomeoParams {
    /// The `romeo` command line's settings for multi-echo data with a magnitude (ROMEO.jl
    /// 08776d4), where they differ from the API's ([`RomeoParams::default`]):
    /// - weights `romeo3` (phase coherence, phase gradient coherence, magnitude coherence): the
    ///   CLI's `-w romeo` default becomes `romeo3` whenever a magnitude is given, and `romeo4`
    ///   (the API's `:romeo`) only without one (`load_data_and_resolve_args!`,
    ///   `ext/RomeoApp/caller.jl:49-55`);
    /// - `temporal_uncertain_unwrapping` 0, the CLI's default (`argparse.jl:165-172`), where the
    ///   API's is 0.5.
    ///
    /// Everything else (template echo 1, temporal unwrapping, no global correction, one seed)
    /// is the same in both.
    pub fn romeo_cli() -> Self {
        Self { phase_linearity: false, temporal_uncertain_unwrapping: 0.0, ..Self::default() }
    }

    /// Create params matching a ROMEO.jl weight type preset.
    pub fn from_weight_type(wt: RomeoWeightType) -> Self {
        if wt == RomeoWeightType::BestPath {
            return Self { bestpath: true, ..Default::default() };
        }
        let f = wt.weight_flags();
        Self {
            phase_coherence: f[0],
            phase_gradient_coherence: f[1],
            phase_linearity: f[2],
            mag_coherence: f[3],
            mag_weight: f[4],
            mag_weight2: f[5],
            ..Default::default()
        }
    }

    /// Convert the boolean fields to the 6-element flag array used internally.
    pub fn weight_flags(&self) -> [bool; 6] {
        [
            self.phase_coherence,
            self.phase_gradient_coherence,
            self.phase_linearity,
            self.mag_coherence,
            self.mag_weight,
            self.mag_weight2,
        ]
    }
}

const TWO_PI: f64 = 2.0 * PI;

/// Wrap angle to [-π, π]
#[inline]
fn wrap_angle(angle: f64) -> f64 {
    let mut a = angle % TWO_PI;
    if a > PI {
        a -= TWO_PI;
    } else if a < -PI {
        a += TWO_PI;
    }
    a
}

/// Index into a 3D array in Fortran order (column-major, matches NIfTI)
#[inline(always)]
fn idx3d(i: usize, j: usize, k: usize, nx: usize, ny: usize) -> usize {
    i + j * nx + k * nx * ny
}

/// Calculate ROMEO edge weights for phase unwrapping.
///
/// Uses the default Romeo weight type (PC + PGC + PL + MC).
/// Returns weights array of size 3 * nx * ny * nz in C order \[dim\]\[i\]\[j\]\[k\].
/// Weights are 1-255 (valid) or 0 (no edge / masked out).
pub fn calculate_weights_romeo(
    phase: &[f64],
    mag: &[f64],
    phase2: Option<&[f64]>,
    te1: f64,
    te2: f64,
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
) -> Vec<u8> {
    calculate_weights_romeo_with_flags(
        phase, mag, phase2, te1, te2, mask, nx, ny, nz,
        RomeoWeightType::Romeo.weight_flags(),
    )
}

/// Calculate ROMEO edge weights with configurable weight components.
///
/// Legacy 3-flag interface for backward compatibility.
pub fn calculate_weights_romeo_configurable(
    phase: &[f64],
    mag: &[f64],
    phase2: Option<&[f64]>,
    te1: f64,
    te2: f64,
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    use_phase_gradient_coherence: bool,
    use_mag_coherence: bool,
    use_mag_weight: bool,
) -> Vec<u8> {
    let flags = [
        true,                           // phase coherence (always on)
        use_phase_gradient_coherence,
        false,                          // phase linearity (use weight_type for this)
        use_mag_coherence,
        use_mag_weight,
        false,                          // mag_weight2
    ];
    calculate_weights_romeo_with_flags(
        phase, mag, phase2, te1, te2, mask, nx, ny, nz, flags,
    )
}

/// Calculate ROMEO edge weights with full 6-component flag control.
///
/// Matches ROMEO.jl `calculateweights_romeo`. The 6 flags are:
/// 0: phase coherence, 1: phase gradient coherence, 2: phase linearity,
/// 3: magnitude coherence, 4: magnitude weight, 5: magnitude weight 2.
///
/// Each component is scaled as `0.1 + 0.9 * value` (matching ROMEO.jl) to prevent
/// any single component from zeroing the weight entirely.
///
/// Weights are stored as u8: 0 = no edge, 1 = worst valid, 255 = best.
pub fn calculate_weights_romeo_with_flags(
    phase: &[f64],
    mag: &[f64],
    phase2: Option<&[f64]>,
    te1: f64,
    te2: f64,
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    flags: [bool; 6],
) -> Vec<u8> {
    let n_total = nx * ny * nz;
    let mut weights = vec![0u8; 3 * n_total];

    let has_mag = !mag.is_empty();
    let has_phase2 = phase2.is_some();

    // Disable magnitude flags if no magnitude data
    let f_mc  = flags[3] && has_mag;
    let f_mw  = flags[4] && has_mag;
    let f_mw2 = flags[5] && has_mag;
    // Disable PGC if no phase2/TEs
    let f_pgc = flags[1] && has_phase2;
    let f_pl  = flags[2];

    // 95th percentile of magnitude for normalization (matching ROMEO.jl)
    let max_mag = if has_mag && (f_mw || f_mw2) {
        percentile_95(mag, mask)
    } else {
        1.0
    };
    let half_max_mag = 0.5 * max_mag + 1e-12;

    for dim in 0..3_usize {
        // Each dim writes a disjoint n_total block; voxels within it are
        // independent, so fan out across voxels (rayon under `parallel`).
        let block = &mut weights[dim * n_total..(dim + 1) * n_total];
        maybe_par_iter_mut!(block).enumerate().for_each(|(idx, w)| {
            let i = idx % nx;
            let j = (idx / nx) % ny;
            let k = idx / (nx * ny);
            let (ni, nj, nk) = match dim {
                0 => (i + 1, j, k),
                1 => (i, j + 1, k),
                _ => (i, j, k + 1),
            };

            if ni >= nx || nj >= ny || nk >= nz {
                return;
            }

            let idx_n = idx3d(ni, nj, nk, nx, ny);

            // ROMEO.jl tests only the edge's first voxel (`mask[I] && checkbounds(J)`,
            // `src/weights.jl:99-100`), so an edge from the mask's last voxel to the next one
            // outside it is valid (weak: the magnitude is masked, so its magnitude coherence is
            // 0) and the growth can step one voxel out of the mask in +x, +y or +z, and back
            // in through such voxels.
            if mask[idx] == 0 {
                return;
            }

            let mut weight = 1.0_f64;

            // 1. Phase coherence: 1 - |wrap(diff)| / π
            if flags[0] {
                let pc = 1.0 - gamma(phase[idx_n] - phase[idx]).abs() / PI;
                weight *= 0.1 + 0.9 * pc;
            }

            // 2. Phase gradient coherence (multi-echo)
            if f_pgc {
                let phase2_data = phase2.unwrap();
                let wrapped_p1 = gamma(phase[idx_n] - phase[idx]);
                let wrapped_p2 = gamma(phase2_data[idx_n] - phase2_data[idx]);
                // ROMEO.jl: γ(P2[i] - P2[j]) * TEs[1] / TEs[2], left to right; its
                // `max(0, x)` keeps a NaN (Julia's `max` propagates NaN, Rust's `f64::max` drops
                // it), so a NaN phase2 makes the edge invalid there and here.
                let pgc = julia_max0(1.0 - (wrapped_p1 - wrapped_p2 * te1 / te2).abs());
                weight *= 0.1 + 0.9 * pgc;
            }

            // 3. Phase linearity: product of two triplet linearities
            if f_pl {
                let pl = phase_linearity_edge(phase, i, j, k, ni, nj, nk, dim, nx, ny, nz);
                weight *= 0.1 + 0.9 * pl;
            }

            if has_mag {
                // ROMEO.jl masks the magnitude (`mag .* mask`, `parsekwargs`)
                let m1 = mag[idx];
                let m2 = if mask[idx_n] != 0 { mag[idx_n] } else { 0.0 };
                let small = m1.min(m2);
                let big = m1.max(m2);

                // 4. Magnitude coherence: (min/max)²
                if f_mc {
                    // 0/0 is NaN in ROMEO.jl, which makes the edge invalid
                    let mc = if big > 0.0 { (small / big).powi(2) } else { f64::NAN };
                    weight *= 0.1 + 0.9 * mc;
                }

                // 5. Magnitude weight: penalize low signal
                if f_mw {
                    let mw = 0.5 + 0.5 * (small / half_max_mag).min(1.0);
                    weight *= 0.1 + 0.9 * mw;
                }

                // 6. Magnitude weight 2: penalize too-high signal (flow artifacts)
                if f_mw2 {
                    let mw2 = 0.5 + 0.5 * (half_max_mag / big.max(1e-12)).min(1.0);
                    weight *= 0.1 + 0.9 * mw2;
                }
            }

            // Rescale to u8: min valid = 1, best = 255, 0 = no edge
            *w = rescale_weight(weight);
        });
    }

    weights
}

/// Phase linearity for an edge between two voxels.
///
/// Computes the product of two triplet linearities: one looking "behind" the edge
/// and one looking "ahead". Matches ROMEO.jl `phaselinearity(P, i, j)`.
fn phase_linearity_edge(
    phase: &[f64],
    i: usize, j: usize, k: usize,
    ni: usize, nj: usize, nk: usize,
    _dim: usize,
    nx: usize, ny: usize, nz: usize,
) -> f64 {
    let idx = idx3d(i, j, k, nx, ny);
    let idx_n = idx3d(ni, nj, nk, nx, ny);

    // "Behind" triplet: (h, idx, idx_n) where h = 2*idx_pos - idx_n_pos
    let (hi, hj, hk) = (2 * i as i32 - ni as i32, 2 * j as i32 - nj as i32, 2 * k as i32 - nk as i32);
    let pl1 = if hi >= 0 && hi < nx as i32 && hj >= 0 && hj < ny as i32 && hk >= 0 && hk < nz as i32 {
        let h_idx = idx3d(hi as usize, hj as usize, hk as usize, nx, ny);
        phase_linearity_triplet(phase[h_idx], phase[idx], phase[idx_n])
    } else {
        0.9
    };

    // "Ahead" triplet: (idx, idx_n, k) where k = 2*idx_n_pos - idx_pos
    let (ki, kj, kk) = (2 * ni as i32 - i as i32, 2 * nj as i32 - j as i32, 2 * nk as i32 - k as i32);
    let pl2 = if ki >= 0 && ki < nx as i32 && kj >= 0 && kj < ny as i32 && kk >= 0 && kk < nz as i32 {
        let k_idx = idx3d(ki as usize, kj as usize, kk as usize, nx, ny);
        phase_linearity_triplet(phase[idx], phase[idx_n], phase[k_idx])
    } else {
        0.9
    };

    pl1 * pl2
}

/// Phase linearity of three consecutive phase values.
///
/// `max(0, 1 - |wrap(a - 2b + c) / 2|)` — measures how linear the phase is.
/// Matches ROMEO.jl `phaselinearity(P, i, j, k)`.
#[inline]
fn phase_linearity_triplet(a: f64, b: f64, c: f64) -> f64 {
    let second_deriv = wrap_angle(a - 2.0 * b + c);
    // `max(0, …)` with Julia's NaN propagation, then ROMEO.jl's `isnan(pl) && return 0.5`: a NaN
    // phase on either side of an edge (e.g. where MCPC-3D-S's offset is undefined) gives 0.5.
    // With Rust's `f64::max` the NaN became 0 and the 0.5 branch was dead (fixed 2026-11).
    let pl = julia_max0(1.0 - (second_deriv / 2.0).abs());
    if pl.is_nan() { 0.5 } else { pl }
}

/// ROMEO.jl's `γ` (`src/utility.jl`): one wrap only, `x ± 2π` when `x` is outside [−π, π].
/// For the difference of two wrapped phases it is the full wrap; for an unwrapped `phase2` (the
/// previous echo in individual mode) it is not, and ROMEO.jl's weights are this one's.
#[inline]
fn gamma(x: f64) -> f64 {
    if x < -PI {
        x + TWO_PI
    } else if x > PI {
        x - TWO_PI
    } else {
        x
    }
}

/// Julia's `max(0, x)`: NaN stays NaN (Rust's `f64::max(0.0)` would return 0).
#[inline]
fn julia_max0(x: f64) -> f64 {
    if x.is_nan() { x } else { x.max(0.0) }
}

/// Rescale a weight in `[0, 1]` (1 best) to the queue priority, as ROMEO.jl bins it.
///
/// ROMEO.jl's `rescale` (`src/weights.jl:139-146`, 08776d4) maps `w` to the bin
/// `max(round((1 − w)·255), 1)` (1 best, 255 worst, Julia's round-half-to-even) and anything
/// outside `[0, 1]` to 0 (invalid). Here the priority is `256 − bin`, so 255 is best, 1 worst
/// and 0 invalid, with exactly ROMEO.jl's bins (a weight of 0 is the worst valid edge, as there).
/// The bins decide the growth order, so they must be ROMEO.jl's for the same result.
#[inline]
fn rescale_weight(w: f64) -> u8 {
    if (0.0..=1.0).contains(&w) {
        let bin = (((1.0 - w) * 255.0).round_ties_even() as i64).max(1);
        (256 - bin) as u8
    } else {
        0
    }
}

/// Compute 95th percentile of magnitude within mask (matching ROMEO.jl).
fn percentile_95(mag: &[f64], mask: &[u8]) -> f64 {
    let mut values: Vec<f64> = mag.iter().enumerate()
        .filter(|(i, v)| mask[*i] > 0 && v.is_finite())
        .map(|(_, &v)| v)
        .collect();
    if values.is_empty() {
        return 1.0;
    }
    values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let idx = ((values.len() as f64 * 0.95) as usize).min(values.len() - 1);
    values[idx]
}

/// Calculate BestPath edge weights.
///
/// Uses second-order phase differences across all neighbor directions
/// (Abdul-Rahman, https://doi.org/10.1364/AO.46.006623).
/// Matches ROMEO.jl `calculateweights_bestpath`.
pub fn calculate_weights_bestpath(
    phase: &[f64],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
) -> Vec<u8> {
    let n_total = nx * ny * nz;

    // Compute D² for each voxel: sum of squared wrapped second-order differences
    // across all unique neighbor offsets from 26-connected directions
    let strides = [1_i64, nx as i64, (nx * ny) as i64];
    let mut neighbor_offsets: Vec<i64> = Vec::new();
    for dx in -1i64..=1 {
        for dy in -1i64..=1 {
            for dz in -1i64..=1 {
                if dx == 0 && dy == 0 && dz == 0 { continue; }
                let offset = (dx * strides[0] + dy * strides[1] + dz * strides[2]).abs();
                if !neighbor_offsets.contains(&offset) {
                    neighbor_offsets.push(offset);
                }
            }
        }
    }

    let mut d2 = vec![0.0_f64; n_total];
    for &n in &neighbor_offsets {
        let nu = n as usize;
        for i in nu..(n_total - nu) {
            let v1 = wrap_angle(phase[i - nu] - phase[i]);
            let v2 = wrap_angle(phase[i] - phase[i + nu]);
            let diff = v1 - v2;
            d2[i] += diff * diff;
        }
    }

    // R = 1 / sqrt(D²)
    let r: Vec<f64> = d2.iter().map(|&v| {
        let d = v.sqrt();
        if d > 1e-12 { 1.0 / d } else { 0.0 }
    }).collect();

    // Edge weights: R[i] + R[i+n] for each direction
    let mut weights = vec![0u8; 3 * n_total];
    for dim in 0..3_usize {
        for i in 0..nx {
            for j in 0..ny {
                for k in 0..nz {
                    let (ni, nj, nk) = match dim {
                        0 => (i + 1, j, k),
                        1 => (i, j + 1, k),
                        _ => (i, j, k + 1),
                    };
                    if ni >= nx || nj >= ny || nk >= nz { continue; }

                    let idx = idx3d(i, j, k, nx, ny);
                    let idx_n = idx3d(ni, nj, nk, nx, ny);

                    if mask[idx] == 0 || mask[idx_n] == 0 { continue; }

                    let w = r[idx] + r[idx_n];
                    // Scale: lower w = more consistent = higher weight
                    // ROMEO.jl: scale(w) = min(max(round((1 - w/10) * 255), 1), 255)
                    let scaled = ((1.0 - w / 10.0) * 255.0).round().clamp(1.0, 255.0) as u8;
                    let edge_idx = dim * n_total + idx3d(i, j, k, nx, ny);
                    weights[edge_idx] = scaled;
                }
            }
        }
    }

    weights
}

/// Simplified weight calculation for single-echo data (no phase2)
pub fn calculate_weights_single_echo(
    phase: &[f64],
    mag: &[f64],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
) -> Vec<u8> {
    calculate_weights_romeo(phase, mag, None, 1.0, 1.0, mask, nx, ny, nz)
}

/// Calculate per-voxel quality map from ROMEO edge weights
///
/// Computes ROMEO edge weights and then aggregates them per-voxel by averaging
/// the incident edge weights across all 6 neighboring directions (±x, ±y, ±z).
/// This produces a quality map where high values indicate voxels with coherent
/// phase and magnitude, suitable for thresholding into a brain mask.
///
/// Reference: MriResearchTools.jl `romeovoxelquality()` function
///
/// # Arguments
/// * `phase` - Wrapped phase data (nx * ny * nz), first echo
/// * `mag` - Magnitude data (nx * ny * nz), optional (pass empty slice if none)
/// * `phase2` - Second echo phase for gradient coherence (optional)
/// * `te1`, `te2` - Echo times for gradient coherence scaling
/// * `mask` - Binary mask (nx * ny * nz), 1 = process
/// * `nx`, `ny`, `nz` - Array dimensions
///
/// # Returns
/// Quality map of size nx * ny * nz with values in range [0, 100]
#[allow(clippy::erasing_op)]
pub fn voxel_quality_romeo(
    phase: &[f64],
    mag: &[f64],
    phase2: Option<&[f64]>,
    te1: f64,
    te2: f64,
    mask: &[u8],
    grid: &Grid,
) -> Vec<f64> {
    let (nx, ny, nz) = grid.dims;
    let n_total = nx * ny * nz;
    let weights = calculate_weights_romeo(phase, mag, phase2, te1, te2, mask, nx, ny, nz);

    let mut quality = vec![0.0_f64; n_total];

    // Each voxel's quality is an independent read of neighbouring edge
    // weights — fan out across voxels (rayon under the `parallel` feature).
    maybe_par_iter_mut!(quality).enumerate().for_each(|(idx, q)| {
        if mask[idx] == 0 {
            return;
        }
        let i = idx % nx;
        let j = (idx / nx) % ny;
        let k = idx / (nx * ny);

        let mut sum = 0.0_f64;
        let mut count = 0u32;

        // +x edge: stored at (i, j, k) in dim=0
        if i + 1 < nx && mask[idx3d(i + 1, j, k, nx, ny)] != 0 {
            sum += weights[idx] as f64;
            count += 1;
        }
        // -x edge: stored at (i-1, j, k) in dim=0
        if i > 0 && mask[idx3d(i - 1, j, k, nx, ny)] != 0 {
            sum += weights[idx3d(i - 1, j, k, nx, ny)] as f64;
            count += 1;
        }
        // +y edge: stored at (i, j, k) in dim=1
        if j + 1 < ny && mask[idx3d(i, j + 1, k, nx, ny)] != 0 {
            sum += weights[n_total + idx] as f64;
            count += 1;
        }
        // -y edge: stored at (i, j-1, k) in dim=1
        if j > 0 && mask[idx3d(i, j - 1, k, nx, ny)] != 0 {
            sum += weights[n_total + idx3d(i, j - 1, k, nx, ny)] as f64;
            count += 1;
        }
        // +z edge: stored at (i, j, k) in dim=2
        if k + 1 < nz && mask[idx3d(i, j, k + 1, nx, ny)] != 0 {
            sum += weights[2 * n_total + idx] as f64;
            count += 1;
        }
        // -z edge: stored at (i, j, k-1) in dim=2
        if k > 0 && mask[idx3d(i, j, k - 1, nx, ny)] != 0 {
            sum += weights[2 * n_total + idx3d(i, j, k - 1, nx, ny)] as f64;
            count += 1;
        }

        if count > 0 {
            // Normalize from 0-255 to 0-1, then scale to 0-100
            *q = (sum / count as f64) / 255.0 * 100.0;
        }
    });

    quality
}

// =========================================================================
// Phase unwrapping functions
// =========================================================================

/// Find a seed point for region growing: the centre of mass of the mask, or - when that
/// falls outside the mask, as it does for a disconnected or concave one - the nearest
/// voxel that is inside it.
pub fn find_seed_point(mask: &[u8], nx: usize, ny: usize, nz: usize) -> (usize, usize, usize) {
    let mut sum_i = 0usize;
    let mut sum_j = 0usize;
    let mut sum_k = 0usize;
    let mut count = 0usize;

    for k in 0..nz {
        for j in 0..ny {
            for i in 0..nx {
                let idx = idx3d(i, j, k, nx, ny);
                if mask[idx] > 0 {
                    sum_i += i;
                    sum_j += j;
                    sum_k += k;
                    count += 1;
                }
            }
        }
    }

    if count == 0 {
        return (nx / 2, ny / 2, nz / 2);
    }

    let (ci, cj, ck) = (sum_i / count, sum_j / count, sum_k / count);
    if mask[idx3d(ci, cj, ck, nx, ny)] > 0 {
        return (ci, cj, ck);
    }

    // The centre of mass of a mask that is disconnected or merely concave can fall outside
    // it - two hemispheres, a horseshoe - and the region growers treat a seed outside the
    // mask as nothing to do and return the input still wrapped. Fall back to the nearest
    // voxel that is actually in the mask. For a mask whose centre of mass is inside, which
    // is the usual case, nothing above this line changes.
    let mut best = (ci, cj, ck);
    let mut best_d2 = f64::INFINITY;
    for k in 0..nz {
        for j in 0..ny {
            for i in 0..nx {
                if mask[idx3d(i, j, k, nx, ny)] == 0 {
                    continue;
                }
                let d2 = (i as f64 - ci as f64).powi(2)
                    + (j as f64 - cj as f64).powi(2)
                    + (k as f64 - ck as f64).powi(2);
                if d2 < best_d2 {
                    best_d2 = d2;
                    best = (i, j, k);
                }
            }
        }
    }
    best
}

/// Compute weights from params (dispatches BestPath vs ROMEO).
fn compute_weights_from_params(
    params: &RomeoParams,
    phase: &[f64],
    mag: &[f64],
    phase2: Option<&[f64]>,
    te1: f64, te2: f64,
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
) -> Vec<u8> {
    if params.bestpath {
        calculate_weights_bestpath(phase, mask, nx, ny, nz)
    } else {
        calculate_weights_romeo_with_flags(
            phase, mag, phase2, te1, te2, mask, nx, ny, nz,
            params.weight_flags(),
        )
    }
}

/// Unwrap a single 3D phase volume using ROMEO.
///
/// Computes ROMEO edge weights and performs region-growing phase unwrapping.
///
/// # Arguments
/// * `phase` - Wrapped phase data (nx * ny * nz)
/// * `mag` - Magnitude data (pass empty slice if none)
/// * `phase2` - Optional second echo phase for gradient coherence weights
/// * `te1`, `te2` - Echo times for phase gradient coherence scaling
/// * `mask` - Binary mask (1 = process)
/// * `params` - ROMEO parameters (weight_type used for weight calculation)
/// * `nx`, `ny`, `nz` - Array dimensions
///
/// # Returns
/// Unwrapped phase (same size as input)
pub fn unwrap_romeo(
    phase: &[f64],
    mag: &[f64],
    phase2: Option<&[f64]>,
    te1: f64,
    te2: f64,
    mask: &[u8],
    params: &RomeoParams,
    grid: &Grid,
) -> Vec<f64> {
    let (nx, ny, nz) = grid.dims;
    let weights = compute_weights_from_params(
        params, phase, mag, phase2, te1, te2,
        mask, nx, ny, nz,
    );

    let mut unwrapped = phase.to_vec();
    let n_total = nx * ny * nz;

    // ROMEO.jl's region growing (`grow_region_unwrap!`), seeded as ROMEO.jl seeds (see the
    // module docs, *Seed*), for one seed as for several.
    let mut visited = vec![0u8; n_total];
    let num_regions = grow_region_unwrap_full(
        &mut unwrapped, &weights, mask, &mut visited,
        nx, ny, nz, params.wrap_addition, params.max_seeds,
        phase2, if phase2.is_some() { Some((te1, te2)) } else { None },
    );

    finish_3d_unwrap(&mut unwrapped, &mut visited, num_regions, &weights, mask, params, nx, ny, nz);
    unwrapped
}

/// What ROMEO.jl's 3D `unwrap!` does after the growth (`src/unwrapping.jl`, `src/algorithm.jl`):
/// region merging and correction, then, with `correct_global`, `wrapped .-= 2π·median(round(φ[mask]
/// / 2π))`. Its multi-echo modes call that same 3D `unwrap!` for the template echo and for every
/// echo in individual mode, so all three paths finish here.
#[allow(clippy::too_many_arguments)]
fn finish_3d_unwrap(
    unwrapped: &mut [f64],
    visited: &mut [u8],
    num_regions: u8,
    weights: &[u8],
    mask: &[u8],
    params: &RomeoParams,
    nx: usize, ny: usize, nz: usize,
) {
    if params.merge_regions && num_regions > 1 {
        let remaining = merge_regions_post(
            unwrapped, visited, num_regions, weights, nx, ny, nz,
        );
        if params.correct_regions {
            // correct_regions on remaining regions
            for &r in &remaining {
                correct_regions(unwrapped, visited, r);
            }
        }
    } else if params.correct_regions && num_regions > 0 {
        correct_regions(unwrapped, visited, num_regions);
    }

    // correctglobal for 3D: subtract median n·2π offset
    if params.correct_global {
        correct_global_offset(unwrapped, mask);
    }
}

/// Unwrap 4D multi-echo phase data using ROMEO.
///
/// Supports two modes matching ROMEO.jl:
///
/// **Template-based** (default, `individual=false`):
/// Spatially unwraps one template echo, then temporally unwraps all others
/// by scaling with TE ratios. Optionally re-unwraps uncertain voxels spatially.
///
/// **Individual** (`individual=true`):
/// Spatially unwraps each echo independently, then optionally corrects global
/// 2π offsets between echoes using median wrap counting.
///
/// # Arguments
/// * `phases` - Wrapped phase for each echo \[n_echoes\]\[nx*ny*nz\]
/// * `mags` - Magnitude for each echo (pass empty `&[]` if none)
/// * `tes` - Echo times (units don't matter, only ratios are used)
/// * `mask` - Binary mask (nx * ny * nz)
/// * `params` - ROMEO parameters
/// * `nx`, `ny`, `nz` - Array dimensions
///
/// # Returns
/// Unwrapped phase for each echo
pub fn unwrap_romeo_multi_echo<P: AsRef<[f64]> + Sync, M: AsRef<[f64]> + Sync>(
    phases: &[P],
    mags: &[M],
    tes: &[f64],
    mask: &[u8],
    params: &RomeoParams,
    grid: &Grid,
) -> Vec<Vec<f64>> {
    let (nx, ny, nz) = grid.dims;
    let n_echoes = phases.len();
    assert!(n_echoes > 0, "phases must have at least one echo");
    assert_eq!(n_echoes, tes.len(), "phases and tes must have same length");

    if n_echoes == 1 {
        let mag = if mags.is_empty() { &[] as &[f64] } else { mags[0].as_ref() };
        return vec![unwrap_romeo(
            phases[0].as_ref(), mag, None, 0.0, 0.0,
            mask, params, grid,
        )];
    }

    if params.individual {
        unwrap_individual(phases, mags, tes, mask, params, nx, ny, nz)
    } else {
        unwrap_template(phases, mags, tes, mask, params, nx, ny, nz)
    }
}

// =========================================================================
// Template-based multi-echo unwrapping
// =========================================================================

/// Template-based multi-echo unwrapping.
///
/// 1. Spatially unwrap the template echo using ROMEO weights
/// 2. Temporally unwrap all other echoes outward from the template
/// 3. Optionally re-unwrap uncertain voxels spatially
fn unwrap_template<P: AsRef<[f64]>, M: AsRef<[f64]>>(
    phases: &[P],
    mags: &[M],
    tes: &[f64],
    mask: &[u8],
    params: &RomeoParams,
    nx: usize, ny: usize, nz: usize,
) -> Vec<Vec<f64>> {
    let n_echoes = phases.len();
    let n_total = nx * ny * nz;
    let template = params.template.min(n_echoes - 1);

    // Select phase2 reference (matching ROMEO.jl p2ref default)
    let p2ref = if template == 0 { 1 } else { template - 1 };

    // Calculate weights using template echo + p2ref
    let template_mag = if mags.is_empty() { &[] as &[f64] } else { mags[template].as_ref() };
    let weights = compute_weights_from_params(
        params,
        phases[template].as_ref(),
        template_mag,
        Some(phases[p2ref].as_ref()),
        tes[template], tes[p2ref],
        mask, nx, ny, nz,
    );

    // Spatially unwrap template echo
    let mut result: Vec<Vec<f64>> = phases.iter().map(|p| p.as_ref().to_vec()).collect();
    // ROMEO.jl: `unwrap!(view(wrapped, :, :, :, template); args..., weights)` with phase2 and
    // TEs, so the seed's 2π count is set from phase2 (`seedcorrection!`).
    let mut visited = vec![0u8; n_total];
    let num_regions = grow_region_unwrap_full(
        &mut result[template], &weights, mask, &mut visited,
        nx, ny, nz, params.wrap_addition, params.max_seeds,
        Some(phases[p2ref].as_ref()), Some((tes[template], tes[p2ref])),
    );
    // The rest of that 3D `unwrap!`: region handling and, with `correct_global`, the template's
    // global turn (until 2026-11 neither ran on the template echo here).
    finish_3d_unwrap(&mut result[template], &mut visited, num_regions, &weights, mask, params, nx, ny, nz);

    // Temporally unwrap other echoes outward from template
    // Order: template-1, template-2, ..., 0, template+1, template+2, ..., n-1
    // (matching ROMEO.jl: [(template-1):-1:1; (template+1):length(TEs)])
    let echo_order: Vec<usize> = (0..template).rev()
        .chain((template + 1)..n_echoes)
        .collect();

    for ieco in echo_order {
        let iref = if ieco < template { ieco + 1 } else { ieco - 1 };
        let te_ratio = tes[ieco] / tes[iref];

        // Temporal unwrap each voxel: align to TE-scaled reference
        for i in 0..n_total {
            if mask[i] > 0 {
                let ref_value = result[iref][i] * te_ratio;
                result[ieco][i] = unwrap_voxel(result[ieco][i], ref_value);
            }
        }

        // Fallback: re-unwrap uncertain voxels spatially
        if params.temporal_uncertain_unwrapping > 0.0 {
            // Build scaled reference for quality assessment
            let ref_scaled: Vec<f64> = (0..n_total).map(|i| {
                if mask[i] > 0 { result[iref][i] * te_ratio } else { 0.0 }
            }).collect();

            temporal_uncertain_rewrap(
                &mut result[ieco],
                &ref_scaled,
                &weights,
                mask,
                params.temporal_uncertain_unwrapping,
                nx, ny, nz,
            );
        }
    }

    result
}

// =========================================================================
// Individual multi-echo unwrapping
// =========================================================================

/// Individual multi-echo unwrapping: ROMEO.jl's `unwrap_individual!` (`src/unwrapping.jl`,
/// v1.6.0).
///
/// Each echo is spatially unwrapped on its own, with its own magnitude and the previous echo (echo
/// 2 for echo 1) as `phase2` for the gradient-coherence weight and the seed's turn. With
/// `correct_global`, as in ROMEO.jl, two corrections follow:
/// 1. per echo, inside the 3D `unwrap!` that `unwrap_individual!` calls with all its keyword
///    arguments: `wrapped .-= 2π·median(round(φ[mask] / 2π))` (finite voxels only), so each
///    echo's median turn is 0;
/// 2. then across echoes, `correct_multi_echo_wraps!` ([`correct_multi_echo_wraps`]).
///
/// Step 1 was missing here until 2026-11; on the data where it was measured (Q0748 audit,
/// workstream C) every echo's median turn was already 0, so it changed nothing there, but in
/// general it sets echo 1's global turn and with it every echo's.
///
/// **`phase2` and threads.** ROMEO.jl runs the echoes in a `Threads.@threads` loop and reads
/// `phase2 = wrapped[:,:,:,e2]` while the loop unwraps `wrapped` in place, so whether echo `i`
/// sees echo `i − 1` wrapped or already unwrapped depends on the thread schedule (≈1 % of voxels of
/// a 16-echo 7 T scan's last echo varied between runs; Q0748 audit). Here the echoes run in order,
/// which is ROMEO.jl with one thread, the deterministic reading of its code: echo `i ≥ 2` sees
/// echo `i − 1` unwrapped (and globally corrected), echo 1 sees echo 2 wrapped. It matters beyond
/// the seed's turn because ROMEO.jl's `γ` wraps only once (`src/utility.jl`), so the
/// gradient-coherence weight of an unwrapped `phase2` differs where neighbours differ by more than
/// 3π (on `tests/romeo_validity_romeojl_parity.rs`'s phantom, 46 voxels of echo 3 changed their
/// turns between the two readings). Until 2026-11 QSM.rs always read the wrapped echo (and ran the
/// echoes in parallel), which is what ROMEO.jl reads when the echoes start at the same time.
///
/// **Medians.** ROMEO.jl's `median` of an even number of turn counts averages the two middle
/// ones, which can give a half turn (a shift by π); here the upper middle count is taken, so the
/// correction is always a whole number of turns. The two agree whenever the middle counts are
/// equal, as they are on any real mask.
fn unwrap_individual<P: AsRef<[f64]> + Sync, M: AsRef<[f64]> + Sync>(
    phases: &[P],
    mags: &[M],
    tes: &[f64],
    mask: &[u8],
    params: &RomeoParams,
    nx: usize, ny: usize, nz: usize,
) -> Vec<Vec<f64>> {
    let n_echoes = phases.len();

    // ROMEO.jl's loop, in echo order (as with one Julia thread): `unwrap!` works on `wrapped` in
    // place, so echo i ≥ 2's `phase2` is echo i − 1 as already unwrapped (and, with
    // `correct_global`, already corrected); echo 1's is echo 2, still wrapped.
    let mut result: Vec<Vec<f64>> = phases.iter().map(|p| p.as_ref().to_vec()).collect();
    for i in 0..n_echoes {
        let e2 = if i == 0 { 1 } else { i - 1 };
        let phase2 = result[e2].clone();
        let mag = if mags.is_empty() { &[] as &[f64] } else { mags[i].as_ref() };
        let weights = compute_weights_from_params(
            params, &result[i], mag, Some(&phase2), tes[i], tes[e2], mask, nx, ny, nz,
        );
        let mut visited = vec![0u8; nx * ny * nz];
        let num_regions = grow_region_unwrap_full(
            &mut result[i], &weights, mask, &mut visited,
            nx, ny, nz, params.wrap_addition, params.max_seeds,
            Some(&phase2), Some((tes[i], tes[e2])),
        );
        // ROMEO.jl's 3D `unwrap!` inside `unwrap_individual!`: regions, then `correctglobal`
        finish_3d_unwrap(&mut result[i], &mut visited, num_regions, &weights, mask, params, nx, ny, nz);
    }

    if params.correct_global {
        correct_multi_echo_wraps(&mut result, tes, mask);
    }

    result
}

// =========================================================================
// Helper functions
// =========================================================================

/// Temporal unwrap: remove 2π wraps by comparing to a reference value.
///
/// Matches ROMEO.jl: `unwrapvoxel(new, old) = new - 2π * round((new - old) / 2π)`
#[inline]
fn unwrap_voxel(new: f64, old: f64) -> f64 {
    new - TWO_PI * ((new - old) / TWO_PI).round()
}

/// Correct global 2π offsets between echoes using median wrap counting.
///
/// For each successive echo pair, computes the median number of 2π wraps
/// between the TE-scaled reference and the current echo, then corrects.
///
/// Matches ROMEO.jl `correct_multi_echo_wraps!`.
pub fn correct_multi_echo_wraps(
    phases: &mut [Vec<f64>],
    tes: &[f64],
    mask: &[u8],
) {
    let n_total = phases[0].len();

    for ieco in 1..phases.len() {
        let iref = ieco - 1;
        let te_ratio = tes[ieco] / tes[iref];

        // Collect wrap counts for all masked voxels
        let mut wrap_counts: Vec<f64> = Vec::new();
        for i in 0..n_total {
            if mask[i] > 0 {
                let expected = phases[iref][i] * te_ratio;
                let nwraps = ((expected - phases[ieco][i]) / TWO_PI).round();
                if nwraps.is_finite() {
                    wrap_counts.push(nwraps);
                }
            }
        }

        if wrap_counts.is_empty() {
            continue;
        }

        // Median wrap count
        wrap_counts.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let median_wraps = wrap_counts[wrap_counts.len() / 2];

        if median_wraps.abs() > 0.1 {
            let correction = TWO_PI * median_wraps;
            for i in 0..n_total {
                if mask[i] > 0 {
                    phases[ieco][i] += correction;
                }
            }
        }
    }
}

/// Correct global 2π offset for a single 3D volume.
///
/// Subtracts `2π * median(round(phase / 2π))` from the entire volume,
/// bringing the median phase closest to 0. Matches ROMEO.jl `correctglobal`.
fn correct_global_offset(phase: &mut [f64], mask: &[u8]) {
    let mut wraps: Vec<f64> = phase.iter().enumerate()
        .filter(|(i, v)| mask[*i] > 0 && v.is_finite())
        .map(|(_, &v)| (v / TWO_PI).round())
        .collect();

    if wraps.is_empty() {
        return;
    }

    wraps.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let median_wraps = wraps[wraps.len() / 2];

    if median_wraps.abs() > 0.1 {
        let correction = TWO_PI * median_wraps;
        for i in 0..phase.len() {
            if mask[i] > 0 {
                phase[i] -= correction;
            }
        }
    }
}

/// Correct each region's median to nearest 0 by subtracting n·2π.
///
/// Matches ROMEO.jl `correct_regions!`.
fn correct_regions(phase: &mut [f64], visited: &[u8], num_regions: u8) {
    for region in 1..=num_regions {
        let mut wraps: Vec<f64> = Vec::new();
        for i in 0..phase.len() {
            if visited[i] == region && phase[i].is_finite() {
                wraps.push((phase[i] / TWO_PI).round());
            }
        }
        if wraps.is_empty() { continue; }
        wraps.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let median_wraps = wraps[wraps.len() / 2];
        if median_wraps.abs() > 0.1 {
            let correction = TWO_PI * median_wraps;
            for i in 0..phase.len() {
                if visited[i] == region {
                    phase[i] -= correction;
                }
            }
        }
    }
}

/// Merge adjacent regions by computing inter-region phase offsets.
///
/// Calculates weighted phase offset between adjacent region pairs,
/// then iteratively merges regions starting from the largest.
/// Matches ROMEO.jl `merge_regions!`.
fn merge_regions_post(
    phase: &mut [f64],
    visited: &mut [u8],
    num_regions: u8,
    weights: &[u8],
    nx: usize, ny: usize, nz: usize,
) -> Vec<u8> {
    let n_total = nx * ny * nz;
    let nr = num_regions as usize;
    if nr <= 1 { return vec![1]; }

    // Count region sizes
    let mut region_size = vec![0usize; nr + 1]; // 1-indexed
    for &v in visited.iter() {
        if v > 0 && (v as usize) <= nr {
            region_size[v as usize] += 1;
        }
    }

    // Compute weighted offset between adjacent regions
    let mut offsets = vec![0.0_f64; (nr + 1) * (nr + 1)];
    let mut offset_counts = vec![0i64; (nr + 1) * (nr + 1)];
    let flat = |r1: usize, r2: usize| r1 * (nr + 1) + r2;

    for dim in 0..3_usize {
        for i in 0..nx {
            for j in 0..ny {
                for k in 0..nz {
                    let (ni, nj, nk) = match dim {
                        0 => (i + 1, j, k),
                        1 => (i, j + 1, k),
                        _ => (i, j, k + 1),
                    };
                    if ni >= nx || nj >= ny || nk >= nz { continue; }
                    let idx = idx3d(i, j, k, nx, ny);
                    let idx_n = idx3d(ni, nj, nk, nx, ny);
                    let ri = visited[idx] as usize;
                    let rj = visited[idx_n] as usize;
                    if ri == 0 || rj == 0 || ri == rj { continue; }

                    let edge_idx = dim * n_total + idx3d(i, j, k, nx, ny);
                    let w = 255u16.saturating_sub(weights[edge_idx] as u16);
                    let w = if w == 255 { 0 } else { w as i64 };

                    offsets[flat(ri, rj)] += (phase[idx] - phase[idx_n]) * w as f64;
                    offset_counts[flat(ri, rj)] += w;
                }
            }
        }
    }

    // Symmetrize
    for i in 1..=nr {
        for j in i..=nr {
            offset_counts[flat(i, j)] += offset_counts[flat(j, i)];
            offset_counts[flat(j, i)] = offset_counts[flat(i, j)];
            offsets[flat(i, j)] -= offsets[flat(j, i)];
            offsets[flat(j, i)] = -offsets[flat(i, j)];
        }
    }

    // Iteratively merge: process largest uncorrected region first
    let mut corrected = vec![false; nr + 1];
    let mut remaining = Vec::new();

    while corrected[1..=nr].iter().any(|&c| !c) {
        // Find largest uncorrected region
        let mut best_region = 0;
        let mut best_size = 0;
        for r in 1..=nr {
            if !corrected[r] && region_size[r] > best_size {
                best_size = region_size[r];
                best_region = r;
            }
        }
        if best_region == 0 { break; }
        corrected[best_region] = true;
        remaining.push(best_region as u8);

        // Find regions to merge: corrected[i] && !corrected[j] && offset_counts > 0
        // Sort by offset_counts (highest first = best connection)
        let mut merge_pairs: Vec<(usize, usize, i64)> = Vec::new();
        for i in 1..=nr {
            for j in 1..=nr {
                if corrected[i] && !corrected[j] && offset_counts[flat(i, j)] > 0 {
                    merge_pairs.push((i, j, offset_counts[flat(i, j)]));
                }
            }
        }
        merge_pairs.sort_by_key(|a| std::cmp::Reverse(a.2));

        for &(ri, rj, count) in &merge_pairs {
            if corrected[rj] { continue; }
            let offset = (offsets[flat(ri, rj)] / count as f64 / TWO_PI).round();
            if offset != 0.0 {
                let correction = offset * TWO_PI;
                for v in 0..n_total {
                    if visited[v] == rj as u8 {
                        phase[v] += correction;
                        visited[v] = ri as u8;
                    }
                }
            }
            corrected[rj] = true;
            offset_counts[flat(ri, rj)] = -1;
            offset_counts[flat(rj, ri)] = -1;
        }
    }

    remaining
}

/// Re-unwrap voxels that are uncertain after temporal unwrapping.
///
/// Computes a quality metric comparing the unwrapped phase to the TE-scaled
/// reference. Voxels with low quality (likely wrap errors) are re-unwrapped
/// spatially using the certain voxels as seeds.
///
/// Matches ROMEO.jl `temporal_uncertain_unwrapping!`.
fn temporal_uncertain_rewrap(
    phase: &mut [f64],
    ref_scaled: &[f64],
    weights: &[u8],
    mask: &[u8],
    threshold: f64,
    nx: usize, ny: usize, nz: usize,
) {
    let n_total = nx * ny * nz;

    // Compute quality: compare phase/2 with ref_scaled/2 using voxel quality
    // (halving increases sensitivity to single-wrap errors)
    // Matches ROMEO.jl: unwrapped_quality(phase, refphase) = voxelquality(phase/2; phase2=refphase/2, TEs=[1,1])
    let half_phase: Vec<f64> = phase.iter().map(|&v| v * 0.5).collect();
    let half_ref: Vec<f64> = ref_scaled.iter().map(|&v| v * 0.5).collect();

    let grid = crate::Grid::new(nx, ny, nz, 1.0, 1.0, 1.0);
    let quality = voxel_quality_romeo(
        &half_phase, &[], Some(&half_ref), 1.0, 1.0,
        mask, &grid,
    );

    // Build visited mask:
    // quality > threshold*100 → certain (visited=2, acts as seed)
    // quality <= threshold*100 → uncertain (visited=1, needs re-unwrapping)
    // outside mask → 0
    let threshold_scaled = threshold * 100.0; // voxel_quality_romeo returns 0-100
    let mut visited = vec![0u8; n_total];
    let mut any_uncertain = false;
    let mut any_certain = false;

    for i in 0..n_total {
        if mask[i] == 0 {
            visited[i] = 0;
        } else if quality[i] > threshold_scaled {
            visited[i] = 2;
            any_certain = true;
        } else {
            visited[i] = 1;
            any_uncertain = true;
        }
    }

    if !any_uncertain || !any_certain {
        return;
    }

    // Re-unwrap uncertain voxels using spatial growing from the certain boundary
    grow_region_unwrap_from_visited(phase, weights, &mut visited, nx, ny, nz);
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Grid;

    fn grid(n: usize) -> Grid { Grid::new(n, n, n, 1.0, 1.0, 1.0) }

    #[test]
    fn find_seed_point_returns_a_voxel_inside_the_mask() {
        // The growers treat a seed outside the mask as nothing to do and hand back the input
        // still wrapped, so a centre of mass in a hollow is a silent total failure. Two
        // ordinary brain-mask shapes put it there.
        let (nx, ny, nz) = (16usize, 16usize, 4usize);
        let at = |m: &mut Vec<u8>, i: usize, j: usize, k: usize| m[i + j * nx + k * nx * ny] = 1;

        // two separated blobs: the centre of mass sits in the gap between them
        let mut split = vec![0u8; nx * ny * nz];
        for k in 0..nz {
            for j in 5..11 {
                for i in 1..5 { at(&mut split, i, j, k); }
                for i in 11..15 { at(&mut split, i, j, k); }
            }
        }
        let (i, j, k) = find_seed_point(&split, nx, ny, nz);
        assert_eq!(split[idx3d(i, j, k, nx, ny)], 1, "seed ({i},{j},{k}) is outside the mask");

        // a horseshoe: one connected component, centre of mass in the hollow
        let mut horseshoe = vec![0u8; nx * ny * nz];
        for k in 0..nz {
            for j in 2..14 {
                for i in 2..5 { at(&mut horseshoe, i, j, k); }
                for i in 11..14 { at(&mut horseshoe, i, j, k); }
            }
            for i in 2..14 { at(&mut horseshoe, i, 2, k); at(&mut horseshoe, i, 3, k); }
        }
        let (i, j, k) = find_seed_point(&horseshoe, nx, ny, nz);
        assert_eq!(horseshoe[idx3d(i, j, k, nx, ny)], 1, "seed ({i},{j},{k}) is outside the mask");

        // a convex mask is unaffected: the seed is exactly the centre of mass
        let mut block = vec![0u8; nx * ny * nz];
        for k in 1..3 {
            for j in 4..12 {
                for i in 4..12 { at(&mut block, i, j, k); }
            }
        }
        assert_eq!(find_seed_point(&block, nx, ny, nz), (7, 7, 1));
    }

    #[test]
    fn romeo_unwraps_a_mask_whose_centre_of_mass_is_outside_it() {
        // The behavioural consequence: before the fallback this returned the input untouched.
        let (nx, ny, nz) = (16usize, 16usize, 2usize);
        let mut mask = vec![0u8; nx * ny * nz];
        for k in 0..nz {
            for j in 5..11 {
                for i in 1..5 { mask[i + j * nx + k * nx * ny] = 1; }
                for i in 11..15 { mask[i + j * nx + k * nx * ny] = 1; }
            }
        }
        // a ramp steep enough to wrap across each blob but not between samples
        let mut truth = vec![0.0; nx * ny * nz];
        for k in 0..nz {
            for j in 0..ny {
                for i in 0..nx {
                    truth[i + j * nx + k * nx * ny] = 1.2 * i as f64;
                }
            }
        }
        let wrapped: Vec<f64> = truth.iter().map(|&v| wrap_angle(v)).collect();
        assert!(
            wrapped.iter().zip(&truth).any(|(w, t)| (w - t).abs() > 1.0),
            "the ramp does not wrap, so this proves nothing"
        );

        let grid = Grid::new(nx, ny, nz, 1.0, 1.0, 1.0);
        let out = unwrap_romeo(
            &wrapped, &[], None, 0.0, 0.0, &mask, &RomeoParams::default(), &grid,
        );
        let changed = out.iter().zip(&wrapped)
            .enumerate()
            .filter(|(i, (o, w))| mask[*i] != 0 && (*o - *w).abs() > 1e-9)
            .count();
        assert!(changed > 0, "nothing in the mask was unwrapped - the seed landed outside it");
    }

    #[test]
    fn test_wrap_angle() {
        assert!((wrap_angle(0.0) - 0.0).abs() < 1e-10);
        assert!((wrap_angle(PI) - PI).abs() < 1e-10);
        assert!((wrap_angle(-PI) - (-PI)).abs() < 1e-10);
        assert!((wrap_angle(TWO_PI) - 0.0).abs() < 1e-10);
        assert!((wrap_angle(3.0 * PI) - PI).abs() < 1e-10);
        assert!((wrap_angle(-3.0 * PI) - (-PI)).abs() < 1e-10);
    }

    #[test]
    fn test_weights_constant_phase() {
        // Constant phase should give high weights (phase coherence = 1)
        let n = 4;
        let phase = vec![0.0; n * n * n];
        let mag = vec![1.0; n * n * n];
        let mask = vec![1u8; n * n * n];

        let weights = calculate_weights_single_echo(&phase, &mag, &mask, n, n, n);

        // All interior weights should be 255 (constant phase, uniform magnitude)
        let mut high_weight_count = 0;
        for &w in weights.iter() {
            if w == 255 {
                high_weight_count += 1;
            }
        }
        assert!(high_weight_count > 0, "Should have some high weights for constant phase");
    }

    #[test]
    fn test_weights_wrapped_jump() {
        // Phase with 2π jump should give low weights at jump location
        let n = 4;
        let mut phase = vec![0.0; n * n * n];

        // Create a 2π jump in x direction at i=2
        for i in 2..n {
            for j in 0..n {
                for k in 0..n {
                    phase[idx3d(i, j, k, n, n)] = TWO_PI;
                }
            }
        }

        let mask = vec![1u8; n * n * n];
        let _weights = calculate_weights_single_echo(&phase, &[], &mask, n, n, n);

        // Weight at x=1 to x=2 edge should be low (wrapped difference = 0, but that's ok)
        // Actually, for a 2π jump, the wrapped difference is 0, so coherence is 1
        // This is correct - ROMEO uses wrapped differences, not raw differences
    }

    #[test]
    fn test_weights_mask() {
        // Weights should be 0 where mask is 0
        let n = 4;
        let phase = vec![0.5; n * n * n];
        let mut mask = vec![1u8; n * n * n];

        // Set some voxels outside mask
        mask[0] = 0;
        mask[1] = 0;

        let weights = calculate_weights_single_echo(&phase, &[], &mask, n, n, n);

        // Edges connected to masked-out voxels should be 0
        // Weight at edge (0,0,0)-(1,0,0) should be 0 since idx 0 is masked out
        assert_eq!(weights[0], 0);  // x-direction edge at (0,0,0)
    }

    #[test]
    fn test_voxel_quality_constant_phase() {
        // Constant phase + uniform magnitude → all quality values should be 100
        let n = 4;
        let phase = vec![0.0; n * n * n];
        let mag = vec![1.0; n * n * n];
        let mask = vec![1u8; n * n * n];

        let quality = voxel_quality_romeo(&phase, &mag, None, 1.0, 1.0, &mask, &grid(n));

        assert_eq!(quality.len(), n * n * n);

        // Interior voxels should have high quality (may be < 100 due to
        // phaselinearity boundary fallback at edges of small test cube)
        let interior_q = quality[idx3d(1, 1, 1, n, n)];
        assert!(interior_q > 90.0,
                "Interior voxel quality should be >90, got {}", interior_q);
    }

    #[test]
    fn test_voxel_quality_masked() {
        // Masked-out voxels should have quality = 0
        let n = 4;
        let phase = vec![0.0; n * n * n];
        let mut mask = vec![1u8; n * n * n];
        mask[idx3d(1, 1, 1, n, n)] = 0;

        let quality = voxel_quality_romeo(&phase, &[], None, 1.0, 1.0, &mask, &grid(n));

        assert_eq!(quality[idx3d(1, 1, 1, n, n)], 0.0,
                   "Masked-out voxel should have quality 0");
    }

    #[test]
    fn test_voxel_quality_range() {
        // Quality values should be in range [0, 100]
        let n = 6;
        let phase: Vec<f64> = (0..n * n * n).map(|i| (i as f64) * 0.7).collect();
        let mag: Vec<f64> = (0..n * n * n).map(|i| (i as f64) / (n * n * n) as f64).collect();
        let mask = vec![1u8; n * n * n];

        let quality = voxel_quality_romeo(&phase, &mag, None, 1.0, 1.0, &mask, &grid(n));

        for &q in quality.iter() {
            assert!((0.0..=100.0).contains(&q),
                    "Quality should be in [0, 100], got {}", q);
        }
    }

    // =========================================================================
    // find_seed_point
    // =========================================================================

    #[test]
    fn test_find_seed_point_full_mask() {
        let (nx, ny, nz) = (8, 8, 8);
        let mask = vec![1u8; nx * ny * nz];
        let (si, sj, sk) = find_seed_point(&mask, nx, ny, nz);
        assert_eq!(si, 3);
        assert_eq!(sj, 3);
        assert_eq!(sk, 3);
    }

    #[test]
    fn test_find_seed_point_empty_mask() {
        let (nx, ny, nz) = (8, 8, 8);
        let mask = vec![0u8; nx * ny * nz];
        let (si, sj, sk) = find_seed_point(&mask, nx, ny, nz);
        assert_eq!(si, 4);
        assert_eq!(sj, 4);
        assert_eq!(sk, 4);
    }

    // =========================================================================
    // unwrap_romeo (3D)
    // =========================================================================

    #[test]
    fn test_unwrap_romeo_3d_smooth() {
        let n = 8;
        let total = n * n * n;
        let phase: Vec<f64> = (0..total).map(|i| {
            (i % n) as f64 / n as f64 * 0.5
        }).collect();
        let mag = vec![1.0; total];
        let mask = vec![1u8; total];

        let unwrapped = unwrap_romeo(
            &phase, &mag, None, 0.0, 0.0,
            &mask, &RomeoParams::default(), &grid(n),
        );

        assert_eq!(unwrapped.len(), total);
        for &v in &unwrapped {
            assert!(v.is_finite());
        }
    }

    #[test]
    fn test_unwrap_romeo_3d_with_wrap() {
        // Phase with a 2π wrap that should be removed
        let n = 8;
        let total = n * n * n;
        let mut phase = vec![0.0; total];

        // Create a smooth ramp 0..0.7 in x, then add 2π at x>=4
        for i in 0..n {
            for j in 0..n {
                for k in 0..n {
                    let idx = idx3d(i, j, k, n, n);
                    let val = 0.1 * i as f64;
                    phase[idx] = if i >= 4 { val + TWO_PI } else { val };
                }
            }
        }

        let mask = vec![1u8; total];
        let unwrapped = unwrap_romeo(
            &phase, &[], None, 0.0, 0.0,
            &mask, &RomeoParams::default(), &grid(n),
        );

        // After unwrapping, the 2π jump should be removed
        // Check that adjacent voxels don't differ by more than π
        for j in 0..n {
            for k in 0..n {
                for i in 1..n {
                    let curr = unwrapped[idx3d(i, j, k, n, n)];
                    let prev = unwrapped[idx3d(i - 1, j, k, n, n)];
                    let diff = (curr - prev).abs();
                    assert!(diff < PI + 0.5,
                        "Adjacent voxels differ by {:.2} at x={}", diff, i);
                }
            }
        }
    }

    // =========================================================================
    // unwrap_romeo_multi_echo (4D template-based)
    // =========================================================================

    #[test]
    fn test_multi_echo_template_output_sizes() {
        let n = 8;
        let total = n * n * n;
        let tes = [5.0, 10.0, 15.0];
        let slope = 0.05;

        let phases: Vec<Vec<f64>> = tes.iter().map(|&te| {
            (0..total).map(|i| {
                let x = (i % n) as f64;
                wrap_angle(slope * x * te)
            }).collect()
        }).collect();
        let mags: Vec<Vec<f64>> = tes.iter().map(|_| vec![1.0; total]).collect();
        let mask = vec![1u8; total];

        let result = unwrap_romeo_multi_echo(
            &phases, &mags, &tes, &mask,
            &RomeoParams::default(), &grid(n),
        );

        assert_eq!(result.len(), 3);
        for echo in &result {
            assert_eq!(echo.len(), total);
            for &v in echo {
                assert!(v.is_finite());
            }
        }
    }

    #[test]
    fn test_multi_echo_template_temporal_consistency() {
        // After template-based unwrapping, echoes should be temporally consistent:
        // phase[e] ≈ phase[0] * TE[e] / TE[0]
        let n = 8;
        let total = n * n * n;
        let tes = [5.0, 10.0, 15.0];
        let slope = 0.03; // small slope so no wrapping needed

        let phases: Vec<Vec<f64>> = tes.iter().map(|&te| {
            (0..total).map(|i| {
                let x = (i % n) as f64;
                slope * x * te // no wrapping needed
            }).collect()
        }).collect();
        let mags: Vec<Vec<f64>> = tes.iter().map(|_| vec![1.0; total]).collect();
        let mask = vec![1u8; total];

        let result = unwrap_romeo_multi_echo(
            &phases, &mags, &tes, &mask,
            &RomeoParams::default(), &grid(n),
        );

        // Check TE-scaling consistency at interior voxels
        for i in 0..total {
            if mask[i] > 0 && result[0][i].abs() > 1e-10 {
                let ratio_12 = result[1][i] / result[0][i];
                let expected_ratio = tes[1] / tes[0];
                assert!((ratio_12 - expected_ratio).abs() < 0.5,
                    "Echo 1/0 ratio {:.3} vs expected {:.3} at voxel {}",
                    ratio_12, expected_ratio, i);
            }
        }
    }

    #[test]
    fn test_multi_echo_single_echo_fallback() {
        let n = 4;
        let total = n * n * n;
        let phase = vec![0.1; total];
        let mask = vec![1u8; total];

        let result = unwrap_romeo_multi_echo(
            &[&phase[..]], &[] as &[&[f64]], &[5.0], &mask,
            &RomeoParams::default(), &grid(n),
        );

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].len(), total);
    }

    // =========================================================================
    // unwrap_romeo_multi_echo (individual mode)
    // =========================================================================

    #[test]
    fn test_multi_echo_individual_mode() {
        let n = 8;
        let total = n * n * n;
        let tes = [5.0, 10.0, 15.0];
        let slope = 0.03;

        let phases: Vec<Vec<f64>> = tes.iter().map(|&te| {
            (0..total).map(|i| {
                let x = (i % n) as f64;
                slope * x * te
            }).collect()
        }).collect();
        let mags: Vec<Vec<f64>> = tes.iter().map(|_| vec![1.0; total]).collect();
        let mask = vec![1u8; total];

        let params = RomeoParams {
            individual: true,
            correct_global: true,
            ..Default::default()
        };

        let result = unwrap_romeo_multi_echo(
            &phases, &mags, &tes, &mask,
            &params, &grid(n),
        );

        assert_eq!(result.len(), 3);
        for echo in &result {
            assert_eq!(echo.len(), total);
            for &v in echo {
                assert!(v.is_finite());
            }
        }
    }

    // =========================================================================
    // correct_multi_echo_wraps
    // =========================================================================

    #[test]
    fn test_correct_multi_echo_wraps_no_correction_needed() {
        let n = 64;
        let tes = [5.0, 10.0, 15.0];
        let mask = vec![1u8; n];

        // Perfectly scaled echoes: no correction needed
        let mut phases: Vec<Vec<f64>> = tes.iter().map(|&te| {
            vec![0.1 * te; n]
        }).collect();

        let original: Vec<Vec<f64>> = phases.clone();
        correct_multi_echo_wraps(&mut phases, &tes, &mask);

        for e in 0..3 {
            for i in 0..n {
                assert!((phases[e][i] - original[e][i]).abs() < 1e-10,
                    "No correction should be applied");
            }
        }
    }

    #[test]
    fn test_correct_multi_echo_wraps_with_offset() {
        let n = 64;
        let tes = [5.0, 10.0];
        let mask = vec![1u8; n];

        // Echo 0: phase = 0.5 for all voxels
        // Echo 1: should be 1.0 (TE ratio = 2) but is offset by +2π
        let mut phases = vec![
            vec![0.5; n],
            vec![1.0 + TWO_PI; n],
        ];

        correct_multi_echo_wraps(&mut phases, &tes, &mask);

        // After correction, echo 1 should be close to 1.0
        for i in 0..n {
            assert!((phases[1][i] - 1.0).abs() < 0.1,
                "Expected ~1.0, got {}", phases[1][i]);
        }
    }

    // =========================================================================
    // RomeoWeightType
    // =========================================================================

    #[test]
    fn test_weight_type_flags() {
        // Romeo: PC + PGC + PL + MC
        let f = RomeoWeightType::Romeo.weight_flags();
        assert_eq!(f, [true, true, true, true, false, false]);

        // Romeo2: PC + MC only
        let f = RomeoWeightType::Romeo2.weight_flags();
        assert_eq!(f, [true, false, false, true, false, false]);

        // Romeo6: all 6
        let f = RomeoWeightType::Romeo6.weight_flags();
        assert_eq!(f, [true, true, true, true, true, true]);
    }

    #[test]
    fn test_romeo4_is_romeo_alias() {
        assert_eq!(
            RomeoWeightType::Romeo.weight_flags(),
            RomeoWeightType::Romeo4.weight_flags(),
        );
    }

    #[test]
    fn test_params_weight_flags_roundtrip() {
        // Default params should match Romeo preset
        let default_params = RomeoParams::default();
        assert_eq!(default_params.weight_flags(), RomeoWeightType::Romeo.weight_flags());
        assert!(!default_params.bestpath);

        // from_weight_type preserves flags
        for wt in [RomeoWeightType::Romeo, RomeoWeightType::Romeo2,
                    RomeoWeightType::Romeo3, RomeoWeightType::Romeo6] {
            let params = RomeoParams::from_weight_type(wt);
            assert_eq!(params.weight_flags(), wt.weight_flags());
            assert!(!params.bestpath);
        }

        // BestPath sets the bestpath flag
        let bp = RomeoParams::from_weight_type(RomeoWeightType::BestPath);
        assert!(bp.bestpath);
    }

    #[test]
    fn test_params_custom_weight_combination() {
        // Custom combination: only phase coherence + mag weight (no preset matches this)
        let params = RomeoParams {
            phase_coherence: true,
            phase_gradient_coherence: false,
            phase_linearity: false,
            mag_coherence: false,
            mag_weight: true,
            mag_weight2: false,
            ..Default::default()
        };
        assert_eq!(params.weight_flags(), [true, false, false, false, true, false]);
    }

    // =========================================================================
    // Template selection
    // =========================================================================

    #[test]
    fn test_multi_echo_template_selection() {
        let n = 8;
        let total = n * n * n;
        let tes = [5.0, 10.0, 15.0];
        let phases: Vec<Vec<f64>> = tes.iter().map(|&te| {
            (0..total).map(|i| 0.02 * (i % n) as f64 * te).collect()
        }).collect();
        let mags: Vec<Vec<f64>> = tes.iter().map(|_| vec![1.0; total]).collect();
        let mask = vec![1u8; total];

        // Use echo 2 (index 1) as template
        let params = RomeoParams {
            template: 1,
            ..Default::default()
        };

        let result = unwrap_romeo_multi_echo(
            &phases, &mags, &tes, &mask,
            &params, &grid(n),
        );

        assert_eq!(result.len(), 3);
        for echo in &result {
            for &v in echo {
                assert!(v.is_finite());
            }
        }
    }

    // =========================================================================
    // Phase linearity
    // =========================================================================

    #[test]
    fn test_phase_linearity_triplet_linear() {
        // Perfectly linear: second derivative = 0 → linearity = 1.0
        assert!((phase_linearity_triplet(0.1, 0.2, 0.3) - 1.0).abs() < 1e-6);
        assert!((phase_linearity_triplet(-0.5, 0.0, 0.5) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_phase_linearity_triplet_jump() {
        // Large second derivative → low linearity
        let pl = phase_linearity_triplet(0.0, 0.0, PI);
        assert!(pl < 0.5, "Jump should give low linearity, got {}", pl);
    }

    // =========================================================================
    // BestPath weights
    // =========================================================================

    #[test]
    fn test_bestpath_weights_constant_phase() {
        let n = 6;
        let total = n * n * n;
        let phase = vec![0.0; total];
        let mask = vec![1u8; total];

        let weights = calculate_weights_bestpath(&phase, &mask, n, n, n);
        assert_eq!(weights.len(), 3 * total);

        // Constant phase → D=0 → R=0 → edge weight should be low
        // (BestPath assigns low weight to ambiguous/constant regions)
    }

    #[test]
    fn test_bestpath_weights_smooth_gradient() {
        let n = 8;
        let total = n * n * n;
        // Small smooth gradient in x
        let phase: Vec<f64> = (0..total).map(|i| 0.1 * (i % n) as f64).collect();
        let mask = vec![1u8; total];

        let weights = calculate_weights_bestpath(&phase, &mask, n, n, n);

        // The previous assertion here was `w <= 255` on a `u8`, which is true by type and so
        // could never fail. What is actually worth pinning is that the weights respond to the
        // phase at all, which needs two assertions pulling in opposite directions -- a bound
        // alone is satisfied by a function that ignores its input and returns a constant.
        //
        // `scale(w) = clamp((1 - w/10) * 255, 1, 255)` with `w` the summed second-order phase
        // inconsistency, so a smooth field sits near the top of the range. Note the range is
        // *not* very discriminating: measured means are 255.0 constant, 249.8 for this
        // gradient and 247.6 for uniform noise, so a high mean on its own says little.
        let assigned: Vec<u8> = weights.iter().copied().filter(|&w| w > 0).collect();
        assert!(!assigned.is_empty(), "no edge weight was assigned at all");
        assert!(
            assigned.iter().all(|&w| w >= 200),
            "a smooth gradient should score near the top of the range; min was {:?}",
            assigned.iter().min()
        );

        // A perfectly flat field has zero second-order difference everywhere, so every assigned
        // edge scores exactly 255. This gradient wraps every `n` voxels, and that discontinuity
        // pulls some edges below 255. A function returning a constant passes the bound above
        // and fails here.
        let flat = calculate_weights_bestpath(&vec![0.5f64; total], &mask, n, n, n);
        assert!(
            flat.iter().filter(|&&w| w > 0).all(|&w| w == 255),
            "a constant field should score a flat 255"
        );
        assert!(
            assigned.iter().any(|&w| w < 255),
            "the gradient's wrap discontinuity should pull some edges below 255"
        );
    }

    // =========================================================================
    // wrap_addition
    // =========================================================================

    #[test]
    fn test_unwrap_with_wrap_addition() {
        let n = 8;
        let total = n * n * n;
        let phase: Vec<f64> = (0..total).map(|i| {
            let x = (i % n) as f64;
            wrap_angle(0.1 * x)
        }).collect();
        let mask = vec![1u8; total];

        let params = RomeoParams {
            wrap_addition: 1.0,
            ..Default::default()
        };

        let unwrapped = unwrap_romeo(
            &phase, &[], None, 0.0, 0.0,
            &mask, &params, &grid(n),
        );

        for &v in &unwrapped {
            assert!(v.is_finite());
        }
    }

    // =========================================================================
    // Multi-seed
    // =========================================================================

    #[test]
    fn test_unwrap_multi_seed() {
        let n = 8;
        let total = n * n * n;
        let phase: Vec<f64> = (0..total).map(|i| {
            wrap_angle(0.05 * (i % n) as f64)
        }).collect();
        let mask = vec![1u8; total];

        let params = RomeoParams {
            max_seeds: 4,
            ..Default::default()
        };

        let unwrapped = unwrap_romeo(
            &phase, &[], None, 0.0, 0.0,
            &mask, &params, &grid(n),
        );

        for &v in &unwrapped {
            assert!(v.is_finite());
        }
    }

    #[test]
    fn test_unwrap_multi_seed_with_merge() {
        let n = 8;
        let total = n * n * n;
        let phase: Vec<f64> = (0..total).map(|i| {
            wrap_angle(0.05 * (i % n) as f64)
        }).collect();
        let mask = vec![1u8; total];

        let params = RomeoParams {
            max_seeds: 4,
            merge_regions: true,
            correct_regions: true,
            ..Default::default()
        };

        let unwrapped = unwrap_romeo(
            &phase, &[], None, 0.0, 0.0,
            &mask, &params, &grid(n),
        );

        for &v in &unwrapped {
            assert!(v.is_finite());
        }
    }

    // =========================================================================
    // correctglobal for 3D
    // =========================================================================

    #[test]
    fn test_correct_global_offset() {
        let n = 64;
        let mask = vec![1u8; n];
        // Phase that's all shifted by 2π
        let mut phase: Vec<f64> = vec![TWO_PI + 0.1; n];

        correct_global_offset(&mut phase, &mask);

        // After correction, should be close to 0.1
        for &v in &phase {
            assert!((v - 0.1).abs() < 0.2, "Expected ~0.1, got {}", v);
        }
    }

    #[test]
    fn test_unwrap_romeo_with_correctglobal() {
        let n = 4;
        let total = n * n * n;
        let phase = vec![0.1; total];
        let mask = vec![1u8; total];

        let params = RomeoParams {
            correct_global: true,
            ..Default::default()
        };

        let unwrapped = unwrap_romeo(
            &phase, &[], None, 0.0, 0.0,
            &mask, &params, &grid(n),
        );

        for &v in &unwrapped {
            assert!(v.is_finite());
        }
    }

    // =========================================================================
    // Weight type: BestPath integration
    // =========================================================================

    #[test]
    fn test_unwrap_romeo_bestpath() {
        let n = 8;
        let total = n * n * n;
        let phase: Vec<f64> = (0..total).map(|i| {
            wrap_angle(0.1 * (i % n) as f64)
        }).collect();
        let mask = vec![1u8; total];

        let params = RomeoParams {
            bestpath: true,
            ..Default::default()
        };

        let unwrapped = unwrap_romeo(
            &phase, &[], None, 0.0, 0.0,
            &mask, &params, &grid(n),
        );

        for &v in &unwrapped {
            assert!(v.is_finite());
        }
    }

    // =========================================================================
    // rescale_weight
    // =========================================================================

    #[test]
    fn test_rescale_weight() {
        // ROMEO.jl: bin = max(round((1 - w) * 255), 1); priority = 256 - bin
        assert_eq!(rescale_weight(1.0), 255);    // best (bin 1)
        assert_eq!(rescale_weight(0.5), 128);    // bin round(127.5) = 128 (ties to even)
        assert_eq!(rescale_weight(0.0), 1);      // worst valid (bin 255)
        assert_eq!(rescale_weight(-0.1), 0);     // invalid
        assert_eq!(rescale_weight(f64::NAN), 0); // invalid
    }
}