/*!
Traits for mathematical functions.
*/

use std::marker::PhantomData;
use crate::types::{Float};
use serde::Serialize;
use crate::loc::Loc;

/// Trait for application of `Self` as a mathematical function or operator on `X`.
pub trait Apply<X> {
    type Output;

    /// Compute the value of `self` at `x`.
    fn apply(&self, x : X) -> Self::Output;
}

/// This blanket implementation is a workaround helper to Rust trait system limitations.
///
/// It is introduced because the trait system does not allow blanket implementations of both
/// [`Apply<X>`] and [`Apply<&'a X>`]. With this, the latter is implemented automatically for
/// the reference, which can be sufficient to apply the operation in another blanket implementation.
impl<'a, T, X> Apply<X> for &'a T where T : Apply<X> {
    type Output = <T as Apply<X>>::Output;

    #[inline]
    fn apply(&self, x : X) -> Self::Output {
        (*self).apply(x)
    }
}

/// A mapping from `Domain` to `Codomain`.
///
/// This is automatically implemented when the relevant [`Apply`] are implemented.
pub trait Mapping<Domain> : Apply<Domain, Output=Self::Codomain>
                            + for<'a> Apply<&'a Domain, Output=Self::Codomain> {
    type Codomain;
}

impl<Domain, Codomain, T> Mapping<Domain> for T
where T : Apply<Domain, Output=Codomain> + for<'a> Apply<&'a Domain, Output=Codomain> {
    type Codomain = Codomain;
}


pub trait RealMapping<F : Float, const N : usize>
: Mapping<Loc<F, N>, Codomain = F> {}

impl<F : Float, T, const N : usize> RealMapping<F, N> for T
where T : Mapping<Loc<F, N>, Codomain = F> {}


/// Trait for calculation the differential of `Self` as a mathematical function on `X`.
pub trait Differentiable<X> : Sized {
    type Output;

    /// Compute the differential of `self` at `x`.
    fn differential(&self, x : X) -> Self::Output;
}

/// A differentiable mapping from `Domain` to [`Mapping::Codomain`], with differentials
/// `Differential`.
///
/// This is automatically implemented when the relevant [`Differentiate`] are implemented.
pub trait DifferentiableMapping<Domain>
: Mapping<Domain>
  + Differentiable<Domain, Output=Self::Differential>
  + for<'a> Differentiable<&'a Domain, Output=Self::Differential> {
    type Differential;

    /// Form the differential mapping of `self`.
    fn diff(self) -> Differential<Domain, Self> {
        Differential{ g : self, _space : PhantomData }
    }
}


impl<Domain, Differential, T> DifferentiableMapping<Domain> for T
where T : Mapping<Domain>
          + Differentiable<Domain, Output=Differential>
          + for<'a> Differentiable<&'a Domain, Output=Differential> {
    type Differential = Differential;
}

/// A sum of [`Mapping`]s.
#[derive(Serialize, Debug, Clone)]
pub struct Sum<Domain, M : Mapping<Domain>> {
    components : Vec<M>,
    _domain : PhantomData<Domain>,
}

impl<Domain, M : Mapping<Domain>> Sum<Domain, M> {
    /// Construct from an iterator.
    pub fn new<I : Iterator<Item = M>>(iter : I) -> Self {
        Sum { components : iter.collect(), _domain : PhantomData }
    }
}


impl<Domain : Copy, M> Apply<Domain> for Sum<Domain, M>
where M : Mapping<Domain>,
      M::Codomain : std::iter::Sum {
    type Output = M::Codomain;

    fn apply(&self, x : Domain) -> Self::Output {
        self.components.iter().map(|c| c.apply(x)).sum()
    }
}

impl<Domain, M> Differentiable<Domain> for Sum<Domain, M>
where M : DifferentiableMapping<Domain>,
      M :: Codomain : std::iter::Sum,
      M :: Differential : std::iter::Sum,
      Domain : Copy {

    type Output = M::Differential;

    fn differential(&self, x : Domain) -> Self::Output {
        self.components.iter().map(|c| c.differential(x)).sum()
    }
}

/// Container for the differential [`Mapping`] of a [`Differentiable`] mapping.
pub struct Differential<X, G : DifferentiableMapping<X>> {
    g : G,
    _space : PhantomData<X>
}

impl<X, G : DifferentiableMapping<X>> Apply<X> for Differential<X, G> {
    type Output = G::Differential;

    #[inline]
    fn apply(&self, x : X) -> G::Differential {
        self.g.differential(x)
    }
}

impl<'a, X, G : DifferentiableMapping<X>> Apply<&'a X> for Differential<X, G> {
    type Output = G::Differential;

    #[inline]
    fn apply(&self, x : &'a X) -> G::Differential {
        self.g.differential(x)
    }
}

    #[inline]
    fn apply(&self, x : X) -> Self::Output {
        self.g.differential(x)
    }
}
