/*!
Norms, projections, etc.
*/

use crate::euclidean::*;
use crate::instance::Ownable;
use crate::linops::{ClosedVectorSpace, VectorSpace};
use crate::mapping::{Instance, Mapping, Space};
use crate::types::*;
use serde::{Deserialize, Serialize};
use std::marker::PhantomData;

//
// Abstract norms
//

#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
/// Helper structure to convert a [`NormExponent`] into a [`Mapping`]
pub struct NormMapping<F: Float, E: NormExponent> {
    pub(crate) exponent: E,
    _phantoms: PhantomData<F>,
}

/// An exponent for norms.
///
// Just a collection of desirable attributes for a marker type
pub trait NormExponent: Copy {
    /// Return the norm as a mappin
    fn as_mapping<F: Float>(self) -> NormMapping<F, Self> {
        NormMapping { exponent: self, _phantoms: PhantomData }
    }
}

/// Exponent type for the 1-[`Norm`].
#[derive(Copy, Debug, Clone, Serialize, Eq, PartialEq)]
pub struct L1;
impl NormExponent for L1 {}

/// Exponent type for the 2-[`Norm`].
#[derive(Copy, Debug, Clone, Serialize, Eq, PartialEq)]
pub struct L2;
impl NormExponent for L2 {}

/// Exponent type for the ∞-[`Norm`].
#[derive(Copy, Debug, Clone, Serialize, Eq, PartialEq)]
pub struct Linfinity;
impl NormExponent for Linfinity {}

/// Exponent type for 2,1-[`Norm`].
/// (1-norm over a domain Ω, 2-norm of a vector at each point of the domain.)
#[derive(Copy, Debug, Clone, Serialize, Eq, PartialEq)]
pub struct L21;
impl NormExponent for L21 {}

/// Norms for pairs (a, b). ‖(a,b)‖ = ‖(‖a‖_A, ‖b‖_B)‖_J
/// For use with [`crate::direct_product::Pair`]
#[derive(Copy, Debug, Clone, Serialize, Eq, PartialEq)]
pub struct PairNorm<A, B, J>(pub A, pub B, pub J);

impl<A, B, J> NormExponent for PairNorm<A, B, J>
where
    A: NormExponent,
    B: NormExponent,
    J: NormExponent,
{
}

/// A Huber/Moreau–Yosida smoothed [`L1`] norm. (Not a norm itself.)
///
/// The parameter γ of this type is the smoothing factor. Zero means no smoothing, and higher
/// values more smoothing. Behaviour with γ < 0 is undefined.
#[derive(Copy, Debug, Clone, Serialize, Eq, PartialEq)]
pub struct HuberL1<F: Float>(pub F);
impl<F: Float> NormExponent for HuberL1<F> {}

/// A Huber/Moreau–Yosida smoothed [`L21`] norm. (Not a norm itself.)
///
/// The parameter γ of this type is the smoothing factor. Zero means no smoothing, and higher
/// values more smoothing. Behaviour with γ < 0 is undefined.
#[derive(Copy, Debug, Clone, Serialize, Eq, PartialEq)]
pub struct HuberL21<F: Float>(pub F);
impl<F: Float> NormExponent for HuberL21<F> {}

/// A normed space (type) with exponent or other type `Exponent` for the norm.
///
/// Use as
/// ```
/// # use alg_tools::norms::{Norm, L1, L2, Linfinity};
/// # use alg_tools::loc::Loc;
/// let x = Loc([1.0, 2.0, 3.0]);
///
/// println!("{}, {} {}", x.norm(L1), x.norm(L2), x.norm(Linfinity))
/// ```
pub trait Norm<Exponent: NormExponent, F: Num = f64> {
    /// Calculate the norm.
    fn norm(&self, _p: Exponent) -> F;
}

/// Indicates that the `Self`-[`Norm`] is dominated by the `Exponent`-`Norm` on the space
/// `Elem` with the corresponding field `F`.
pub trait Dominated<F: Num, Exponent: NormExponent, Elem> {
    /// Indicates the factor $c$ for the inequality $‖x‖ ≤ C ‖x‖_p$.
    fn norm_factor(&self, p: Exponent) -> F;
    /// Given a norm-value $‖x‖_p$, calculates $C‖x‖_p$ such that $‖x‖ ≤ C‖x‖_p$
    #[inline]
    fn from_norm(&self, p_norm: F, p: Exponent) -> F {
        p_norm * self.norm_factor(p)
    }
}

/// Trait for distances with respect to a norm.
pub trait Dist<Exponent: NormExponent, F: Num = f64>: Norm<Exponent, F> + Space {
    /// Calculate the distance
    fn dist<I: Instance<Self>>(&self, other: I, _p: Exponent) -> F;
}

/// Trait for Euclidean projections to the `Exponent`-[`Norm`]-ball.
///
/// Use as
/// ```
/// # use alg_tools::norms::{Projection, L2, Linfinity};
/// # use alg_tools::loc::Loc;
/// let x = Loc([1.0, 2.0, 3.0]);
///
/// println!("{:?}, {:?}", x.proj_ball(1.0, L2), x.proj_ball(0.5, Linfinity));
/// ```
pub trait Projection<F: Num, Exponent: NormExponent>: Ownable + Norm<Exponent, F> {
    /// Projection of `self` to the `q`-norm-ball of radius ρ.
    fn proj_ball(self, ρ: F, q: Exponent) -> Self::OwnedVariant;
}

pub trait ProjectionMut<F: Num, Exponent: NormExponent>: Projection<F, Exponent> {
    /// In-place projection of `self` to the `q`-norm-ball of radius ρ.
    fn proj_ball_mut(&mut self, ρ: F, q: Exponent);
}

/*impl<F : Float, E : Euclidean<F>> Norm<L2, F> for E {
    #[inline]
    fn norm(&self, _p : L2) -> F { self.norm2() }

    fn dist(&self, other : &Self, _p : L2) -> F { self.dist2(other) }
}*/

impl<F: Float, E: Euclidean<F> + Norm<L2, F>> Projection<F, L2> for E {
    #[inline]
    fn proj_ball(self, ρ: F, _p: L2) -> Self::OwnedVariant {
        self.proj_ball2(ρ)
    }
}

impl<F: Float, E: EuclideanMut<F> + Norm<L2, F>> ProjectionMut<F, L2> for E {
    #[inline]
    fn proj_ball_mut(&mut self, ρ: F, _p: L2) {
        self.proj_ball2_mut(ρ)
    }
}

impl<F: Float> HuberL1<F> {
    fn apply(self, xnsq: F) -> F {
        let HuberL1(γ) = self;
        let xn = xnsq.sqrt();
        if γ == F::ZERO {
            xn
        } else {
            if xn > γ {
                xn - γ / F::TWO
            } else if xn < (-γ) {
                -xn - γ / F::TWO
            } else {
                xnsq / (F::TWO * γ)
            }
        }
    }
}

impl<F: Float, E: Euclidean<F> + Normed<F, NormExp = L2>> Norm<HuberL1<F>, F> for E {
    fn norm(&self, huber: HuberL1<F>) -> F {
        huber.apply(self.norm2_squared())
    }
}

impl<F: Float, E: Euclidean<F> + Normed<F, NormExp = L2>> Dist<HuberL1<F>, F> for E {
    fn dist<I: Instance<Self>>(&self, other: I, huber: HuberL1<F>) -> F {
        huber.apply(self.dist2_squared(other))
    }
}

// impl<F : Float, E : Norm<L2, F>> Norm<L21, F> for Vec<E> {
//     fn norm(&self, _l21 : L21) -> F {
//         self.iter().map(|e| e.norm(L2)).sum()
//     }
// }

// impl<F : Float, E : Dist<F, L2>> Dist<L21, F> for Vec<E> {
//     fn dist<I : Instance<Self>>(&self, other : I, _l21 : L21) -> F {
//         self.iter().zip(other.iter()).map(|(e, g)| e.dist(g, L2)).sum()
//     }
// }

impl<E, F, Domain> Mapping<Domain> for NormMapping<F, E>
where
    F: Float,
    E: NormExponent,
    Domain: Space,
    Domain::Principal: Norm<E, F>,
{
    type Codomain = F;

    #[inline]
    fn apply<I: Instance<Domain>>(&self, x: I) -> F {
        x.eval(|r| r.norm(self.exponent))
    }
}

pub trait Normed<F: Num = f64>: Space + Norm<Self::NormExp, F> {
    type NormExp: NormExponent;

    fn norm_exponent(&self) -> Self::NormExp;

    #[inline]
    fn norm_(&self) -> F {
        self.norm(self.norm_exponent())
    }

    // fn similar_origin(&self) -> Self;

    fn is_zero(&self) -> bool {
        self.norm_() == F::ZERO
    }
}

pub trait HasDual<F: Num = f64>: Normed<F> + VectorSpace<Field = F> {
    type DualSpace: Normed<F> + ClosedVectorSpace<Field = F>;

    fn dual_origin(&self) -> <Self::DualSpace as VectorSpace>::PrincipalV;
}

/// Automatically implemented trait for reflexive spaces
pub trait Reflexive<F: Num = f64>: HasDual<F>
where
    Self::DualSpace: HasDual<F, DualSpace = Self::Principal>,
{
}

impl<F: Num, X: HasDual<F>> Reflexive<F> for X where
    X::DualSpace: HasDual<F, DualSpace = Self::Principal>
{
}

pub trait HasDualExponent: NormExponent {
    type DualExp: NormExponent;

    fn dual_exponent(&self) -> Self::DualExp;
}

impl HasDualExponent for L2 {
    type DualExp = L2;

    #[inline]
    fn dual_exponent(&self) -> Self::DualExp {
        L2
    }
}

impl HasDualExponent for L1 {
    type DualExp = Linfinity;

    #[inline]
    fn dual_exponent(&self) -> Self::DualExp {
        Linfinity
    }
}

impl HasDualExponent for Linfinity {
    type DualExp = L1;

    #[inline]
    fn dual_exponent(&self) -> Self::DualExp {
        L1
    }
}

#[macro_export]
macro_rules! impl_weighted_norm {
    ($exponent : ty) => {
        impl<C, F, D> Norm<F, Weighted<$exponent, C>> for D
        where
            F: Float,
            D: Norm<$exponent, F>,
            C: Constant<Type = F>,
        {
            fn norm(&self, e: Weighted<$exponent, C>) -> F {
                let v = e.weight.value();
                assert!(v > F::ZERO);
                v * self.norm(e.base_fn)
            }
        }

        impl<C: Constant> NormExponent for Weighted<$exponent, C> {}

        impl<C: Constant> HasDualExponent for Weighted<$exponent, C>
        where
            $exponent: HasDualExponent,
        {
            type DualExp = Weighted<<$exponent as HasDualExponent>::DualExp, C::Type>;

            fn dual_exponent(&self) -> Self::DualExp {
                Weighted {
                    weight: C::Type::ONE / self.weight.value(),
                    base_fn: self.base_fn.dual_exponent(),
                }
            }
        }

        impl<C, F, T> Projection<F, Weighted<$exponent, C>> for T
        where
            T: Projection<F, $exponent>,
            F: Float,
            C: Constant<Type = F>,
        {
            fn proj_ball(self, ρ: F, q: Weighted<$exponent, C>) -> Self {
                self.proj_ball(ρ / q.weight.value(), q.base_fn)
            }

            fn proj_ball_mut(&mut self, ρ: F, q: Weighted<$exponent, C>) {
                self.proj_ball_mut(ρ / q.weight.value(), q.base_fn)
            }
        }
    };
}

//impl_weighted_norm!(L1);
//impl_weighted_norm!(L2);
//impl_weighted_norm!(Linfinity);
