/*!
Euclidean spaces.
*/

use crate::instance::Instance;
use crate::linops::{VectorSpace, AXPY};
use crate::norms::{HasDual, Norm, Normed, Reflexive, L2};
use crate::types::*;

pub mod wrap;

// TODO: Euclidean & EuclideanMut
//
/// Space (type) with Euclidean and vector space structure
///
/// The type should implement vector space operations (addition, subtraction, scalar
/// multiplication and scalar division) along with their assignment versions, as well
/// as an inner product.
// TODO: remove F parameter, use VectorSpace::Field
pub trait Euclidean<F: Float = f64>:
    VectorSpace<Field = F, PrincipalV = Self::PrincipalE> + Reflexive<F, DualSpace = Self::PrincipalE>
{
    /// Principal form of the space; always equal to [`crate::linops::Space::Principal`] and
    /// [`VectorSpace::PrincipalV`], but with more traits guaranteed.
    type PrincipalE: ClosedEuclidean<F>;

    // Inner product
    fn dot<I: Instance<Self>>(&self, other: I) -> F;

    /// Calculate the square of the 2-norm, $\frac{1}{2}\\|x\\|_2^2$, where `self` is $x$.
    ///
    /// This is not automatically implemented to avoid imposing
    /// `for <'a> &'a Self : Instance<Self>` trait bound bloat.
    fn norm2_squared(&self) -> F;

    /// Calculate the square of the 2-norm divided by 2, $\frac{1}{2}\\|x\\|_2^2$,
    /// where `self` is $x$.
    #[inline]
    fn norm2_squared_div2(&self) -> F {
        self.norm2_squared() / F::TWO
    }

    /// Calculate the 2-norm $‖x‖_2$, where `self` is $x$.
    #[inline]
    fn norm2(&self) -> F {
        self.norm2_squared().sqrt()
    }

    /// Calculate the 2-distance squared $\\|x-y\\|_2^2$, where `self` is $x$.
    fn dist2_squared<I: Instance<Self>>(&self, y: I) -> F;

    /// Calculate the 2-distance $\\|x-y\\|_2$, where `self` is $x$.
    #[inline]
    fn dist2<I: Instance<Self>>(&self, y: I) -> F {
        self.dist2_squared(y).sqrt()
    }

    /// Projection to the 2-ball.
    #[inline]
    fn proj_ball2(self, ρ: F) -> Self::PrincipalV {
        let r = self.norm2();
        if r > ρ {
            self * (ρ / r)
        } else {
            self.into_owned()
        }
    }
}

pub trait ClosedEuclidean<F: Float = f64>:
    Instance<Self> + Euclidean<F, PrincipalE = Self>
{
}
impl<F: Float, X: Instance<X> + Euclidean<F, PrincipalE = Self>> ClosedEuclidean<F> for X {}

// TODO: remove F parameter, use AXPY::Field
pub trait EuclideanMut<F: Float = f64>: Euclidean<F> + AXPY<Field = F> {
    /// In-place projection to the 2-ball.
    #[inline]
    fn proj_ball2_mut(&mut self, ρ: F) {
        let r = self.norm2();
        if r > ρ {
            *self *= ρ / r
        }
    }
}

impl<X, F: Float> EuclideanMut<F> for X where X: Euclidean<F> + AXPY<Field = F> {}

/// Trait for [`Euclidean`] spaces with dimensions known at compile time.
pub trait StaticEuclidean<F: Float = f64>: Euclidean<F> {
    /// Returns the origin
    fn origin() -> <Self as Euclidean<F>>::PrincipalE;
}

macro_rules! scalar_euclidean {
    ($f:ident) => {
        impl VectorSpace for $f {
            type Field = $f;
            type PrincipalV = $f;

            #[inline]
            fn similar_origin(&self) -> Self::PrincipalV {
                0.0
            }
        }
        impl AXPY for $f {
            #[inline]
            fn axpy<I: Instance<$f>>(&mut self, α: $f, x: I, β: $f) {
                *self = β * *self + α * x.own()
            }

            #[inline]
            fn set_zero(&mut self) {
                *self = 0.0
            }
        }

        impl Norm<L2, $f> for $f {
            fn norm(&self, _p: L2) -> $f {
                self.abs()
            }
        }

        impl Normed<$f> for $f {
            type NormExp = L2;

            fn norm_exponent(&self) -> Self::NormExp {
                L2
            }
        }

        impl HasDual<$f> for $f {
            type DualSpace = $f;

            #[inline]
            fn dual_origin(&self) -> $f {
                0.0
            }
        }

        impl Euclidean<$f> for $f {
            type PrincipalE = $f;

            #[inline]
            fn dot<I: Instance<Self>>(&self, other: I) -> $f {
                *self * other.own()
            }

            #[inline]
            fn norm2_squared(&self) -> $f {
                *self * *self
            }

            #[inline]
            fn dist2_squared<I: Instance<Self>>(&self, y: I) -> $f {
                let d = *self - y.own();
                d * d
            }
        }
    };
}

scalar_euclidean!(f64);
scalar_euclidean!(f32);
