deconvolution 0.2.1

use image::DynamicImage;
use ndarray::Array3;

use super::rl::{PoissonEm, PoissonRegularization, run_poisson_em, run_poisson_em_array3};
use crate::{ChannelMode, Error, Kernel2D, Kernel3D, RangePolicy, Result, SolveReport};

#[derive(Debug, Clone, PartialEq)]
/// Configuration for classical maximum-likelihood estimation.
pub struct Cmle {
    iterations: usize,
    relative_update_tolerance: Option<f32>,
    filter_epsilon: f32,
    snr: f32,
    acuity: f32,
    channel_mode: ChannelMode,
    range_policy: RangePolicy,
    collect_history: bool,
}

impl Default for Cmle {
    fn default() -> Self {
        Self {
            iterations: 30,
            relative_update_tolerance: None,
            filter_epsilon: 1e-6,
            snr: 40.0,
            acuity: 1.0,
            channel_mode: ChannelMode::Independent,
            range_policy: RangePolicy::PreserveInput,
            collect_history: true,
        }
    }
}

impl Cmle {
    /// Create a CMLE config with default SNR and acuity settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the maximum EM iteration count.
    pub fn iterations(mut self, value: usize) -> Self {
        self.iterations = value;
        self
    }

    /// Set the relative update stopping tolerance.
    ///
    /// `None` disables this stopping criterion.
    pub fn relative_update_tolerance(mut self, value: Option<f32>) -> Self {
        self.relative_update_tolerance = value;
        self
    }

    /// Set the positive denominator floor used in multiplicative updates.
    pub fn filter_epsilon(mut self, value: f32) -> Self {
        self.filter_epsilon = value;
        self
    }

    /// Set the positive signal-to-noise ratio used to derive readout noise.
    pub fn snr(mut self, value: f32) -> Self {
        self.snr = value;
        self
    }

    /// Set the positive acuity factor used to derive damping.
    pub fn acuity(mut self, value: f32) -> Self {
        self.acuity = value;
        self
    }

    /// Set how `image::DynamicImage` channels are restored.
    pub fn channel_mode(mut self, value: ChannelMode) -> Self {
        self.channel_mode = value;
        self
    }

    /// Set output range handling after restoration.
    pub fn range_policy(mut self, value: RangePolicy) -> Self {
        self.range_policy = value;
        self
    }

    /// Enable or disable objective and residual history in [`SolveReport`].
    pub fn collect_history(mut self, value: bool) -> Self {
        self.collect_history = value;
        self
    }
}

#[derive(Debug, Clone, PartialEq)]
/// Configuration for Gaussian maximum-likelihood estimation.
pub struct Gmle {
    iterations: usize,
    relative_update_tolerance: Option<f32>,
    filter_epsilon: f32,
    snr: f32,
    acuity: f32,
    roughness: f32,
    tv_epsilon: f32,
    channel_mode: ChannelMode,
    range_policy: RangePolicy,
    collect_history: bool,
}

impl Default for Gmle {
    fn default() -> Self {
        Self {
            iterations: 20,
            relative_update_tolerance: Some(1e-4),
            filter_epsilon: 1e-6,
            snr: 24.0,
            acuity: 0.85,
            roughness: 1.0,
            tv_epsilon: 1e-3,
            channel_mode: ChannelMode::Independent,
            range_policy: RangePolicy::PreserveInput,
            collect_history: true,
        }
    }
}

impl Gmle {
    /// Create a GMLE config with default noise and TV settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the maximum EM iteration count.
    pub fn iterations(mut self, value: usize) -> Self {
        self.iterations = value;
        self
    }

    /// Set the relative update stopping tolerance.
    ///
    /// `None` disables this stopping criterion.
    pub fn relative_update_tolerance(mut self, value: Option<f32>) -> Self {
        self.relative_update_tolerance = value;
        self
    }

    /// Set the positive denominator floor used in multiplicative updates.
    pub fn filter_epsilon(mut self, value: f32) -> Self {
        self.filter_epsilon = value;
        self
    }

    /// Set the positive signal-to-noise ratio used to derive readout noise.
    pub fn snr(mut self, value: f32) -> Self {
        self.snr = value;
        self
    }

    /// Set the positive acuity factor used to derive damping.
    pub fn acuity(mut self, value: f32) -> Self {
        self.acuity = value;
        self
    }

    /// Set the nonnegative roughness factor used to derive TV weight.
    pub fn roughness(mut self, value: f32) -> Self {
        self.roughness = value;
        self
    }

    /// Set the positive epsilon used in TV gradient magnitude.
    pub fn tv_epsilon(mut self, value: f32) -> Self {
        self.tv_epsilon = value;
        self
    }

    /// Set how `image::DynamicImage` channels are restored.
    pub fn channel_mode(mut self, value: ChannelMode) -> Self {
        self.channel_mode = value;
        self
    }

    /// Set output range handling after restoration.
    pub fn range_policy(mut self, value: RangePolicy) -> Self {
        self.range_policy = value;
        self
    }

    /// Enable or disable objective and residual history in [`SolveReport`].
    pub fn collect_history(mut self, value: bool) -> Self {
        self.collect_history = value;
        self
    }
}

#[derive(Debug, Clone, PartialEq)]
/// Configuration for quadratic maximum-likelihood estimation.
pub struct Qmle {
    iterations: usize,
    relative_update_tolerance: Option<f32>,
    filter_epsilon: f32,
    snr: f32,
    acuity: f32,
    channel_mode: ChannelMode,
    range_policy: RangePolicy,
    collect_history: bool,
}

impl Default for Qmle {
    fn default() -> Self {
        Self {
            iterations: 10,
            relative_update_tolerance: Some(2e-4),
            filter_epsilon: 1e-6,
            snr: 60.0,
            acuity: 1.1,
            channel_mode: ChannelMode::Independent,
            range_policy: RangePolicy::PreserveInput,
            collect_history: true,
        }
    }
}

impl Qmle {
    /// Create a QMLE config with default SNR and acuity settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the maximum EM iteration count.
    pub fn iterations(mut self, value: usize) -> Self {
        self.iterations = value;
        self
    }

    /// Set the relative update stopping tolerance.
    ///
    /// `None` disables this stopping criterion.
    pub fn relative_update_tolerance(mut self, value: Option<f32>) -> Self {
        self.relative_update_tolerance = value;
        self
    }

    /// Set the positive denominator floor used in multiplicative updates.
    pub fn filter_epsilon(mut self, value: f32) -> Self {
        self.filter_epsilon = value;
        self
    }

    /// Set the positive signal-to-noise ratio used to derive readout noise.
    pub fn snr(mut self, value: f32) -> Self {
        self.snr = value;
        self
    }

    /// Set the positive acuity factor used to derive damping.
    pub fn acuity(mut self, value: f32) -> Self {
        self.acuity = value;
        self
    }

    /// Set how `image::DynamicImage` channels are restored.
    pub fn channel_mode(mut self, value: ChannelMode) -> Self {
        self.channel_mode = value;
        self
    }

    /// Set output range handling after restoration.
    pub fn range_policy(mut self, value: RangePolicy) -> Self {
        self.range_policy = value;
        self
    }

    /// Enable or disable objective and residual history in [`SolveReport`].
    pub fn collect_history(mut self, value: bool) -> Self {
        self.collect_history = value;
        self
    }
}

/// Restore an image with classical maximum-likelihood estimation.
///
/// # Errors
///
/// Returns an error for invalid PSFs, empty or non-finite image data, invalid
/// SNR, acuity, or EM settings, or non-finite iterative updates.
pub fn cmle(image: &DynamicImage, psf: &Kernel2D) -> Result<(DynamicImage, SolveReport)> {
    cmle_with(image, psf, &Cmle::new())
}

/// Restore an image with CMLE and explicit settings.
///
/// # Errors
///
/// Returns an error for invalid PSFs, empty or non-finite image data, invalid
/// SNR, acuity, or EM settings, or non-finite iterative updates.
pub fn cmle_with(
    image: &DynamicImage,
    psf: &Kernel2D,
    config: &Cmle,
) -> Result<(DynamicImage, SolveReport)> {
    validate_cmle(config)?;
    let poisson = cmle_poisson_em(config)?;
    run_poisson_em(image, psf, &poisson, PoissonRegularization::None)
}

pub(crate) fn cmle_array3_with(
    volume: &Array3<f32>,
    psf: &Kernel3D,
    config: &Cmle,
) -> Result<(Array3<f32>, SolveReport)> {
    validate_cmle(config)?;
    let poisson = cmle_poisson_em(config)?;
    run_poisson_em_array3(volume, psf, &poisson, PoissonRegularization::None)
}

/// Restore an image with Gaussian maximum-likelihood estimation.
///
/// # Errors
///
/// Returns an error for invalid PSFs, empty or non-finite image data, invalid
/// SNR, acuity, roughness, TV, or EM settings, or non-finite iterative updates.
pub fn gmle(image: &DynamicImage, psf: &Kernel2D) -> Result<(DynamicImage, SolveReport)> {
    gmle_with(image, psf, &Gmle::new())
}

/// Restore an image with GMLE and explicit settings.
///
/// # Errors
///
/// Returns an error for invalid PSFs, empty or non-finite image data, invalid
/// SNR, acuity, roughness, TV, or EM settings, or non-finite iterative updates.
pub fn gmle_with(
    image: &DynamicImage,
    psf: &Kernel2D,
    config: &Gmle,
) -> Result<(DynamicImage, SolveReport)> {
    validate_gmle(config)?;
    let (poisson, regularization) = gmle_poisson_em(config)?;
    run_poisson_em(image, psf, &poisson, regularization)
}

pub(crate) fn gmle_array3_with(
    volume: &Array3<f32>,
    psf: &Kernel3D,
    config: &Gmle,
) -> Result<(Array3<f32>, SolveReport)> {
    validate_gmle(config)?;
    let (poisson, regularization) = gmle_poisson_em(config)?;
    run_poisson_em_array3(volume, psf, &poisson, regularization)
}

/// Restore an image with quadratic maximum-likelihood estimation.
///
/// # Errors
///
/// Returns an error for invalid PSFs, empty or non-finite image data, invalid
/// SNR, acuity, or EM settings, or non-finite iterative updates.
pub fn qmle(image: &DynamicImage, psf: &Kernel2D) -> Result<(DynamicImage, SolveReport)> {
    qmle_with(image, psf, &Qmle::new())
}

/// Restore an image with QMLE and explicit settings.
///
/// # Errors
///
/// Returns an error for invalid PSFs, empty or non-finite image data, invalid
/// SNR, acuity, or EM settings, or non-finite iterative updates.
pub fn qmle_with(
    image: &DynamicImage,
    psf: &Kernel2D,
    config: &Qmle,
) -> Result<(DynamicImage, SolveReport)> {
    validate_qmle(config)?;
    let poisson = qmle_poisson_em(config)?;
    run_poisson_em(image, psf, &poisson, PoissonRegularization::None)
}

pub(crate) fn qmle_array3_with(
    volume: &Array3<f32>,
    psf: &Kernel3D,
    config: &Qmle,
) -> Result<(Array3<f32>, SolveReport)> {
    validate_qmle(config)?;
    let poisson = qmle_poisson_em(config)?;
    run_poisson_em_array3(volume, psf, &poisson, PoissonRegularization::None)
}

fn cmle_poisson_em(config: &Cmle) -> Result<PoissonEm> {
    Ok(PoissonEm {
        iterations: config.iterations,
        relative_update_tolerance: config.relative_update_tolerance,
        filter_epsilon: config.filter_epsilon,
        damping: acuity_to_damping(config.acuity)?,
        weights: None,
        readout_noise: snr_to_readout_noise(config.snr)?,
        positivity: true,
        channel_mode: config.channel_mode,
        range_policy: config.range_policy,
        collect_history: config.collect_history,
    })
}

fn gmle_poisson_em(config: &Gmle) -> Result<(PoissonEm, PoissonRegularization)> {
    let snr_noise = snr_to_readout_noise(config.snr)?;
    let noise_damping = (2.0 / config.snr.max(1.0)).sqrt().clamp(0.0, 1.0);
    let damping = acuity_to_damping(config.acuity)?
        .unwrap_or(0.0)
        .max(noise_damping);
    let tv_weight = ((config.roughness / config.snr.max(1.0)) * 0.8).clamp(0.0, 0.15);
    Ok((
        PoissonEm {
            iterations: config.iterations,
            relative_update_tolerance: config.relative_update_tolerance,
            filter_epsilon: config.filter_epsilon,
            damping: Some(damping),
            weights: None,
            readout_noise: snr_noise,
            positivity: true,
            channel_mode: config.channel_mode,
            range_policy: config.range_policy,
            collect_history: config.collect_history,
        },
        PoissonRegularization::Tv {
            weight: tv_weight,
            epsilon: config.tv_epsilon,
        },
    ))
}

fn qmle_poisson_em(config: &Qmle) -> Result<PoissonEm> {
    Ok(PoissonEm {
        iterations: config.iterations,
        relative_update_tolerance: config.relative_update_tolerance,
        filter_epsilon: config.filter_epsilon,
        damping: acuity_to_damping(config.acuity)?,
        weights: None,
        readout_noise: snr_to_readout_noise(config.snr)? * 0.5,
        positivity: true,
        channel_mode: config.channel_mode,
        range_policy: config.range_policy,
        collect_history: config.collect_history,
    })
}

fn snr_to_readout_noise(snr: f32) -> Result<f32> {
    if !snr.is_finite() || snr <= 0.0 {
        return Err(Error::InvalidParameter);
    }
    let noise = (1.0 / snr).clamp(0.0, 0.25);
    if !noise.is_finite() {
        return Err(Error::NonFiniteInput);
    }
    Ok(noise)
}

fn acuity_to_damping(acuity: f32) -> Result<Option<f32>> {
    if !acuity.is_finite() || acuity <= 0.0 {
        return Err(Error::InvalidParameter);
    }
    if acuity >= 1.0 {
        return Ok(None);
    }
    let damping = (1.0 / acuity - 1.0).clamp(0.0, 1.0);
    if !damping.is_finite() {
        return Err(Error::NonFiniteInput);
    }
    Ok(Some(damping))
}

fn validate_cmle(config: &Cmle) -> Result<()> {
    if config.iterations == 0 {
        return Err(Error::InvalidParameter);
    }
    if let Some(tol) = config.relative_update_tolerance
        && (!tol.is_finite() || tol < 0.0)
    {
        return Err(Error::InvalidParameter);
    }
    if !config.filter_epsilon.is_finite() || config.filter_epsilon <= 0.0 {
        return Err(Error::InvalidParameter);
    }
    snr_to_readout_noise(config.snr)?;
    let _ = acuity_to_damping(config.acuity)?;
    Ok(())
}

fn validate_gmle(config: &Gmle) -> Result<()> {
    if config.iterations == 0 {
        return Err(Error::InvalidParameter);
    }
    if let Some(tol) = config.relative_update_tolerance
        && (!tol.is_finite() || tol < 0.0)
    {
        return Err(Error::InvalidParameter);
    }
    if !config.filter_epsilon.is_finite() || config.filter_epsilon <= 0.0 {
        return Err(Error::InvalidParameter);
    }
    if !config.roughness.is_finite() || config.roughness < 0.0 {
        return Err(Error::InvalidParameter);
    }
    if !config.tv_epsilon.is_finite() || config.tv_epsilon <= 0.0 {
        return Err(Error::InvalidParameter);
    }
    snr_to_readout_noise(config.snr)?;
    let _ = acuity_to_damping(config.acuity)?;
    Ok(())
}

fn validate_qmle(config: &Qmle) -> Result<()> {
    if config.iterations == 0 {
        return Err(Error::InvalidParameter);
    }
    if let Some(tol) = config.relative_update_tolerance
        && (!tol.is_finite() || tol < 0.0)
    {
        return Err(Error::InvalidParameter);
    }
    if !config.filter_epsilon.is_finite() || config.filter_epsilon <= 0.0 {
        return Err(Error::InvalidParameter);
    }
    snr_to_readout_noise(config.snr)?;
    let _ = acuity_to_damping(config.acuity)?;
    Ok(())
}