/*!
Quadratic functions  of the form $\frac{1}{2}\|Ax-b\|_2^2$ for an operator $A$
to a [`Euclidean`] space.
*/

#![allow(non_snake_case)]

use super::{DifferentiableImpl, LipschitzDifferentiableImpl, Mapping};
use crate::convex::ConvexMapping;
use crate::euclidean::Euclidean;
use crate::instance::{Instance, Space};
use crate::linops::{BoundedLinear, Linear, Preadjointable};
use crate::norms::{Norm, NormExponent, Normed, L2};
use crate::types::Float;
use std::marker::PhantomData;

/// Functions of the form $\frac{1}{2}\|Ax-b\|_2^2$ for an operator $A$
/// to a [`Euclidean`] space.
#[derive(Clone, Copy)]
pub struct Quadratic<'a, F: Float, Domain: Space, A: Mapping<Domain>> {
    opA: &'a A,
    b: &'a <A as Mapping<Domain>>::Codomain,
    _phantoms: PhantomData<F>,
}

#[allow(non_snake_case)]
impl<'a, F: Float, Domain: Space, A: Mapping<Domain>> Quadratic<'a, F, Domain, A> {
    pub fn new(opA: &'a A, b: &'a A::Codomain) -> Self {
        Quadratic {
            opA,
            b,
            _phantoms: PhantomData,
        }
    }

    pub fn operator(&self) -> &'a A {
        self.opA
    }

    pub fn data(&self) -> &'a <A as Mapping<Domain>>::Codomain {
        self.b
    }
}

//+ AdjointProductBoundedBy<RNDM<F, N>, P, FloatType = F>,

impl<'a, F: Float, X: Space, A: Mapping<X>> Mapping<X> for Quadratic<'a, F, X, A>
where
    A::Codomain: Euclidean<F>,
{
    type Codomain = F;

    fn apply<I: Instance<X>>(&self, x: I) -> F {
        // TODO: possibly (if at all more effcient) use GEMV once generalised
        // to not require preallocation. However, Rust should be pretty efficient
        // at not doing preallocations or anything here, as the result of self.opA.apply()
        // can be consumed, so maybe GEMV is no use.
        (self.opA.apply(x) - self.b).norm2_squared() / F::TWO
    }
}

impl<'a, F: Float, X: Normed<F>, A: Linear<X>> ConvexMapping<X, F> for Quadratic<'a, F, X, A> where
    A::Codomain: Euclidean<F>
{
}

impl<'a, F, X, A> DifferentiableImpl<X> for Quadratic<'a, F, X, A>
where
    F: Float,
    X: Space,
    <A as Mapping<X>>::Codomain: Euclidean<F>,
    A: Linear<X> + Preadjointable<X>,
    <<A as Mapping<X>>::Codomain as Euclidean<F>>::Output: Instance<<A as Mapping<X>>::Codomain>,
{
    type Derivative = A::PreadjointCodomain;

    fn differential_impl<I: Instance<X>>(&self, x: I) -> Self::Derivative {
        // TODO: possibly (if at all more effcient) use GEMV once generalised
        // to not require preallocation. However, Rust should be pretty efficient
        // at not doing preallocations or anything here, as the result of self.opA.apply()
        // can be consumed, so maybe GEMV is no use.
        self.opA.preadjoint().apply(self.opA.apply(x) - self.b)
    }
}

impl<'a, F, X, ExpX, A> LipschitzDifferentiableImpl<X, ExpX> for Quadratic<'a, F, X, A>
where
    F: Float,
    X: Space + Clone + Norm<F, ExpX>,
    ExpX: NormExponent,
    A: Clone + BoundedLinear<X, ExpX, L2, F>,
    Self: DifferentiableImpl<X>,
{
    type FloatType = F;

    fn diff_lipschitz_factor(&self, seminorm: ExpX) -> Option<F> {
        Some(self.opA.opnorm_bound(seminorm, L2).powi(2))
    }
}
