"""
FORCE: Forward modeling for complex microstructure estimation

This module provides signal simulation for matching-based diffusion MRI
reconstruction using multi-compartment tissue models.

The FORCE method generates a library of simulated signals representing
various tissue configurations (white matter with 1-3 fibers, gray matter,
CSF) and uses this library to reconstruct microstructural parameters
from real diffusion MRI data.

References
----------
:footcite:p:`Shah2025`
"""

import ast
from concurrent.futures import ProcessPoolExecutor, as_completed
import gc
import multiprocessing as mp
import os
import sys
import tempfile

import numpy as np

from dipy.data import default_sphere
from dipy.reconst.bingham import _single_bingham_to_sf as bingham_to_sf
from dipy.sims.voxel import all_tensor_evecs
from dipy.utils.logging import logger
from dipy.utils.multiproc import determine_num_processes


def dispersion_lut(target_sphere, odi_list):
    """Generate spherical functions for all directions and ODI values.
    Currently uses Bingham distribution, but can be extended to other
    dispersion models in the future.

    Parameters
    ----------
    target_sphere : ndarray (N, 3)
        Unit sphere vertices.
    odi_list : ndarray
        List of orientation dispersion index values.

    Returns
    -------
    SF : dict
        Nested dictionary mapping (vertex_index, odi) to spherical function (SF)
        values.
    """
    dispersion_sf_lut = {}
    for i in range(len(target_sphere)):
        vertex_key = tuple(target_sphere[i])
        dispersion_sf_lut[i] = {}

        for odi in odi_list:
            k = 1 / np.tan(np.pi / 2 * odi)
            evecs = all_tensor_evecs(vertex_key)
            major_axis, minor_axis = evecs[:, 1], evecs[:, 2]

            dispersion_sf_lut[i][odi] = bingham_to_sf(
                1, k, k, major_axis, minor_axis, target_sphere
            )

    return dispersion_sf_lut


def smallest_shell_bval(bvals, *, b0_threshold=50, shell_tolerance=50, n=1):
    """Find the smallest non-zero b-value shell.

    Parameters
    ----------
    bvals : ndarray
        B-values array.
    b0_threshold : float, optional
        Maximum b-value for b0 volumes.
    shell_tolerance : float, optional
        Tolerance for grouping shells.
    n : int, optional
        Number of smallest shells to return.

    Returns
    -------
    min_shell : float
        Smallest non-zero shell b-value.
    shell_mask : ndarray
        Boolean mask for volumes in that shell.
    """
    bvals = np.asarray(bvals, dtype=float)
    non_b0 = bvals > b0_threshold
    if not np.any(non_b0):
        raise ValueError("No non-b0 volumes found.")
    rounded = np.round(bvals[non_b0] / shell_tolerance) * shell_tolerance
    unique_shells = np.unique(rounded)
    if len(unique_shells) < n:
        raise ValueError(
            f"Only {len(unique_shells)} unique shells found, but n={n} requested."
        )
    else:
        min_shells = unique_shells[:n]

    rounded_all = np.round(bvals / shell_tolerance) * shell_tolerance
    shell_mask = np.isin(rounded_all, min_shells)
    return min_shells, shell_mask


def init_worker(base_seed=None):
    """Initialize a ProcessPoolExecutor worker with a unique RNG state.

    With ``base_seed=None``, seed = PID + high-resolution time so every
    worker has a different stream. With ``initargs=(base_seed,)`` for
    reproducibility, seed = base_seed + PID so workers differ but the run
    is reproducible for a fixed worker count.

    Parameters
    ----------
    base_seed : int, optional
        Base seed for reproducible worker initialization.
    """
    import random
    import time

    if base_seed is not None:
        seed = (int(base_seed) + os.getpid()) % (2**32)
    else:
        # Unique per process and run: PID + high-resolution time
        seed = (os.getpid() * (2**16) + int(time.perf_counter_ns())) % (2**32)
    np.random.seed(seed)
    random.seed(seed)


def _create_memmap(output_dir, name, dtype, shape):
    """Create a memory-mapped file."""
    path = os.path.join(output_dir, f"{name}.mmap")
    return np.memmap(path, mode="w+", dtype=dtype, shape=shape), path


def _generate_batch_worker(
    start_idx,
    batch_size,
    sphere,
    evecs,
    bingham,
    odi,
    bval,
    bvec,
    wm_thresh,
    tort,
    memmap_info,
    diffusivity_cfg,
):
    """Worker function for parallel batch generation.

    Opens memmaps by path in each process and writes directly.
    """
    from dipy.sims._force_core import create_mixed_signal, set_diffusivity_ranges
    from dipy.sims._multi_tensor_omp import multi_tensor

    # Unpack memmap info
    (
        signals_path,
        signals_shape,
        signals_dtype,
        labels_path,
        labels_shape,
        labels_dtype,
        num_fibers_path,
        num_fibers_shape,
        num_fibers_dtype,
        dispersion_path,
        dispersion_shape,
        dispersion_dtype,
        wm_fraction_path,
        wm_fraction_shape,
        wm_fraction_dtype,
        gm_fraction_path,
        gm_fraction_shape,
        gm_fraction_dtype,
        csf_fraction_path,
        csf_fraction_shape,
        csf_fraction_dtype,
        nd_path,
        nd_shape,
        nd_dtype,
        odfs_path,
        odfs_shape,
        odfs_dtype,
        ufa_wm_path,
        ufa_wm_shape,
        ufa_wm_dtype,
        ufa_voxel_path,
        ufa_voxel_shape,
        ufa_voxel_dtype,
        fraction_array_path,
        fraction_array_shape,
        fraction_array_dtype,
        wm_disp_path,
        wm_disp_shape,
        wm_disp_dtype,
        wm_d_par_path,
        wm_d_par_shape,
        wm_d_par_dtype,
        wm_d_perp_path,
        wm_d_perp_shape,
        wm_d_perp_dtype,
        gm_d_par_path,
        gm_d_par_shape,
        gm_d_par_dtype,
        csf_d_par_path,
        csf_d_par_shape,
        csf_d_par_dtype,
        f_ins_path,
        f_ins_shape,
        f_ins_dtype,
    ) = memmap_info

    # Apply diffusivity ranges in worker
    set_diffusivity_ranges(**diffusivity_cfg)

    # Open memmaps in read-write mode
    signals_mm = np.memmap(
        signals_path, mode="r+", dtype=signals_dtype, shape=signals_shape
    )
    labels_mm = np.memmap(
        labels_path, mode="r+", dtype=labels_dtype, shape=labels_shape
    )
    num_fibers_mm = np.memmap(
        num_fibers_path, mode="r+", dtype=num_fibers_dtype, shape=num_fibers_shape
    )
    dispersion_mm = np.memmap(
        dispersion_path, mode="r+", dtype=dispersion_dtype, shape=dispersion_shape
    )
    wm_fraction_mm = np.memmap(
        wm_fraction_path, mode="r+", dtype=wm_fraction_dtype, shape=wm_fraction_shape
    )
    gm_fraction_mm = np.memmap(
        gm_fraction_path, mode="r+", dtype=gm_fraction_dtype, shape=gm_fraction_shape
    )
    csf_fraction_mm = np.memmap(
        csf_fraction_path, mode="r+", dtype=csf_fraction_dtype, shape=csf_fraction_shape
    )
    nd_mm = np.memmap(nd_path, mode="r+", dtype=nd_dtype, shape=nd_shape)
    odfs_mm = np.memmap(odfs_path, mode="r+", dtype=odfs_dtype, shape=odfs_shape)
    ufa_wm_mm = np.memmap(
        ufa_wm_path, mode="r+", dtype=ufa_wm_dtype, shape=ufa_wm_shape
    )
    ufa_voxel_mm = np.memmap(
        ufa_voxel_path, mode="r+", dtype=ufa_voxel_dtype, shape=ufa_voxel_shape
    )
    fraction_array_mm = np.memmap(
        fraction_array_path,
        mode="r+",
        dtype=fraction_array_dtype,
        shape=fraction_array_shape,
    )
    wm_disp_mm = np.memmap(
        wm_disp_path, mode="r+", dtype=wm_disp_dtype, shape=wm_disp_shape
    )
    wm_d_par_mm = np.memmap(
        wm_d_par_path, mode="r+", dtype=wm_d_par_dtype, shape=wm_d_par_shape
    )
    wm_d_perp_mm = np.memmap(
        wm_d_perp_path, mode="r+", dtype=wm_d_perp_dtype, shape=wm_d_perp_shape
    )
    gm_d_par_mm = np.memmap(
        gm_d_par_path, mode="r+", dtype=gm_d_par_dtype, shape=gm_d_par_shape
    )
    csf_d_par_mm = np.memmap(
        csf_d_par_path, mode="r+", dtype=csf_d_par_dtype, shape=csf_d_par_shape
    )
    f_ins_mm = np.memmap(f_ins_path, mode="r+", dtype=f_ins_dtype, shape=f_ins_shape)

    for i in range(batch_size):
        idx = start_idx + i
        res = create_mixed_signal(
            sphere, evecs, bingham, odi, bval, bvec, multi_tensor, wm_thresh, tort
        )
        (
            signals_mm[idx],
            labels_mm[idx],
            num_fibers_mm[idx],
            dispersion_mm[idx],
            wm_fraction_mm[idx],
            gm_fraction_mm[idx],
            csf_fraction_mm[idx],
            nd_mm[idx],
            odfs_mm[idx],
            ufa_wm_mm[idx],
            ufa_voxel_mm[idx],
            fraction_array_mm[idx],
            wm_disp_mm[idx],
            wm_d_par_mm[idx],
            wm_d_perp_mm[idx],
            gm_d_par_mm[idx],
            csf_d_par_mm[idx],
            f_ins_mm[idx],
        ) = res

    # Flush memmaps
    for mm in [
        signals_mm,
        labels_mm,
        num_fibers_mm,
        dispersion_mm,
        wm_fraction_mm,
        gm_fraction_mm,
        csf_fraction_mm,
        nd_mm,
        odfs_mm,
        ufa_wm_mm,
        ufa_voxel_mm,
        fraction_array_mm,
        wm_disp_mm,
        wm_d_par_mm,
        wm_d_perp_mm,
        gm_d_par_mm,
        csf_d_par_mm,
        f_ins_mm,
    ]:
        mm.flush()

    return batch_size


def _main_is_guarded():
    """Return True if ``__main__`` has a top-level ``if __name__ == '__main__':`` guard.

    Parses the source of the ``__main__`` module with :mod:`ast` and inspects
    only the direct children of the module node (top-level statements).  Returns
    ``True`` also when the main module was invoked with ``python -m pkg.mod``
    (``__spec__`` is not None), or from an interactive session (no ``__file__``),
    because those cases do not cause spawn re-execution problems.

    Returns
    -------
    guarded : bool
        ``True`` if spawning workers is safe without causing recursive
        re-execution of the caller's script.
    """
    main = sys.modules.get("__main__")
    if main is None:
        return True
    if getattr(main, "__spec__", None) is not None:
        # Invoked as `python -m pkg.mod` — spawn imports the module, no re-run
        return True
    main_file = getattr(main, "__file__", None)
    if main_file is None:
        # Interactive session or frozen executable
        return True
    try:
        with open(main_file) as fh:
            source = fh.read()
        tree = ast.parse(source)
    except (OSError, SyntaxError):
        return False
    for node in ast.iter_child_nodes(tree):
        if not isinstance(node, ast.If):
            continue
        test = node.test
        if (
            isinstance(test, ast.Compare)
            and isinstance(test.left, ast.Name)
            and test.left.id == "__name__"
            and len(test.ops) == 1
            and isinstance(test.ops[0], ast.Eq)
            and len(test.comparators) == 1
            and isinstance(test.comparators[0], ast.Constant)
            and test.comparators[0].value == "__main__"
        ):
            return True
    return False


def generate_force_simulations(
    gtab,
    *,
    num_simulations=100000,
    output_dir=None,
    num_cpus=1,
    batch_size=1000,
    wm_threshold=0.5,
    tortuosity=False,
    odi_range=(0.01, 0.3),
    num_odi_values=10,
    diffusivity_config=None,
    dtype=np.float32,
    compute_dti=True,
    compute_dki=False,
    verbose=True,
):
    """Generate FORCE simulations.

    Creates a library of simulated diffusion MRI signals with
    corresponding microstructural parameters for matching-based
    reconstruction.

    Parameters
    ----------
    gtab : GradientTable
        Gradient table with b-values and b-vectors.
    num_simulations : int, optional
        Number of simulated voxels.
    output_dir : str, optional
        Directory for output files. If None, uses a temporary directory.
    num_cpus : int or None, optional
        Number of CPU cores for parallel processing.
        ``1`` runs in-process (no subprocess, safe everywhere).
        ``None`` or ``-1`` uses all available cores.
        Values ``< -1`` leave that many cores idle
        (e.g. ``-2`` = all minus one). ``0`` raises ``ValueError``.
    batch_size : int, optional
        Batch size for processing.
    wm_threshold : float, optional
        Minimum WM fraction to include fiber labels.
    tortuosity : bool, optional
        Use tortuosity constraint for perpendicular diffusivity.
    odi_range : tuple, optional
        (min, max) orientation dispersion index range.
    num_odi_values : int, optional
        Number of ODI values to sample.
    diffusivity_config : dict, optional
        Custom diffusivity ranges.
    dtype : dtype, optional
        Data type for outputs.
    compute_dti : bool, optional
        Compute DTI metrics (FA, MD, RD).
    compute_dki : bool, optional
        Compute DKI metrics (AK, RK, MK, KFA).
    verbose : bool, optional
        Enable progress output.

    Returns
    -------
    simulations : dict
        Simulations containing signals and parameters.
    """
    from tqdm import tqdm

    from dipy.reconst import dti
    from dipy.sims._force_core import set_diffusivity_ranges

    num_cpus = determine_num_processes(num_cpus)

    # Setup output directory
    _tmpdir_ctx = None
    if output_dir is None:
        _tmpdir_ctx = tempfile.TemporaryDirectory(prefix="force_sim_")
        output_dir = _tmpdir_ctx.name
    os.makedirs(output_dir, exist_ok=True)

    # Setup diffusivity config
    if diffusivity_config is None:
        diffusivity_config = {
            "wm_d_par_range": (2.0e-3, 3.0e-3),
            "wm_d_perp_range": (0.3e-3, 1.5e-3),
            "gm_d_iso_range": (0.7e-3, 1.2e-3),
            "csf_d": 3.0e-3,
        }

    set_diffusivity_ranges(**diffusivity_config)

    # Setup sphere and ODI values
    sphere = default_sphere
    target_sphere = np.ascontiguousarray(sphere.vertices, dtype=np.float64)
    n_dirs = target_sphere.shape[0]

    bvals = np.ascontiguousarray(gtab.bvals.astype(np.float64))
    bvecs = np.ascontiguousarray(gtab.bvecs.astype(np.float64))
    n_bvals = bvals.shape[0]

    # Pre-compute eigenvectors and dispersion lookup table
    evecs = np.array(
        [all_tensor_evecs(tuple(point)) for point in target_sphere],
        dtype=np.float64,
    )
    odi_list = np.linspace(odi_range[0], odi_range[1], num_odi_values).astype(
        np.float64
    )
    dispersion_sf = dispersion_lut(target_sphere, odi_list)

    label_dtype = np.uint8

    # Create memmaps for large arrays
    memmap_paths = {}

    def create_mm(name, dt, shape):
        mm, path = _create_memmap(output_dir, name, dt, shape)
        memmap_paths[name] = path
        return mm

    signals_mm = create_mm("signals", dtype, (num_simulations, n_bvals))
    labels_mm = create_mm("labels", label_dtype, (num_simulations, n_dirs))
    num_fibers_mm = create_mm("num_fibers", dtype, (num_simulations,))
    dispersion_mm = create_mm("dispersion", dtype, (num_simulations,))
    wm_fraction_mm = create_mm("wm_fraction", dtype, (num_simulations,))
    gm_fraction_mm = create_mm("gm_fraction", dtype, (num_simulations,))
    csf_fraction_mm = create_mm("csf_fraction", dtype, (num_simulations,))
    nd_mm = create_mm("nd", dtype, (num_simulations,))
    odfs_mm = create_mm("odfs", np.float16, (num_simulations, n_dirs))
    ufa_wm_mm = create_mm("ufa_wm", dtype, (num_simulations,))
    ufa_voxel_mm = create_mm("ufa_voxel", dtype, (num_simulations,))
    fraction_array_mm = create_mm("fraction_array", dtype, (num_simulations, 3))
    wm_disp_mm = create_mm("wm_disp", dtype, (num_simulations,))
    wm_d_par_mm = create_mm("wm_d_par", dtype, (num_simulations,))
    wm_d_perp_mm = create_mm("wm_d_perp", dtype, (num_simulations,))
    gm_d_par_mm = create_mm("gm_d_par", dtype, (num_simulations,))
    csf_d_par_mm = create_mm("csf_d_par", dtype, (num_simulations,))
    f_ins_mm = create_mm("f_ins", dtype, (num_simulations, 3))

    memmaps = [
        signals_mm,
        labels_mm,
        num_fibers_mm,
        dispersion_mm,
        wm_fraction_mm,
        gm_fraction_mm,
        csf_fraction_mm,
        nd_mm,
        odfs_mm,
        ufa_wm_mm,
        ufa_voxel_mm,
        fraction_array_mm,
        wm_disp_mm,
        wm_d_par_mm,
        wm_d_perp_mm,
        gm_d_par_mm,
        csf_d_par_mm,
        f_ins_mm,
    ]

    # Build batch specs
    num_batches_full = num_simulations // batch_size
    remainder = num_simulations % batch_size
    total_batches = num_batches_full + (1 if remainder > 0 else 0)

    batch_specs = []
    current_start = 0
    for batch_idx in range(total_batches):
        bs = batch_size if batch_idx < num_batches_full else remainder
        batch_specs.append((current_start, bs))
        current_start += bs

    # Pack memmap info for workers
    memmap_info = (
        memmap_paths["signals"],
        (num_simulations, n_bvals),
        dtype,
        memmap_paths["labels"],
        (num_simulations, n_dirs),
        label_dtype,
        memmap_paths["num_fibers"],
        (num_simulations,),
        dtype,
        memmap_paths["dispersion"],
        (num_simulations,),
        dtype,
        memmap_paths["wm_fraction"],
        (num_simulations,),
        dtype,
        memmap_paths["gm_fraction"],
        (num_simulations,),
        dtype,
        memmap_paths["csf_fraction"],
        (num_simulations,),
        dtype,
        memmap_paths["nd"],
        (num_simulations,),
        dtype,
        memmap_paths["odfs"],
        (num_simulations, n_dirs),
        np.float16,
        memmap_paths["ufa_wm"],
        (num_simulations,),
        dtype,
        memmap_paths["ufa_voxel"],
        (num_simulations,),
        dtype,
        memmap_paths["fraction_array"],
        (num_simulations, 3),
        dtype,
        memmap_paths["wm_disp"],
        (num_simulations,),
        dtype,
        memmap_paths["wm_d_par"],
        (num_simulations,),
        dtype,
        memmap_paths["wm_d_perp"],
        (num_simulations,),
        dtype,
        memmap_paths["gm_d_par"],
        (num_simulations,),
        dtype,
        memmap_paths["csf_d_par"],
        (num_simulations,),
        dtype,
        memmap_paths["f_ins"],
        (num_simulations, 3),
        dtype,
    )

    # Run simulations with progress bar
    pbar = tqdm(total=num_simulations, desc="Simulating", disable=not verbose)

    if num_cpus == 1:
        # Serial in-process path — no subprocess spawned, safe on all platforms
        init_worker()
        for start_idx, bs in batch_specs:
            batch_done = _generate_batch_worker(
                start_idx,
                bs,
                target_sphere,
                evecs,
                dispersion_sf,
                odi_list,
                bvals,
                bvecs,
                wm_threshold,
                tortuosity,
                memmap_info,
                diffusivity_config,
            )
            pbar.update(batch_done)
    elif sys.platform == "win32":
        if _main_is_guarded():
            ctx = mp.get_context("spawn")
            with ProcessPoolExecutor(
                max_workers=num_cpus, initializer=init_worker, mp_context=ctx
            ) as executor:
                futures = {
                    executor.submit(
                        _generate_batch_worker,
                        start_idx,
                        bs,
                        target_sphere,
                        evecs,
                        dispersion_sf,
                        odi_list,
                        bvals,
                        bvecs,
                        wm_threshold,
                        tortuosity,
                        memmap_info,
                        diffusivity_config,
                    ): (start_idx, bs)
                    for start_idx, bs in batch_specs
                }
                for future in as_completed(futures):
                    batch_done = future.result()
                    pbar.update(batch_done)
        else:
            logger.warning(
                "Parallel simulation on Windows requires your script to be "
                "protected with `if __name__ == '__main__':`. "
                "Falling back to serial execution."
            )
            init_worker()
            for start_idx, bs in batch_specs:
                batch_done = _generate_batch_worker(
                    start_idx,
                    bs,
                    target_sphere,
                    evecs,
                    dispersion_sf,
                    odi_list,
                    bvals,
                    bvecs,
                    wm_threshold,
                    tortuosity,
                    memmap_info,
                    diffusivity_config,
                )
                pbar.update(batch_done)
    else:
        # fork: workers are copies of the parent — __main__ is NOT re-run,
        # so no `if __name__ == '__main__':` guard is needed in the caller's
        # script. Safe for CLI scientific-computing scripts on macOS/Linux
        # (no GUI/ObjC).
        ctx = mp.get_context("fork")
        with ProcessPoolExecutor(
            max_workers=num_cpus, initializer=init_worker, mp_context=ctx
        ) as executor:
            futures = {
                executor.submit(
                    _generate_batch_worker,
                    start_idx,
                    bs,
                    target_sphere,
                    evecs,
                    dispersion_sf,
                    odi_list,
                    bvals,
                    bvecs,
                    wm_threshold,
                    tortuosity,
                    memmap_info,
                    diffusivity_config,
                ): (start_idx, bs)
                for start_idx, bs in batch_specs
            }

            for future in as_completed(futures):
                batch_done = future.result()
                pbar.update(batch_done)

    pbar.close()

    # Flush memmaps
    for mm in memmaps:
        mm.flush()

    # Compute DTI metrics
    fa_dti = np.zeros(num_simulations, dtype=dtype)
    md_dti = np.zeros(num_simulations, dtype=dtype)
    rd_dti = np.zeros(num_simulations, dtype=dtype)

    # Compute DKI metrics
    ak_arr = np.zeros(num_simulations, dtype=dtype)
    rk_arr = np.zeros(num_simulations, dtype=dtype)
    mk_arr = np.zeros(num_simulations, dtype=dtype)
    kfa_arr = np.zeros(num_simulations, dtype=dtype)

    if compute_dti:
        from dipy.core.gradients import gradient_table as gt_func

        min_b, shell_mask = smallest_shell_bval(bvals)
        b0_mask = bvals <= gtab.b0_threshold
        use_mask = shell_mask | b0_mask

        gtab_small = gt_func(bvals[use_mask], bvecs=bvecs[use_mask])
        dti_model = dti.TensorModel(gtab_small)

        dti_batch_size = 2000
        pbar = tqdm(total=num_simulations, desc="DTI fitting", disable=not verbose)

        for start in range(0, num_simulations, dti_batch_size):
            end = min(start + dti_batch_size, num_simulations)
            data_batch = signals_mm[start:end][:, use_mask]
            dti_fit = dti_model.fit(data_batch)
            fa_dti[start:end] = dti_fit.fa.astype(dtype)
            md_dti[start:end] = dti_fit.md.astype(dtype)
            rd_dti[start:end] = dti_fit.rd.astype(dtype)
            pbar.update(end - start)

        pbar.close()

    if compute_dki:
        from dipy.core.gradients import gradient_table as gt_func
        from dipy.reconst.dki import DiffusionKurtosisModel

        min_bs, shell_mask = smallest_shell_bval(bvals, n=2)
        b0_mask = bvals <= gtab.b0_threshold
        use_mask = shell_mask | b0_mask

        gtab_dki = gt_func(bvals[use_mask], bvecs=bvecs[use_mask])
        dki_model = DiffusionKurtosisModel(gtab_dki)

        dki_batch_size = 2000
        pbar = tqdm(
            total=num_simulations, desc="Kurtosis Estimation", disable=not verbose
        )

        for start in range(0, num_simulations, dki_batch_size):
            end = min(start + dki_batch_size, num_simulations)
            data_batch = signals_mm[start:end][:, use_mask]
            dki_fit = dki_model.multi_fit(data_batch)[0]
            ak_arr[start:end] = dki_fit.ak().astype(dtype)
            rk_arr[start:end] = dki_fit.rk().astype(dtype)
            mk_arr[start:end] = dki_fit.mk().astype(dtype)
            kfa_arr[start:end] = dki_fit.kfa.astype(dtype)
            pbar.update(end - start)

        pbar.close()

    # Build output simulations dict
    simulations = {
        "signals": np.array(signals_mm),
        "labels": np.array(labels_mm),
        "num_fibers": np.array(num_fibers_mm),
        "dispersion": np.array(dispersion_mm),
        "wm_fraction": np.array(wm_fraction_mm),
        "gm_fraction": np.array(gm_fraction_mm),
        "csf_fraction": np.array(csf_fraction_mm),
        "nd": np.array(nd_mm),
        "odfs": np.array(odfs_mm),
        "fa": fa_dti,
        "md": md_dti,
        "rd": rd_dti,
        "ufa_wm": np.array(ufa_wm_mm),
        "ufa_voxel": np.array(ufa_voxel_mm),
        "fraction_array": np.array(fraction_array_mm),
    }

    if compute_dki:
        simulations.update(
            {
                "ak": ak_arr,
                "rk": rk_arr,
                "mk": mk_arr,
                "kfa": kfa_arr,
            }
        )

    # Cleanup memmaps
    for mm in memmaps:
        if hasattr(mm, "base") and hasattr(mm.base, "close"):
            mm.base.close()
        del mm
    gc.collect()

    for _name, path in memmap_paths.items():
        try:
            if os.path.exists(path):
                os.remove(path)
        except OSError:
            pass

    gc.collect()

    if _tmpdir_ctx is not None:
        _tmpdir_ctx.cleanup()

    return simulations


def save_force_simulations(simulations, output_path):
    """Save FORCE simulations to compressed NPZ file.

    Parameters
    ----------
    simulations : dict
        FORCE simulations.
    output_path : str
        Path for output file (.npz extension).
    """
    np.savez_compressed(output_path, **simulations)


def load_force_simulations(input_path):
    """Load FORCE simulations from NPZ file.

    Parameters
    ----------
    input_path : str
        Path to simulations file.

    Returns
    -------
    simulations : dict
        FORCE simulations.
    """
    data = np.load(input_path, allow_pickle=False)
    return {key: data[key] for key in data.files}


def validate_diffusivity_config(config):
    """Validate diffusivity configuration dictionary.

    Parameters
    ----------
    config : dict
        Diffusivity configuration with keys:
        wm_d_par_range, wm_d_perp_range, gm_d_iso_range, csf_d.

    Returns
    -------
    valid : bool
        True if configuration is valid.

    Raises
    ------
    ValueError
        If configuration is invalid.
    """
    required_keys = ["wm_d_par_range", "wm_d_perp_range", "gm_d_iso_range", "csf_d"]

    for key in required_keys:
        if key not in config:
            raise ValueError(f"Missing required key: {key}")

    for key in ["wm_d_par_range", "wm_d_perp_range", "gm_d_iso_range"]:
        val = config[key]
        if hasattr(val, "__len__"):
            if len(val) != 2:
                raise ValueError(f"{key} must be scalar or (min, max) tuple")
            if val[0] > val[1]:
                raise ValueError(f"{key} min must be <= max")

    csf_d = config["csf_d"]
    if not isinstance(csf_d, (int, float)) or csf_d <= 0:
        raise ValueError("csf_d must be a positive number")

    return True


def get_default_diffusivity_config():
    """Get default diffusivity configuration.

    Returns
    -------
    config : dict
        Default diffusivity ranges for FORCE simulation.
    """
    return {
        "wm_d_par_range": (2.0e-3, 3.0e-3),
        "wm_d_perp_range": (0.3e-3, 1.5e-3),
        "gm_d_iso_range": (0.7e-3, 1.2e-3),
        "csf_d": 3.0e-3,
    }
