"""
Synthetic data generation for testing, validation, ML training data generation.

This module provides tools for generating realistic simulated spectroscopy
data from models, with support for different detector types and noise models.
Use for:
- Testing fitting algorithms with known ground truth
- Exploring parameter sensitivity and identifiability
- Optimizing experimental design (SNR requirements)
- Generating training data for machine learning
- Validating analysis pipelines

Key Features
------------
- Two detector types: analog and photon counting
- Multiple noise models: Poisson, Gaussian, or none
- 1D and 2D spectrum simulation
- Batch generation for statistical analysis
- Parameter sweeping (grid/random/uniform) for ML training

Detector Types
--------------
**Analog Detectors** (CCD, photodiodes, lock-in amplifiers):
- Continuous signal output
- Additive noise (Gaussian or Poisson)
- Noise level controlled by noise_level parameter

**Photon Counting** (APD, photomultiplier, event mode):
- Discrete photon events
- Shot noise inherent (Poisson statistics)
- Count rate determines signal-to-noise ratio

Workflow
--------
1. Testing and Validation

   - Create model with trspecfit.mcp.Model
   - Initialize Simulator with model and noise parameters
   - Generate data with simulate_1d() or simulate_2d()
     OR generate multiple realizations with simulate_n()
   - Save data and ground truth with save_data()
   - Fit simulated data to validate fitting pipeline

2. Machine Learning Training Data Generation

   - Create model with trspecfit.mcp.Model
   - Define parameter space using trspecfit.utils.sweep.ParameterSweep
   - Initialize Simulator with model and noise parameters
   - Generate multiple realizations (n) for each parameter combination
     (data, ground truth, and relevant metadata get saved automatically)

Examples
--------
See examples/simulator/ directory for complete workflows.
"""

import json
import warnings
from pathlib import Path
from typing import cast

import h5py
import matplotlib.pyplot as plt
import numpy as np

from trspecfit.mcp import Model
from trspecfit.utils import plot as uplt
from trspecfit.utils.hdf5 import require_group
from trspecfit.utils.sweep import ParameterSweep


#
#
class Simulator:
    """
    Simulate 2D time- and energy-resolved spectroscopy data with noise.

    This class generates synthetic data based on a model, adding realistic noise
    to simulate experimental measurements. Supports both analog detectors (with
    additive noise) and photon counting detectors (with shot noise).

    Parameters
    ----------
    model : Model
        Model instance from trspecfit.mcp with defined components and parameters.
        Must have energy and time axes set before simulation.
    detection : {'analog', 'photon_counting'}, default='analog'
        Detection technique to simulate:
        - 'analog': Continuous signal with additive noise
        - 'photon_counting': Discrete photon events with Poisson statistics
    noise_level : float, default=0.05
        Noise amplitude for analog detectors (0.0-1.0 for relative noise).
        Larger values = more noise. Ignored for photon_counting.
    noise_type : {'poisson', 'gaussian', 'none'}, default='poisson'
        Type of noise for analog detectors:
        - 'poisson': Shot noise (realistic for low light)
        - 'gaussian': White noise (simpler, faster)
        - 'none': No noise (testing and debugging)
        Ignored for photon_counting (always Poisson).
    counts_per_delay : int, optional
        Total photon count per time delay (photon_counting only).
        Directly sets signal-to-noise ratio. Mutually exclusive with
        count_rate + integration_time.
    count_rate : float, optional
        Photon count rate in Hz (photon_counting only).
        Combined with integration_time to compute counts_per_delay.
    integration_time : float, optional
        Integration time per delay point in seconds (photon_counting only).
        Combined with count_rate to compute counts_per_delay.
    seed : int, optional
        Random seed for reproducibility. If None, uses random initialization.

    Attributes
    ----------
    model : Model
        Model instance used for simulation
    detection : str
        Detection type ('analog' or 'photon_counting')
    seed : int or None
        Random seed value
    noise_level : float
        Analog detector noise level
    noise_type : str
        Analog detector noise type
    counts_per_delay : int
        Photon counting detector count budget
    count_rate : float or None
        Photon counting detector rate
    integration_time : float or None
        Photon counting integration time per delay
    data_clean : ndarray or None
        Most recently generated clean (noiseless) data
    data_noisy : ndarray or None
        Most recently generated noisy data
    noise : ndarray or None
        Most recently generated noise component (noisy - clean)

    Examples
    --------
    See examples/simulator/ directory for complete workflows.

    Notes
    -----
    **Analog vs. Photon Counting:**

    Analog detectors (CCD, photodiode, lock-in):
    - Pros: High dynamic range, simple operation
    - Cons: Read noise, dark current
    - Simulation: Continuous signal + additive noise

    Photon counting (APD, PMT, event mode):
    - Pros: No read noise, single-photon sensitivity
    - Cons: Dead time, count rate limits, pulse pileup
    - Simulation: Discrete events following Poisson statistics

    **Noise Level Selection:**

    For analog detectors, noise_level is relative to signal:
    - 0.01 (1%): Very clean, ideal conditions
    - 0.05 (5%): Typical good data
    - 0.10 (10%): Moderate noise, still fittable
    - 0.20 (20%): Challenging, may need averaging

    For photon counting, SNR set by counts_per_delay:
    - 100 counts: SNR ~ 10 (marginal)
    - 1000 counts: SNR ~ 32 (good)
    - 10000 counts: SNR ~ 100 (excellent)

    **Photon Counting Parameter Resolution:**

    The simulator resolves photon counting parameters as:
    1. If counts_per_delay specified directly → use it
    2. Else if count_rate and integration_time specified → compute counts_per_delay
    3. Else → estimate from model scale (prints warning)

    The third case assumes model amplitudes represent realistic count rates,
    which may not be true. Always specify counts_per_delay or (count_rate,
    integration_time) explicitly for accurate photon counting simulation.

    **Memory Usage:**

    Large 2D datasets can use significant memory:
    - Single dataset: ~8 MB per 1000×500 spectrum (float64)
    - simulate_n(n=100): ~800 MB for same size
    - Consider smaller grids or batch processing for large n

    See Also
    --------
    trspecfit.mcp.Model : Model class for simulation
    simulate_1d : Generate 1D spectrum
    simulate_2d : Generate 2D spectrum
    simulate_n : Generate multiple realizations
    save_data : Save simulated data to HDF5
    """

    #
    def __init__(
        self,
        model: Model,
        detection: str = "analog",
        noise_level: float = 0.05,
        noise_type: str = "poisson",
        counts_per_delay: int | None = None,
        count_rate: float | None = None,
        integration_time: float | None = None,
        seed: int | None = None,
    ) -> None:
        """
        Initialize simulator with a model and noise parameters.

        Parameters
        ----------
        model : Model
            Model instance with defined components and parameters.
            Must have model.energy and (for 2D) model.time set.
        detection : {'analog', 'photon_counting'}, default='analog'
            Detection technique to simulate
        noise_level : float, default=0.05
            Noise amplitude for analog detectors (0.0-1.0 relative to signal)
        noise_type : {'poisson', 'gaussian', 'none'}, default='poisson'
            Noise type for analog detectors
        counts_per_delay : int, optional
            Total photons per delay (photon_counting only)
        count_rate : float, optional
            Photon rate in Hz (photon_counting only)
        integration_time : float, optional
            Integration time per delay in seconds (photon_counting only)
        seed : int, optional
            Random seed for reproducibility

        Raises
        ------
        ValueError
            If detection type is invalid
            If counts_per_delay ≤ 0 (after estimation)

        Examples
        --------
        >>> # Analog with default settings
        >>> sim = Simulator(model)

        >>> # Analog with custom noise
        >>> sim = Simulator(model, noise_level=0.1, noise_type='gaussian')

        >>> # Photon counting with direct count specification
        >>> sim = Simulator(model, detection='photon_counting',
        ...                 counts_per_delay=5000)

        >>> # Photon counting with rate specification
        >>> sim = Simulator(model, detection='photon_counting',
        ...                 count_rate=1e6, integration_time=0.01)

        >>> # Reproducible simulation
        >>> sim = Simulator(model, seed=42)
        """

        self.model = model
        self.detection = detection.lower()
        self.seed = seed

        # Analog detector parameters
        self.noise_level = noise_level
        self.noise_type = noise_type.lower()

        # Photon counting parameters
        self.counts_per_delay: int | None = counts_per_delay
        self.count_rate: float | None = count_rate
        self.integration_time: float | None = integration_time

        # Validate and resolve photon counting parameters
        if self.detection == "photon_counting":
            self._resolve_photon_counting_params()

        # Validate detection type
        if self.detection not in ["analog", "photon_counting"]:
            raise ValueError(
                f"detection must be 'analog' or 'photon_counting', got '{detection}'"
            )

        # Set random seed if provided
        self.rng = np.random.default_rng(seed)

        # Storage for simulated data
        self.data_clean: np.ndarray | None = None  # Without noise
        self.data_noisy: np.ndarray | None = None  # With noise
        self.noise: np.ndarray | None = None  # Just the noise component

    #
    def __repr__(self) -> str:
        return f"Simulator(model='{self.model.name}', detection='{self.detection}')"

    #
    def _resolve_photon_counting_params(self) -> None:
        """
        Resolve photon counting parameters. User can specify either:
        - counts_per_delay directly, OR
        - count_rate and integration_time (will compute counts_per_delay), OR
        - nothing (will estimate from model scale with warning)
        """

        if self.counts_per_delay is not None:
            # Direct specification takes precedence
            if self.count_rate is not None or self.integration_time is not None:
                warnings.warn(
                    "counts_per_delay specified directly."
                    " Ignoring count_rate and integration_time.",
                    stacklevel=2,
                )
        elif self.count_rate is not None and self.integration_time is not None:
            # Calculate from rate and time
            self.counts_per_delay = int(self.count_rate * self.integration_time)
        else:
            # Estimate from model: assume clean_data represents count rates
            # Generate clean data to get the scale
            if self.model.time is not None and len(self.model.time) > 0:
                # 2D data
                self.model.create_value_2d()
                clean_data = self.model.value_2d
            else:
                # 1D data
                self.model.create_value_1d()
                clean_data = self.model.value_1d
            if clean_data is None:
                raise RuntimeError(
                    "Model evaluation did not produce clean data"
                    " for photon counting estimation"
                )

            # Estimate counts_per_delay as the average total signal per time step
            # This assumes the model amplitudes are in a realistic count rate range
            signal_positive = np.abs(clean_data)
            if clean_data.ndim == 2:
                row_totals = np.sum(signal_positive, axis=1)
                self.counts_per_delay = int(np.mean(row_totals))
            else:
                self.counts_per_delay = int(np.sum(signal_positive))

            warnings.warn(
                f"No photon count specified for photon_counting detection. "
                f"Estimating from model: {self.counts_per_delay:.2e} counts/delay. "
                f"For accurate simulation, specify counts_per_delay"
                f" or (count_rate, integration_time). "
                f"This estimate assumes your model amplitudes"
                f" represent realistic count rates.",
                stacklevel=2,
            )

        # Ensure counts_per_delay is positive
        if self.counts_per_delay <= 0:
            raise ValueError(
                f"counts_per_delay must be positive, got {self.counts_per_delay}.\n"
                f"Model may have negative or zero signal."
                f" Check your model or specify counts_per_delay manually."
            )

    #
    def generate_clean_data(self, dim: int = 2, t_ind: int = 0) -> np.ndarray:
        """
        Generate clean data from model (no noise).

        Parameters
        ----------
        dim : int
            Dimension (1 for 1D, 2 for 2D).
        t_ind : int
            Time index for 1D simulations (ignored for 2D).

        Returns
        -------
        ndarray
            Clean data array (1D or 2D depending on dim).
        """

        if dim == 1:
            self.model.create_value_1d(t_ind=t_ind, return_1d=False)
            assert self.model.value_1d is not None
            self.data_clean = self.model.value_1d.copy()
        elif dim == 2:
            self.model.create_value_2d()
            assert self.model.value_2d is not None
            self.data_clean = self.model.value_2d.copy()
        else:
            raise ValueError(f"dim must be 1 or 2, got {dim}")

        if self.data_clean is None:
            raise RuntimeError("Model evaluation did not produce clean data")
        return self.data_clean

    #
    def add_noise(
        self, clean_data: np.ndarray, dim: int = 2
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Add noise to clean data based on detection technique.

        Parameters
        ----------
        clean_data : ndarray
            Clean data array (1D or 2D).
        dim : int
            Dimension (1 for 1D, 2 for 2D).

        Returns
        -------
        tuple of (ndarray, ndarray)
            Tuple of (noisy_data, noise).
        """

        if self.detection == "analog":
            # Use traditional noise addition
            if dim == 1:
                noise = self._generate_noise_analog_1d(clean_data)
            elif dim == 2:
                noise = self._generate_noise_analog_2d(clean_data)
            else:
                raise ValueError(f"dim must be 1 or 2, got {dim}")

            noisy_data = clean_data + noise

        elif self.detection == "photon_counting":
            # Sample photons according to signal distribution
            if dim == 1:
                noisy_data = self._sample_photons_1d(clean_data)
            elif dim == 2:
                noisy_data = self._sample_photons_2d(clean_data)
            else:
                raise ValueError(f"dim must be 1 or 2, got {dim}")

            noise = noisy_data - clean_data
        else:
            raise ValueError(f"Unknown detection type: {self.detection}")

        return noisy_data, noise

    #
    def simulate_1d(self, t_ind: int = 0) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Simulate 1D spectrum (energy-resolved) at a specific time point.

        Generates a single energy-resolved spectrum from the model at the
        specified time index, adds appropriate noise for the detector type,
        and stores results for later access.

        Parameters
        ----------
        t_ind : int, default=0
            Time index for which to generate spectrum.
            For models without time-dependence, use default 0.

        Returns
        -------
        clean_data : ndarray
            Noiseless spectrum from model (shape: [n_energy])
        noisy_data : ndarray
            Spectrum with added noise (shape: [n_energy])
        noise : ndarray
            Noise component (noisy - clean, shape: [n_energy])

        Examples
        --------
        >>> # Simulate baseline spectrum
        >>> sim = Simulator(model, noise_level=0.05)
        >>> clean, noisy, noise = sim.simulate_1d(t_ind=0)
        >>>
        >>> # Plot comparison
        >>> plt.plot(model.energy, clean, 'k-', label='Clean')
        >>> plt.plot(model.energy, noisy, 'r.', label='Noisy', ms=2)
        >>> plt.legend()

        >>> # Calculate SNR
        >>> snr = sim.get_snr()
        >>> print(f"Signal-to-noise ratio: {snr:.1f}")

        >>> # Simulate different time points
        >>> for t_i in [0, 50, 100]:
        ...     clean, noisy, noise = sim.simulate_1d(t_ind=t_i)
        ...     plt.plot(model.energy, noisy, label=f't={model.time[t_i]:.1f}')

        Notes
        -----
        Results are stored in simulator attributes for later access:
        - self.data_clean: (Last) clean spectrum
        - self.data_noisy: Last noisy spectrum
        - self.noise: Last noise realization

        Can access these without re-simulation:
        >>> sim.simulate_1d()
        >>> snr = sim.get_snr()  # Uses stored data
        >>> sim.plot_comparison(dim=1)  # Uses stored data

        See Also
        --------
        simulate_2d : Simulate full 2D spectrum
        simulate_n : Generate multiple 1D realizations
        plot_comparison : Visualize results
        """

        # Generate clean spectrum from model
        clean_data = self.generate_clean_data(dim=1, t_ind=t_ind)

        # Add noise
        noisy_data, noise = self.add_noise(clean_data, dim=1)

        # Store for later use (e.g., plotting, SNR calculation)
        self.data_clean = clean_data
        self.data_noisy = noisy_data
        self.noise = noise

        return clean_data, noisy_data, noise

    #
    def simulate_2d(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Simulate 2D spectrum (time- and energy-resolved).

        Generates a complete 2D time- and energy-resolved spectrum from the
        model, adds appropriate noise for each time point, and stores results.

        Returns
        -------
        clean_data : ndarray
            Noiseless 2D spectrum from model (shape: [n_time, n_energy])
        noisy_data : ndarray
            2D spectrum with added noise (shape: [n_time, n_energy])
        noise : ndarray
            Noise component (noisy - clean, shape: [n_time, n_energy])

        Examples
        --------
        >>> # Basic 2D simulation
        >>> sim = Simulator(model, noise_level=0.05)
        >>> clean, noisy, noise = sim.simulate_2d()
        >>>
        >>> # Visualize with built-in plotter
        >>> sim.plot_comparison(dim=2)

        >>> # Manual visualization
        >>> fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
        >>> ax1.pcolormesh(model.energy, model.time, clean)
        >>> ax1.set_title('Clean Model')
        >>> ax2.pcolormesh(model.energy, model.time, noisy)
        >>> ax2.set_title(f'Noisy (SNR={sim.get_snr():.1f})')

        >>> # Test fitting on simulated data
        >>> clean, noisy, noise = sim.simulate_2d()
        >>> # ... set up fitting ...
        >>> file.data = noisy  # Use noisy data for fit
        >>> file.fit_2d(model_name='test', stages=2)
        >>> # Compare fitted vs. true parameters

        >>> # Vary noise level to study impact
        >>> for noise_level in [0.01, 0.05, 0.10]:
        ...     sim.set_noise_level(noise_level)
        ...     clean, noisy, noise = sim.simulate_2d()
        ...     snr = sim.get_snr()
        ...     print(f"Noise level {noise_level:.2f}: SNR = {snr:.1f}")

        Notes
        -----
        **Noise Application:**

        For analog detectors, noise is added independently at each pixel.
        For photon counting, photons are distributed across all pixels
        according to the signal probability distribution, then reconverted
        to same scale as input for direct comparison.

        **Performance:**

        Simulation time scales with:
        - Model evaluation time (dominates for complex models)
        - Array size (n_time × n_energy)
        - Noise generation method

        Typical times:
        - Simple model, 200×500 array: ~0.1-1 second
        - Complex model with time-dependence: ~1-10 seconds
        - Photon counting slightly slower than analog

        **Memory:**

        Three arrays stored (clean, noisy, noise), each:
        - Size: n_time × n_energy × 8 bytes (float64)
        - Example: 200×500 = ~2.4 MB per array, ~7.2 MB total

        See Also
        --------
        simulate_1d : Simulate single spectrum
        simulate_n : Generate multiple 2D realizations
        plot_comparison : Visualize results
        save_data : Save to HDF5 file
        """

        # Generate clean spectrum from model
        clean_data = self.generate_clean_data(dim=2)

        # Add noise
        noisy_data, noise = self.add_noise(clean_data, dim=2)

        # Store for later use (e.g., plotting, SNR calculation)
        self.data_clean = clean_data
        self.data_noisy = noisy_data
        self.noise = noise

        return clean_data, noisy_data, noise

    #
    def simulate_n(
        self, n: int, *, dim: int = 2, t_ind: int = 0, show_progress: bool = True
    ) -> tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]:
        """
        Generate n simulated datasets with independent noise realizations.

        Generates the clean data ONCE from the model, then adds n independent
        noise realizations. Use for statistical analysis of fitting algorithms
        and uncertainty quantification or machine learning model training.

        Parameters
        ----------
        n : int
            Number of datasets to generate (must be >= 1)
        dim : {1, 2}, default=2
            Dimensionality:
            - 1: Generate 1D spectra
            - 2: Generate 2D spectra
        t_ind : int, default=0
            Time index for 1D simulations (ignored for dim=2)
        show_progress : bool, default=True
            Print progress updates during generation

        Returns
        -------
        clean_data : ndarray
            Single clean dataset (1D or 2D depending on dim).
            Same for all n realizations (generated once).
        noisy_data_list : list of ndarray
            List of n noisy datasets, each with independent noise.
            Each element has same shape as clean_data.
        noise_list : list of ndarray
            List of n noise realizations (noisy - clean for each dataset).
            Each element has same shape as clean_data.

        Examples
        --------
        >>> # Generate 20 independent noisy datasets
        >>> sim = Simulator(model, noise_level=0.05)
        >>> clean, noisy_list, noise_list = sim.simulate_n(n=20, dim=2)
        >>>
        >>> # Fit each dataset and analyze parameter distribution
        >>> fitted_params = []
        >>> for noisy_data in noisy_list:
        ...     file.data = noisy_data
        ...     file.fit_2d('test', stages=2)
        ...     fitted_params.append(model.lmfit_pars['amplitude'].value)
        >>>
        >>> # Check parameter recovery
        >>> true_value = model.lmfit_pars['amplitude'].value
        >>> mean_fitted = np.mean(fitted_params)
        >>> std_fitted = np.std(fitted_params)
        >>> print(f"True: {true_value:.2f}")
        >>> print(f"Mean fitted: {mean_fitted:.2f} ± {std_fitted:.2f}")

        >>> # Analyze noise statistics
        >>> noise_mean = np.mean(noise_list, axis=0)
        >>> noise_std = np.std(noise_list, axis=0)
        >>>
        >>> # Should be close to zero (unbiased)
        >>> print(f"Noise mean: {np.mean(noise_mean):.2e}")
        >>> # Should match noise_level * signal scale
        >>> print(f"Noise std: {np.mean(noise_std):.2e}")

        >>> # Save multiple realizations for later use
        >>> clean, noisy_list, noise_list = sim.simulate_n(n=100, dim=2)
        >>> sim.save_data(
        ...     filepath='simulations/batch_001.h5',
        ...     n_data=noisy_list
        ... )

        >>> # Test convergence of fitted parameters with n
        >>> for n_datasets in [5, 10, 20, 50]:
        ...     clean, noisy_list, _ = sim.simulate_n(n=n_datasets, dim=2)
        ...     # ... fit each and compute parameter statistics ...
        ...     print(f"n={n_datasets}: parameter std = {param_std:.3f}")

        Notes
        -----
        **Efficiency:**

        Generating clean data once and adding n noise realizations is much
        faster than generating n complete simulations:

        - This method: 1 model evaluation + n noise additions
        - n separate simulate_2d calls: n model evaluations + n noise additions

        For complex models where evaluation is slow, this can save
        minutes to hours of computation time.

        **Statistical Analysis:**

        This function enables:
        - Monte Carlo analysis of fitting uncertainty
        - Algorithm validation (can recover true parameters?)
        - Bias detection (systematic fitting errors)
        - Confidence interval validation (coverage probability)
        - Experimental design optimization (required SNR)

        **Memory Considerations:**

        All n datasets stored in memory as lists:
        - Memory usage: n × (n_time × n_energy × 8 bytes)
        - Example: 100 datasets of 200×500 = ~800 MB

        For very large n or large arrays, consider:
        - Processing in batches
        - Saving to disk incrementally
        - Using generator pattern instead of list

        **Progress Display:**

        When show_progress=True, prints:
        - "Generating clean data from model... Done"
        - "Adding noise to dataset i/n" (updates in place)
        - "Generated n noisy datasets successfully"

        Set show_progress=False for batch processing or when redirecting output.

        See Also
        --------
        simulate_1d : Single 1D simulation
        simulate_2d : Single 2D simulation
        save_data : Save multiple datasets to HDF5
        """

        if n < 1:
            raise ValueError(f"n must be >= 1, got {n}")

        # Generate clean data ONCE
        if show_progress:
            print("Generating clean data from model...", end=" ")
        clean_data = self.generate_clean_data(dim=dim, t_ind=t_ind)
        if show_progress:
            print("Done")

        # Generate n independent noise realizations
        noisy_data_list = []
        noise_list = []

        for i in range(n):
            if show_progress:
                print(f"Adding noise to dataset {i + 1}/{n}", end="\r")

            # Just add noise to the same clean data
            noisy_data, noise = self.add_noise(clean_data, dim=dim)

            noisy_data_list.append(noisy_data)
            noise_list.append(noise)

        if show_progress:
            print(f"Generated {n} noisy datasets successfully")

        # Store last realization for later use
        self.data_clean = clean_data
        self.data_noisy = noisy_data_list[-1]
        self.noise = noise_list[-1]

        return clean_data, noisy_data_list, noise_list

    #
    def _generate_noise_analog_1d(self, signal: np.ndarray) -> np.ndarray:
        """
        Generate 1D noise array for analog detectors.

        Parameters
        ----------
        signal : ndarray
            Clean signal array.

        Returns
        -------
        ndarray
            Noise array with same shape as signal.
        """

        if self.noise_type == "none":
            return np.zeros_like(signal)

        if self.noise_type == "gaussian":
            # Gaussian noise with amplitude proportional to noise_level
            noise_amp = self.noise_level * np.max(np.abs(signal))
            return cast("np.ndarray", self.rng.normal(0, noise_amp, signal.shape))

        if self.noise_type == "poisson":
            # Poisson noise: scale signal to photon counts, add noise, scale back
            # Avoid negative values in Poisson distribution
            signal_positive = np.abs(signal)

            # Scale signal to photon counts (higher = less relative noise)
            scale_factor = 1.0 / (self.noise_level + 1e-10)  # Avoid division by zero
            signal_scaled = signal_positive * scale_factor

            # Generate Poisson noise
            noisy_scaled = self.rng.poisson(signal_scaled)

            # Scale back and compute noise component
            signal_noisy = noisy_scaled / scale_factor
            noise = signal_noisy - signal_positive

            # Restore original sign
            return cast("np.ndarray", noise * np.sign(signal))

        raise ValueError(
            f"Unknown noise type: {self.noise_type}. "
            "Use 'poisson', 'gaussian', or 'none'"
        )

    #
    def _generate_noise_analog_2d(self, signal: np.ndarray) -> np.ndarray:
        """
        Generate 2D noise array for analog detectors.

        Parameters
        ----------
        signal : ndarray
            Clean 2D signal array.

        Returns
        -------
        ndarray
            2D noise array with same shape as signal.
        """

        if self.noise_type == "none":
            return np.zeros_like(signal)

        if self.noise_type == "gaussian":
            # Gaussian noise with amplitude proportional to noise_level
            noise_amp = self.noise_level * np.max(np.abs(signal))
            return cast("np.ndarray", self.rng.normal(0, noise_amp, signal.shape))

        if self.noise_type == "poisson":
            # Poisson noise: scale signal to photon counts, add noise, scale back
            signal_positive = np.abs(signal)

            # Scale signal to photon counts
            scale_factor = 1.0 / (self.noise_level + 1e-10)
            signal_scaled = signal_positive * scale_factor

            # Generate Poisson noise
            noisy_scaled = self.rng.poisson(signal_scaled)

            # Scale back and compute noise component
            signal_noisy = noisy_scaled / scale_factor
            noise = signal_noisy - signal_positive

            # Restore original sign
            return cast("np.ndarray", noise * np.sign(signal))

        raise ValueError(
            f"Unknown noise type: {self.noise_type}. "
            "Use 'poisson', 'gaussian', or 'none'"
        )

    #
    def _sample_photons_1d(self, signal: np.ndarray) -> np.ndarray:
        """
        Sample photons for 1D photon counting detection.

        Each energy pixel independently draws from a Poisson distribution
        where the expected count is proportional to the signal intensity.
        The signal is scaled so the total expected counts across all energy
        pixels equals counts_per_delay.

        Parameters
        ----------
        signal : ndarray
            Clean signal array (represents expected count rate).

        Returns
        -------
        ndarray
            Noisy data array in same units as input signal.
        """

        signal_positive = np.abs(signal)
        total_signal = np.sum(signal_positive)

        if total_signal == 0:
            return np.zeros_like(signal)
        if self.counts_per_delay is None:
            raise ValueError(
                "counts_per_delay must be defined for photon counting simulation"
            )

        # Scale signal to expected photon counts
        # Total expected counts across all pixels = counts_per_delay
        scale_factor = self.counts_per_delay / total_signal
        expected_counts = signal_positive * scale_factor

        # Independent Poisson draw per pixel
        photon_counts = self.rng.poisson(expected_counts)

        # Scale back to original signal units and restore sign
        # (negative signal = bleach/emission, positive = absorption)
        noisy_data = photon_counts / scale_factor * np.sign(signal)

        return cast("np.ndarray", noisy_data)

    #
    def _sample_photons_2d(self, signal: np.ndarray) -> np.ndarray:
        """
        Sample photons for 2D photon counting detection.

        Applies independent Poisson noise per pixel across the full 2D array.
        The signal is scaled so the average total expected counts per time step
        equals counts_per_delay. Time steps with stronger signal naturally
        accumulate more photons (better SNR), matching real experiments where
        each time delay is measured for the same integration time.

        Parameters
        ----------
        signal : ndarray
            Clean 2D signal array (shape: [n_time, n_energy]).

        Returns
        -------
        ndarray
            Noisy 2D data array in same units as input signal.
        """

        signal_positive = np.abs(signal)

        # Average total signal per time step
        row_totals = np.sum(signal_positive, axis=1)
        mean_row_total = np.mean(row_totals)

        if mean_row_total == 0:
            return np.zeros_like(signal)
        if self.counts_per_delay is None:
            raise ValueError(
                "counts_per_delay must be defined for photon counting simulation"
            )

        # Scale so the average row has counts_per_delay total expected counts
        scale_factor = self.counts_per_delay / mean_row_total
        expected_counts = signal_positive * scale_factor

        # Independent Poisson draw per pixel
        photon_counts = self.rng.poisson(expected_counts)

        # Scale back to original signal units and restore sign
        # (negative signal = bleach/emission, positive = absorption)
        noisy_data = photon_counts / scale_factor * np.sign(signal)

        return cast("np.ndarray", noisy_data)

    #
    def set_noise_level(self, noise_level: float) -> None:
        """
        Update noise level (analog detectors only).

        Parameters
        ----------
        noise_level : float
            Standard deviation of Gaussian noise (absolute units).
        """

        if self.detection != "analog":
            warnings.warn(
                "noise_level only applies to analog detection",
                stacklevel=2,
            )
        self.noise_level = noise_level

    #
    def set_noise_type(self, noise_type: str) -> None:
        """
        Update noise type (analog detectors only).

        Parameters
        ----------
        noise_type : str
            Noise distribution: ``'gaussian'`` or ``'uniform'``.
        """

        if self.detection != "analog":
            warnings.warn(
                "noise_type only applies to analog detection",
                stacklevel=2,
            )
        self.noise_type = noise_type.lower()

    #
    def set_counts_per_delay(self, counts_per_delay: int) -> None:
        """
        Update counts per delay (photon counting only).

        Parameters
        ----------
        counts_per_delay : int
            Total photon counts collected per delay step.
        """

        if self.detection != "photon_counting":
            warnings.warn(
                "counts_per_delay only applies to photon_counting detection",
                stacklevel=2,
            )
        self.counts_per_delay = counts_per_delay

    #
    def set_count_rate(
        self, count_rate: float, integration_time: float | None = None
    ) -> None:
        """
        Update count rate (photon counting only).

        Parameters
        ----------
        count_rate : float
            Photon rate in Hz.
        integration_time : float | None
            Integration time per delay in seconds.
            If None, uses existing value.
        """

        if self.detection != "photon_counting":
            warnings.warn(
                "count_rate only applies to photon_counting detection",
                stacklevel=2,
            )
            return

        self.count_rate = count_rate
        if integration_time is not None:
            self.integration_time = integration_time

        if self.integration_time is not None:
            self.counts_per_delay = int(self.count_rate * self.integration_time)
        else:
            raise ValueError(
                "integration_time must be set to calculate"
                " counts_per_delay from count_rate"
            )

    #
    def set_seed(self, seed: int | None) -> None:
        """
        Update random seed.

        Parameters
        ----------
        seed : int or None
            Random seed for reproducibility. None for non-deterministic.
        """

        self.seed = seed
        self.rng = np.random.default_rng(seed)

    #
    def get_snr(self, scale: str = "linear") -> float:
        """
        Calculate Signal-to-Noise Ratio (SNR).

        Computes SNR from the most recently simulated data using power-based
        definition: SNR = signal_power / noise_power.

        Parameters
        ----------
        scale : {'linear', 'dB'}, default='linear'
            Output scale:
            - 'linear': SNR as ratio (e.g., 25.0)
            - 'dB': SNR in decibels (e.g., 13.98 dB)

        Returns
        -------
        float
            SNR value in requested scale.
            Returns np.inf if noise_power is exactly zero.

        Raises
        ------
        ValueError
            If no simulated data available (must call simulate_1d/2d/n first),
            or if scale is not 'linear' or 'dB'.

        Examples
        --------
        >>> # Calculate SNR after simulation
        >>> sim = Simulator(model, noise_level=0.05)
        >>> clean, noisy, noise = sim.simulate_2d()
        >>>
        >>> snr_linear = sim.get_snr(scale='linear')
        >>> print(f"SNR: {snr_linear:.1f}")
        SNR: 25.3
        >>>
        >>> snr_db = sim.get_snr(scale='dB')
        >>> print(f"SNR: {snr_db:.1f} dB")
        SNR: 14.0 dB

        >>> # Compare SNR across noise levels
        >>> for noise_level in [0.01, 0.05, 0.10, 0.20]:
        ...     sim.set_noise_level(noise_level)
        ...     sim.simulate_2d()
        ...     snr = sim.get_snr()
        ...     print(f"Noise {noise_level:.2f}: SNR = {snr:.1f}")
        Noise 0.01: SNR = 625.0
        Noise 0.05: SNR = 25.0
        Noise 0.10: SNR = 6.2
        Noise 0.20: SNR = 1.6

        >>> # Plot SNR vs photon count
        >>> counts = [100, 500, 1000, 5000, 10000]
        >>> snrs = []
        >>> for count in counts:
        ...     sim = Simulator(model, detection='photon_counting',
        ...                     counts_per_delay=count)
        ...     sim.simulate_2d()
        ...     snrs.append(sim.get_snr())
        >>> plt.loglog(counts, snrs, 'o-')
        >>> plt.xlabel('Counts per delay')
        >>> plt.ylabel('SNR')

        Notes
        -----
        **SNR Definition:**

        Uses power-based (energy) definition:

        SNR_linear = (mean(signal²)) / (mean(noise²))
        SNR_dB = 10 × log₁₀(SNR_linear)

        This differs from amplitude-based definition (20 log₁₀) by factor of 2.
        Power-based is standard in signal processing and communications.

        **Interpretation:**

        Linear scale:
        - SNR = 1: Signal and noise have equal power (marginal)
        - SNR = 10: Signal 10× stronger than noise (good)
        - SNR = 100: Signal 100× stronger than noise (excellent)

        dB scale:
        - 0 dB: Equal signal and noise
        - 10 dB: 10× signal power (good)
        - 20 dB: 100× signal power (excellent)
        - Each 10 dB = 10× power ratio

        **Typical Values:**

        For spectroscopy data:
        - SNR < 5 (< 7 dB): Difficult to fit reliably
        - SNR 5-20 (7-13 dB): Good quality, typical experimental data
        - SNR 20-100 (13-20 dB): High quality
        - SNR > 100 (> 20 dB): Exceptional, near ideal

        **Limitations:**

        This is a global SNR averaged over entire spectrum. Local SNR
        may vary significantly, especially for:
        - Weak features vs. strong peaks
        - Time-dependent signals (varying amplitude)
        - Non-uniform noise (detector artifacts)

        For accurate local SNR, compute on regions of interest separately.

        See Also
        --------
        simulate_1d : Must call before get_snr
        simulate_2d : Must call before get_snr
        plot_comparison : Shows SNR in title
        """

        if self.data_clean is None or self.noise is None:
            raise ValueError(
                "No simulated data available. Run simulate_1d or simulate_2d first."
            )
        signal_power = np.mean(self.data_clean**2)
        noise_power = np.mean(self.noise**2)

        if noise_power == 0:
            return float(np.inf)

        if scale == "linear":
            return float(signal_power / noise_power)
        if scale == "dB":
            return float(10 * np.log10(signal_power / noise_power))
        raise ValueError("scale must be either 'linear' or 'dB'")

    #
    def plot_comparison(
        self,
        t_ind: int = 0,
        dim: int = 1,
        snr_scale: str = "linear",
        *,
        save_img: int = 0,
    ) -> None:
        """
        Plot comparison of clean vs noisy data.

        Creates visualization showing clean model data, noisy simulated data,
        and noise component side-by-side. Essential for visually assessing
        simulation quality and noise characteristics.

        Parameters
        ----------
        t_ind : int, default=0
            Time index for 1D plots (ignored for dim=2)
        dim : {1, 2}, default=1
            Dimensionality:
            - 1: Create 1D plot with clean, noisy, and noise curves
            - 2: Create three-panel 2D plot (clean, noisy, noise)
        snr_scale : {'linear', 'dB'}, default='linear'
            Scale for SNR display in title:
            - 'linear': Show as ratio (e.g., "SNR: 25.0 linear")
            - 'dB': Show in decibels (e.g., "SNR: 14.0 dB")
        save_img : int, default=0
            0: display, 1: save+display, -1: save only, -2: close (no display/save)

        Examples
        --------
        >>> # 1D comparison
        >>> sim = Simulator(model, noise_level=0.05)
        >>> sim.simulate_1d(t_ind=0)
        >>> sim.plot_comparison(dim=1)

        >>> # 2D comparison with dB scale
        >>> sim = Simulator(model, noise_level=0.05)
        >>> sim.simulate_2d()
        >>> sim.plot_comparison(dim=2, snr_scale='dB')

        >>> # Compare different noise levels visually
        >>> fig, axes = plt.subplots(3, 1, figsize=(10, 12))
        >>> for i, noise_level in enumerate([0.01, 0.05, 0.10]):
        ...     sim.set_noise_level(noise_level)
        ...     sim.simulate_1d()
        ...     # ... manual plotting on axes[i] ...

        >>> # Check photon counting vs analog
        >>> sim_analog = Simulator(model, detection='analog', noise_level=0.05)
        >>> sim_photon = Simulator(model, detection='photon_counting',
        ...                         counts_per_delay=1000)
        >>> sim_analog.simulate_2d()
        >>> sim_photon.simulate_2d()
        >>> # ... compare visually ...

        Notes
        -----
        **1D Plot Layout:**

        Single plot with three traces:
        - Clean: Black line (ground truth)
        - Noisy: Red scatter points (simulated data)
        - Noise: Gray line (noise component)

        Scatter points for noisy data help visualize noise granularity.

        **2D Plot Layout:**

        Three side-by-side panels:
        - Left: Clean model data
        - Center: Noisy simulated data (with SNR in title)
        - Right: Noise component (difference)

        All use same colormap from model.plot_config for consistency.

        **Visual Assessment:**

        Good simulation should show:
        - Noisy data follows clean data trend
        - Noise is randomly distributed (no patterns)
        - SNR appropriate for intended use case
        - Peak features still distinguishable in noisy data

        If noise dominates signal (SNR << 1), features may be
        completely obscured - increase signal or reduce noise.

        **Configuration:**

        Plot uses model.plot_config for:
        - Axis labels (energy/time labels)
        - Axis direction (e.g., reversed energy)
        - Colormap (for 2D plots)
        - DPI settings

        This ensures consistency with other trspecfit plots.

        See Also
        --------
        simulate_1d : Generate 1D data to plot
        simulate_2d : Generate 2D data to plot
        get_snr : SNR calculation shown in title
        """

        detection_str = f" [{self.detection}]"
        plt_title = (
            f"Simulated Data (SNR: {self.get_snr(scale=snr_scale):.1f}"
            f" {snr_scale}){detection_str}"
        )

        if dim == 1:
            if self.data_clean is None:
                self.simulate_1d(t_ind)
            if self.data_clean is None or self.data_noisy is None or self.noise is None:
                raise RuntimeError("Simulation data not available for plotting")

            # Get config from model
            config = self.model.plot_config

            # Plot with noisy data as scatter
            uplt.plot_1d(
                data=[self.data_clean, self.data_noisy, self.noise],
                x=self.model.energy,
                config=config,
                title=plt_title,
                legend=["Clean", "Noisy", "Noise"],
                linestyles=["-", "", "-"],  # Empty string = no line for noisy data
                markers=[None, "o", None],  # Scatter points for noisy data
                markersizes=[6, 3, 6],  # Smaller markers for noisy data
                save_img=save_img,
            )

        elif dim == 2:
            if self.data_clean is None:
                self.simulate_2d()
            if self.data_clean is None or self.data_noisy is None or self.noise is None:
                raise RuntimeError("Simulation data not available for plotting")

            # Get config from model
            config = self.model.plot_config

            # Create 3-panel plot
            _fig, axes = plt.subplots(1, 3, figsize=(15, 4))

            panels = [
                (self.data_clean, "Clean Model Data"),
                (self.data_noisy, plt_title),
                (self.noise, "Noise (Simulated - Clean)"),
            ]
            for ax, (data, title) in zip(axes, panels, strict=True):
                im = ax.pcolormesh(
                    self.model.energy,
                    self.model.time,
                    data,
                    shading="nearest",
                    cmap=config.z_colormap,
                )
                ax.set_title(title)
                ax.set_xlabel(config.x_label)
                ax.set_ylabel(config.y_label)
                if config.ticksize is not None:
                    ax.tick_params(labelsize=config.ticksize)
                uplt._apply_axis_settings(
                    ax,
                    x_type=config.x_type,
                    x_dir=config.x_dir,
                    y_type=config.y_type,
                    y_dir=config.y_dir,
                    x_lim=config.x_lim,
                    y_lim=config.y_lim,
                )
                plt.colorbar(im, ax=ax)

            plt.tight_layout()
            uplt._finalize_plot(
                save_img=save_img,
                save_path="",
                dpi_save=config.dpi_save,
            )

    #
    def save_data(
        self,
        *,
        filepath: str | None = None,
        save_format: str = "hdf5",
        n_data: list[np.ndarray] | None = None,
        overwrite: bool = True,
        show_output: int = 1,
    ) -> None:
        """
        Save simulated data to file with metadata.

        Exports simulated data in HDF5 format with complete metadata including
        model parameters, noise settings, and experimental axes. Essential for
        sharing simulated datasets and ensuring reproducibility.

        Parameters
        ----------
        filepath : str or Path, optional
            Path where to save data. If None, uses default:
            './simulated_data/simulated_data.h5'
            If provided path doesn't include 'simulated_data' directory,
            it will be automatically placed there.
        save_format : str, default='hdf5'
            File format. Currently only 'hdf5' supported.
            Future: could add .mat, .npz, etc.
        n_data : list of ndarray, optional
            Multiple noisy datasets from simulate_n() to save.
            If None, saves single dataset from simulate_1d() or simulate_2d().
        overwrite : bool, default=True
            If True, overwrite existing files.
            If False, raise FileExistsError if file exists.
        show_output : int, default=1
            Output mode:

            - 0: Silent / programmatic / API mode -- no prints
            - 1: Interactive / notebook / UI mode -- show timing and
              save confirmation

        Raises
        ------
        ValueError
            If no simulated data available (must call simulate first)
        FileExistsError
            If file exists and overwrite=False

        Examples
        --------
        >>> # Save single simulation
        >>> sim = Simulator(model, noise_level=0.05, seed=42)
        >>> clean, noisy, noise = sim.simulate_2d()
        >>> sim.save_data('simulation_001.h5')
        Data saved to: ./simulated_data/simulation_001.h5

        >>> # Save multiple realizations
        >>> clean, noisy_list, noise_list = sim.simulate_n(n=50, dim=2)
        >>> sim.save_data(
        ...     filepath='batch_simulation.h5',
        ...     n_data=noisy_list
        ... )
        Data saved to: ./simulated_data/batch_simulation.h5

        >>> # Prevent accidental overwrites
        >>> sim.save_data('important_data.h5', overwrite=False)
        FileExistsError: File already exists: ./simulated_data/important_data.h5
        Set overwrite=True to overwrite, or provide a different filepath.

        >>> # Load saved data later
        >>> import h5py
        >>> with h5py.File('simulated_data/simulation_001.h5', 'r') as f:
        ...     energy = f['energy'][:]
        ...     time = f['time'][:]
        ...     clean = f['clean_data'][:]
        ...     noisy = f['simulated_data/000000'][:]
        ...
        ...     # Read metadata
        ...     noise_level = f['metadata'].attrs['noise_level']
        ...     model_params = f['metadata'].attrs['model_parameters']

        Notes
        -----
        **HDF5 File Structure:**

        ::

            /
            ├── energy              (dataset: 1D array)
            ├── time                (dataset: 1D array, empty for 1D simulations)
            ├── clean_data          (dataset: 1D or 2D array)
            ├── simulated_data/     (group)
            │   ├── 000000          (dataset: first noisy realization)
            │   ├── 000001          (dataset: second noisy realization)
            │   └── ...
            └── metadata/           (group with [optional]attributes)
                ├── detection       ('analog' or 'photon_counting')
                ├── noise_level     (analog noise level)
                ├── noise_type      (analog noise type)
                ├── counts_per_delay (photon counting counts)
                ├── count_rate      ([optional] photon counting rate)
                ├── integration_time      ([optional] photon counting integration time)
                ├── seed            ([optional] random seed, if set)
                ├── dimension       (1 or 2)
                ├── n_datasets      (number of noisy datasets)
                ├── model_parameters (JSON string of all parameters)
                └── model_name      (model name)

        **Why HDF5?**

        HDF5 format chosen because:
        - Efficient for large multidimensional arrays
        - Self-describing (metadata embedded)
        - Widely supported (Python, MATLAB, Igor, etc.)
        - Allows partial loading (don't need entire file in memory)
        - Standard in scientific computing

        **Model Parameters:**

        All model parameters saved as JSON string in metadata for complete
        reproducibility. Includes:

        - Parameter values
        - vary flags (which parameters were free)
        - Bounds (min/max)
        - Expressions (parameter constraints)

        This allows exact recreation of the model used for simulation.

        **File Organization:**

        Default directory structure::

            project_directory/
            └── simulated_data/
                ├── simulation_001.h5
                ├── simulation_002.h5
                └── batch_001.h5

        Keeps simulated data organized and separate from experimental data.

        **Multiple Datasets:**

        When n_data provided (from simulate_n), all realizations saved in
        simulated_data group with sequential names:

        - 000000, 000001, ..., 000099 for 100 datasets
        - Zero-padded for proper sorting

        Clean data saved once (same for all realizations).

        **Loading Data:**

        Standard h5py usage::

            import h5py

            with h5py.File('simulated_data/data.h5', 'r') as f:
                # Load axes
                energy = f['energy'][:]
                time = f['time'][:]

                # Load clean data
                clean = f['clean_data'][:]

                # Load all noisy datasets
                noisy_datasets = []
                for key in sorted(f['simulated_data'].keys()):
                    noisy_datasets.append(f['simulated_data'][key][:])

                # Load metadata
                detection = f['metadata'].attrs['detection']
                n_datasets = f['metadata'].attrs['n_datasets']

        See Also
        --------
        simulate_n : Generate multiple datasets to save
        simulate_1d : Generate 1D data
        simulate_2d : Generate 2D data
        h5py : Python HDF5 library
        """

        if self.data_noisy is None and n_data is None:
            raise ValueError(
                "No simulated data available."
                " Run simulate_1d, simulate_2d, or simulate_n first."
            )

        # Create simulated_data directory if it doesn't exist
        sim_dir = Path.cwd() / "simulated_data"
        sim_dir.mkdir(parents=True, exist_ok=True)

        # Set default filepath
        if filepath is None:
            if save_format == "hdf5":
                filepath = str(sim_dir / "simulated_data.h5")
        else:
            # User provided filepath - make sure it's in simulated_data directory
            if not Path(filepath).is_relative_to(sim_dir):
                filepath = str(sim_dir / Path(filepath).name)

        # Check if file exists and handle overwrite
        if filepath is None:
            raise ValueError("filepath could not be resolved")
        if Path(filepath).exists() and not overwrite:
            raise FileExistsError(
                f"File already exists: {filepath}\n"
                f"Set overwrite=True to overwrite, or provide a different filepath."
            )

        if save_format == "hdf5":
            self._save_hdf5(filepath, n_data)

        # add other formats here

        else:
            raise ValueError(f"Unknown save format: {save_format}. Use 'hdf5'.")

        if show_output >= 1:
            print(f"Data saved to: {filepath}")

    #
    def _save_hdf5(self, filepath: str, n_data: list[np.ndarray] | None = None) -> None:
        """
        Save data to HDF5 format with proper structure

        HDF5 structure:
            /energy         - energy axis array
            /time           - time axis array (empty if 1D)
            /clean_data     - clean data without noise
            /simulated_data/000000 - first noisy dataset
            /simulated_data/000001 - second noisy dataset
            ...
            /metadata       - group containing simulation parameters
        """

        with h5py.File(filepath, "w") as f:
            # Save axes at root level
            f.create_dataset("energy", data=self.model.energy)

            # Handle time axis - check if it exists and has valid data
            if self.model.time is not None and len(self.model.time) > 0:
                f.create_dataset("time", data=self.model.time)
            else:
                # For 1D simulations, save empty array
                f.create_dataset("time", data=np.array([]))

            # Save clean data at root level
            if self.data_clean is not None:
                f.create_dataset("clean_data", data=self.data_clean)

            # Create simulated_data group for noisy datasets
            sim_group = f.create_group("simulated_data")

            # Save noisy data - either from n_data list or single simulation
            if n_data is not None:
                # Multiple datasets from simulate_n
                for i, noisy_data in enumerate(n_data):
                    dataset_name = f"{i:06d}"
                    sim_group.create_dataset(dataset_name, data=noisy_data)
            elif self.data_noisy is not None:
                # Single dataset
                sim_group.create_dataset("000000", data=self.data_noisy)

            # Save metadata group at root level
            meta = f.create_group("metadata")
            meta.attrs["detection"] = self.detection

            # Save detection-specific parameters
            if self.detection == "analog":
                meta.attrs["noise_level"] = self.noise_level
                meta.attrs["noise_type"] = self.noise_type
            elif self.detection == "photon_counting":
                meta.attrs["counts_per_delay"] = self.counts_per_delay
                if self.count_rate is not None:
                    meta.attrs["count_rate"] = self.count_rate
                if self.integration_time is not None:
                    meta.attrs["integration_time"] = self.integration_time

            if self.seed is not None:
                meta.attrs["seed"] = self.seed

            # Determine dimensionality from clean data
            if self.data_clean is not None:
                if self.data_clean.ndim == 1:
                    meta.attrs["dimension"] = 1
                else:
                    meta.attrs["dimension"] = 2

            # Save number of datasets
            if n_data is not None:
                meta.attrs["n_datasets"] = len(n_data)
            else:
                meta.attrs["n_datasets"] = 1

            # Save model parameters as JSON in metadata
            params_dict = {}
            for par_name in self.model.lmfit_pars:
                par = self.model.lmfit_pars[par_name]
                params_dict[par_name] = {
                    "value": float(par.value),
                    "vary": bool(par.vary),
                    "min": float(par.min) if par.min is not None else None,
                    "max": float(par.max) if par.max is not None else None,
                    "expr": par.expr or None,
                }

            # Save as JSON string in metadata
            meta.attrs["model_parameters"] = json.dumps(params_dict, indent=2)
            meta.attrs["model_name"] = self.model.name

    #
    def simulate_parameter_sweep(
        self,
        parameter_sweep: ParameterSweep,
        n_realizations: int,
        *,
        dim: int = 2,
        filepath: str = "ml_training_data.h5",
        show_progress: bool = True,
    ) -> None:
        """
        Generate ML training dataset by sweeping parameters.

        Processes configurations one at a time, immediately saving to disk.
        Memory usage remains constant regardless of parameter space size.

        Parameters
        ----------
        parameter_sweep : ParameterSweep
            Generator yielding parameter configurations
        n_realizations : int
            Number of noisy realizations per parameter configuration
        dim : {1, 2}, default=2
            Dimensionality of simulated data:
            - 1: Generate 1D spectra
            - 2: Generate 2D spectra
        filepath : str, default='ml_training_data.h5'
            HDF5 file path for output
        show_progress : bool, default=True
            Print progress updates during generation

        Examples
        --------
        >>> # Set up parameter space
        >>> sweep = ParameterSweep(strategy='random', seed=42)
        >>> sweep.add_uniform('GLP_01_A', 5, 30, n_samples=100)
        >>> sweep.add_uniform('GLP_01_x0', 5, 15, n_samples=100)
        >>>
        >>> # Generate dataset
        >>> sim = Simulator(model, noise_level=0.05, seed=42)
        >>> sim.simulate_parameter_sweep(
        ...     parameter_sweep=sweep,
        ...     n_realizations=20,
        ...     filepath='training_data.h5'
        ... )
        Processing config 1/100: {'GLP_01_A': 12.5, 'GLP_01_x0': 8.3}
          Saved config 1 with 20 realizations
        ...
        Parameter sweep complete!
        Generated 100 configs × 20 realizations
        Data saved to: ./simulated_data/training_data.h5

        Notes
        -----
        **Memory Efficiency:**
        Only one configuration is in memory at a time. Each is immediately
        written to disk before processing the next. Total memory usage is
        independent of parameter space size.

        **Resumability:**
        If interrupted, completed configurations are already saved to disk.
        Currently does not support automatic resume (will overwrite file).

        **File Structure:**
        See _initialize_sweep_hdf5 for complete HDF5 structure description.

        See Also
        --------
        ParameterSweep : Define parameter space to sweep
        simulate_n : Generate multiple noisy realizations
        _initialize_sweep_hdf5 : HDF5 file structure
        _append_config_to_hdf5 : Incremental saving logic
        """

        # Convert to Path object
        filepath_obj = Path(filepath)

        # If filepath is just a filename (no directory component),
        # put it in simulated_data subdirectory
        if filepath_obj.parent == Path():
            sim_dir = Path.cwd() / "simulated_data"
            sim_dir.mkdir(parents=True, exist_ok=True)
            filepath = str(sim_dir / filepath_obj.name)
        else:
            # User provided a path with directory, use it as-is but ensure parent exists
            filepath_obj.parent.mkdir(parents=True, exist_ok=True)
            filepath = str(filepath_obj)

        # Get total number of configurations
        n_configs = parameter_sweep.get_n_configs()

        if show_progress:
            print(
                f"Starting parameter sweep:\n"
                f"  Total configurations: {n_configs}\n"
                f"  Realizations per config: {n_realizations}\n"
                f"  Total datasets: {n_configs * n_realizations}\n"
                f"  Output file: {filepath}\n"
            )

        # Initialize HDF5 file with structure
        self._initialize_sweep_hdf5(
            filepath, parameter_sweep, n_realizations, n_configs, dim
        )

        # Process each configuration
        for config_idx, param_config in enumerate(parameter_sweep):
            if show_progress:
                # Format parameters nicely
                param_str = ", ".join(f"{k}={v:.3g}" for k, v in param_config.items())
                print(
                    f"Processing config {config_idx + 1}/{n_configs}: {{{param_str}}}"
                )

            # Update model parameters
            param_names = list(param_config.keys())
            param_values = list(param_config.values())
            self.model.update_value(param_values, par_select=param_names)

            # Generate noisy realizations for this config
            clean, noisy_list, _noise_list = self.simulate_n(
                n=n_realizations,
                dim=dim,
                show_progress=False,  # Don't clutter output
            )

            # Append to HDF5 immediately (memory-efficient)
            self._append_config_to_hdf5(
                filepath, config_idx, param_config, clean, noisy_list
            )

            if show_progress:
                print(
                    f"  ✓ Saved config {config_idx + 1}"
                    f" with {n_realizations} realizations"
                )

        if show_progress:
            print(
                f"\n{'=' * 60}\n"
                f"Parameter sweep complete!\n"
                f"Generated {n_configs} configs × {n_realizations} realizations\n"
                f"Total datasets: {n_configs * n_realizations}\n"
                f"Data saved to: {filepath}\n"
                f"{'=' * 60}"
            )

    #
    def _initialize_sweep_hdf5(
        self,
        filepath: str,
        parameter_sweep: ParameterSweep,
        n_realizations: int,
        n_configs: int,
        dim: int,
    ) -> None:
        """
        Create HDF5 file structure for parameter sweep data.

        File Structure
        --------------
        /
        ├── energy (dataset)              # Energy axis
        ├── time (dataset)                # Time axis (empty for 1D)
        ├── metadata/ (group)             # Sweep metadata
        │   ├── attrs: n_configs, n_realizations_per_config, ...
        │   └── attrs: parameter_space (JSON)
        ├── parameter_configs/ (group)    # Parameter configurations
        │   ├── config_000000/ (group)
        │   │   ├── attrs: GLP_01_A, GLP_01_x0, ...
        │   │   ├── attrs: all_parameter_values (JSON)
        │   │   └── clean (dataset)       # Clean data for this config
        │   ├── config_000001/ (group)
        │   └── ...
        └── simulated_data/ (group)       # Noisy realizations
            ├── config_000000/ (group)
            │   ├── 000000 (dataset)      # First realization
            │   ├── 000001 (dataset)      # Second realization
            │   └── ...
            ├── config_000001/ (group)
            └── ...

        Parameters
        ----------
        filepath : str
            Path to HDF5 file to create
        parameter_sweep : ParameterSweep
            Parameter sweep object (for metadata)
        n_realizations : int
            Number of noisy realizations per config
        n_configs : int
            Total number of parameter configurations
        dim : {1, 2}
            Dimensionality of simulated data
        """

        with h5py.File(filepath, "w") as f:
            # Save axes (same for all configs)
            f.create_dataset("energy", data=self.model.energy)
            if self.model.time is not None and len(self.model.time) > 0:
                f.create_dataset("time", data=self.model.time)
            else:
                f.create_dataset("time", data=np.array([]))

            # Create groups for organization
            f.create_group("parameter_configs")
            f.create_group("simulated_data")

            # Save sweep metadata
            meta = f.create_group("metadata")
            meta.attrs["n_configs"] = n_configs
            meta.attrs["n_realizations_per_config"] = n_realizations
            meta.attrs["total_datasets"] = n_configs * n_realizations

            # Simulator settings
            meta.attrs["detection"] = self.detection
            if self.detection == "analog":
                meta.attrs["noise_level"] = self.noise_level
                meta.attrs["noise_type"] = self.noise_type
            elif self.detection == "photon_counting":
                meta.attrs["counts_per_delay"] = self.counts_per_delay
                if self.count_rate is not None:
                    meta.attrs["count_rate"] = self.count_rate
                if self.integration_time is not None:
                    meta.attrs["integration_time"] = self.integration_time

            if self.seed is not None:
                meta.attrs["seed"] = self.seed

            # Parameter sweep settings
            meta.attrs["sweep_strategy"] = parameter_sweep.strategy
            meta.attrs["sweep_seed"] = (
                "None" if parameter_sweep.seed is None else parameter_sweep.seed
            )

            # Save parameter space definition as JSON
            param_space = {}
            for par_name, spec in parameter_sweep.parameter_specs.items():
                # Convert numpy arrays to lists for JSON serialization
                spec_copy = spec.copy()
                if "values" in spec_copy:
                    spec_copy["values"] = spec_copy["values"].tolist()
                param_space[par_name] = spec_copy

            meta.attrs["parameter_space"] = json.dumps(param_space, indent=2)

            # Save full model definition once (static across all configs)
            model_params = {}
            for par_name in self.model.lmfit_pars:
                par = self.model.lmfit_pars[par_name]
                model_params[par_name] = {
                    "value": float(par.value),
                    "vary": bool(par.vary),
                    "min": float(par.min) if par.min is not None else None,
                    "max": float(par.max) if par.max is not None else None,
                    "expr": par.expr or None,
                }
            meta.attrs["model_parameters"] = json.dumps(model_params, indent=2)
            meta.attrs["model_name"] = self.model.name

            # Dimension matches the actual data written, not model capability
            meta.attrs["dimension"] = dim

    #
    def _append_config_to_hdf5(
        self,
        filepath: str,
        config_idx: int,
        param_config: dict[str, float],
        clean: np.ndarray,
        noisy_list: list[np.ndarray],
    ) -> None:
        """
        Append single parameter configuration and its realizations to HDF5.

        Parameters
        ----------
        filepath : str
            Path to HDF5 file
        config_idx : int
            Configuration index (for naming)
        param_config : dict
            Parameter values for this configuration
        clean : ndarray
            Clean (noiseless) data
        noisy_list : list of ndarray
            List of noisy realizations
        """

        with h5py.File(filepath, "a") as f:
            # Create group for this configuration
            config_name = f"config_{config_idx:06d}"
            configs_group = require_group(f["parameter_configs"], "parameter_configs")
            config_group = configs_group.create_group(config_name)

            # Save swept parameters as attributes
            for par_name, value in param_config.items():
                config_group.attrs[par_name] = float(value)

            # Save all parameter values for this config (full model state)
            param_values = {
                par_name: float(self.model.lmfit_pars[par_name].value)
                for par_name in self.model.lmfit_pars
            }
            config_group.attrs["all_parameter_values"] = json.dumps(param_values)

            # Save clean data for this configuration
            config_group.create_dataset("clean", data=clean)

            # Save noisy realizations
            simulated_group = require_group(f["simulated_data"], "simulated_data")
            data_group = simulated_group.create_group(config_name)
            for real_idx, noisy_data in enumerate(noisy_list):
                data_group.create_dataset(f"{real_idx:06d}", data=noisy_data)
