/*!
Solver for the point source localisation problem using a simplified forward-backward splitting method.

Instead of the $𝒟$-norm of `fb.rs`, this uses a standard Radon norm for the proximal map.
*/

use super::{InsertionConfig, ProxPenalty, ProxTerm, StepLengthBound, StepLengthBoundPD};
use crate::dataterm::QuadraticDataTerm;
use crate::forward_model::ForwardModel;
use crate::measures::merging::SpikeMerging;
use crate::measures::{DiscreteMeasure, Radon};
use crate::regularisation::RadonSquaredRegTerm;
use crate::types::*;
use alg_tools::bounds::MinMaxMapping;
use alg_tools::error::DynResult;
use alg_tools::instance::{Instance, Space};
use alg_tools::iterate::{AlgIterator, AlgIteratorIteration};
use alg_tools::linops::BoundedLinear;
use alg_tools::nalgebra_support::ToNalgebraRealField;
use alg_tools::norms::{Norm, L2};
use numeric_literals::replace_float_literals;
use serde::{Deserialize, Serialize};

/// Radon-norm squared proximal penalty

#[derive(Copy, Clone, Serialize, Deserialize)]
pub struct RadonSquared;

#[replace_float_literals(F::cast_from(literal))]
impl<Domain, F, M, Reg> ProxPenalty<Domain, M, Reg, F> for RadonSquared
where
    Domain: Space + Clone + PartialEq + 'static,
    for<'a> &'a Domain: Instance<Domain>,
    F: Float + ToNalgebraRealField,
    M: MinMaxMapping<Domain, F>,
    Reg: RadonSquaredRegTerm<Domain, F>,
    DiscreteMeasure<Domain, F>: SpikeMerging<F>,
{
    type ReturnMapping = M;

    fn prox_type() -> ProxTerm {
        ProxTerm::RadonSquared
    }

    fn insert_and_reweigh<I>(
        &self,
        μ: &mut DiscreteMeasure<Domain, F>,
        τv: &mut M,
        τ: F,
        ε: F,
        config: &InsertionConfig<F>,
        reg: &Reg,
        _state: &AlgIteratorIteration<I>,
        stats: &mut IterInfo<F>,
    ) -> DynResult<(Option<Self::ReturnMapping>, bool)>
    where
        I: AlgIterator,
    {
        // If no merging heuristic is used, let's be more conservative about spike insertion,
        // and skip it after first round. If merging is done, being more greedy about spike
        // insertion also seems to improve performance.
        let skip_by_rough_check = !config.merging.enabled;
        let violation = reg.find_tolerance_violation(τv, τ, ε, skip_by_rough_check, config);
        reg.solve_oc_radonsq(μ, τv, τ, ε, violation, config, stats);

        Ok((None, true))
    }

    fn merge_spikes(
        &self,
        μ: &mut DiscreteMeasure<Domain, F>,
        τv: &mut M,
        μ_base: &DiscreteMeasure<Domain, F>,
        τ: F,
        ε: F,
        config: &InsertionConfig<F>,
        reg: &Reg,
        fitness: Option<impl Fn(&DiscreteMeasure<Domain, F>) -> F>,
    ) -> usize {
        if config.fitness_merging {
            if let Some(f) = fitness {
                return μ.merge_spikes_fitness(config.merging, f, |&v| v).1;
            }
        }
        μ.merge_spikes(config.merging, |μ_candidate| {
            // Important: μ_candidate's new points are afterwards,
            // and do not conflict with μ_base.
            // TODO: could simplify to requiring μ_base instead of μ_radon.
            // but may complicate with sliding base's exgtra points that need to be
            // after μ_candidate's extra points.
            // TODO: doesn't seem to work, maybe need to merge μ_base as well?
            // Although that doesn't seem to make sense.
            let μ_radon = μ_candidate.sub_matching(μ_base);
            reg.verify_merge_candidate_radonsq(τv, μ_candidate, τ, ε, &config, &μ_radon)
            //let n = μ_candidate.dist_matching(μ_base);
            //reg.find_tolerance_violation_slack(τv, τ, ε, false, config, n).is_none()
        })
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<'a, F, A, Domain> StepLengthBound<F, QuadraticDataTerm<F, Domain, A>> for RadonSquared
where
    F: Float + ToNalgebraRealField,
    Domain: Space + Norm<Radon, F>,
    A: ForwardModel<Domain, F> + BoundedLinear<Domain, Radon, L2, F>,
{
    fn step_length_bound(&self, f: &QuadraticDataTerm<F, Domain, A>) -> DynResult<F> {
        // TODO: direct squared calculation
        Ok(f.operator().opnorm_bound(Radon, L2)?.powi(2))
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<'a, F, A, Domain> StepLengthBoundPD<F, A, DiscreteMeasure<Domain, F>> for RadonSquared
where
    Domain: Space + Clone + PartialEq + 'static,
    F: Float + ToNalgebraRealField,
    A: BoundedLinear<DiscreteMeasure<Domain, F>, Radon, L2, F>,
{
    fn step_length_bound_pd(&self, opA: &A) -> DynResult<F> {
        opA.opnorm_bound(Radon, L2)
    }
}
