src/fe_model/p2_local_model.rs

use crate::types::*;
use crate::loc::Loc;
use crate::sets::{Set,NPolygon,SpannedHalfspace};
use crate::linsolve::*;
use crate::norms::Dot;
use super::base::{LocalModel,RealLocalModel};
use crate::sets::Cube;
use numeric_literals::replace_float_literals;

pub struct Simplex<F : Float, const N : usize, const D : usize>(pub [Loc<F, N>; D]);
pub type PlanarSimplex<F> = Simplex<F,2,3>;
pub type RealInterval<F> = Simplex<F,1,2>;

#[inline]
#[replace_float_literals(F::cast_from(literal))]
pub(crate) fn midpoint<F : Float, const N : usize>(a : &Loc<F,N>, b : &Loc<F,N>) -> Loc<F, N> {
    (a+b)/2.0
}

impl<'a, F : Float> Set<Loc<F,1>> for RealInterval<F> {
    #[inline]
    fn contains(&self, &Loc([x]) : &Loc<F,1>) -> bool {
        let &[Loc([x0]), Loc([x1])] = &self.0;
        (x0 < x && x < x1) || (x1 < x && x < x0)
    }
}

impl<'a, F : Float> Set<Loc<F,2>> for PlanarSimplex<F> {
    #[inline]
    fn contains(&self, x : &Loc<F,2>) -> bool {
        let &[x0, x1, x2] = &self.0;
        NPolygon([[x0, x1].spanned_halfspace(),
                  [x1, x2].spanned_halfspace(),
                  [x2, x0].spanned_halfspace()]).contains(x)
    }
}

trait P2Powers {
    type Output;
    type Diff;
    type Full;
    fn p2powers(&self) -> Self::Output;
    fn p2powers_full(&self) -> Self::Full;
    fn p2powers_diff(&self) -> Self::Diff;
}

#[replace_float_literals(F::cast_from(literal))]
impl<F : Float> P2Powers for Loc<F, 1> {
    type Output = Loc<F, 1>;
    type Full = Loc<F, 3>;
    type Diff = Loc<Loc<F, 1>, 1>;

    #[inline]
    fn p2powers(&self) -> Self::Output {
        let &Loc([x0]) = self;
        [x0*x0].into()
    }

    #[inline]
    fn p2powers_full(&self) -> Self::Full {
        let &Loc([x0]) = self;
        [1.0, x0, x0*x0].into()
    }

    #[inline]
    fn p2powers_diff(&self) -> Self::Diff {
        let &Loc([x0]) = self;
        [[x0+x0].into()].into()
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F : Float> P2Powers for Loc<F, 2> {
    type Output = Loc<F, 3>;
    type Full = Loc<F, 6>;
    type Diff = Loc<Loc<F, 3>, 2>;

    #[inline]
    fn p2powers(&self) -> Self::Output {
        let &Loc([x0, x1]) = self;
        [x0*x0, x0*x1, x1*x1].into()
    }

    #[inline]
    fn p2powers_full(&self) -> Self::Full {
        let &Loc([x0, x1]) = self;
        [1.0, x0, x1, x0*x0, x0*x1, x1*x1].into()
    }

    #[inline]
    fn p2powers_diff(&self) -> Self::Diff {
        let &Loc([x0, x1]) = self;
        [[x0+x0, x1, 0.0].into(), [0.0, x0, x1+x1].into()].into()
    }
}

pub trait P2Model<F : Num, const N : usize> {
    type Model : LocalModel<Loc<F,N>,F>;
    fn p2_model<G : Fn(&Loc<F, N>) -> F>(&self, g : G) -> Self::Model;
}

pub struct P2LocalModel<F : Num, const N : usize, const E : usize, const V : usize, const Q : usize> {
    a0 : F,
    a1 : Loc<F, N>,
    a2 :  Loc<F, E>,
    //node_values : Loc<F, V>,
    //edge_values : Loc<F, Q>,
}

//
// 1D planar model construction
//

impl<F : Float> RealInterval<F> {
    #[inline]
    fn midpoints(&self) -> [Loc<F, 1>; 1] {
        let [ref n0, ref n1] = &self.0;
        let n01 = midpoint(n0, n1);
        [n01]
    }
}

impl<F : Float> P2LocalModel<F, 1, 1, 2, 0> {
    #[inline]
    pub fn new(
        &[n0, n1, n01] : &[Loc<F, 1>; 3],
        &[v0, v1, v01] : &[F; 3],
    ) -> Self {
        let p = move |x : &Loc<F, 1>, v : F| {
            let Loc([c, d, e]) = x.p2powers_full();
            [c, d, e, v]
        };
        let [a0, a1, a11] = linsolve([
            p(&n0, v0),
            p(&n1, v1),
            p(&n01, v01)
        ]);
        P2LocalModel {
            a0 : a0,
            a1 : [a1].into(),
            a2 : [a11].into(),
            //node_values : [v0, v1].into(),
            //edge_values: [].into(),
        }
    }
}

impl<F : Float> P2Model<F,1> for RealInterval<F> {
    type Model = P2LocalModel<F, 1, 1, 2, 0>;

    #[inline]
    fn p2_model<G : Fn(&Loc<F, 1>) -> F>(&self, g : G) -> Self::Model {
        let [n01] =  self.midpoints();
        let [n0, n1] = self.0;
        let vals = [g(&n0), g(&n1), g(&n01)];
        let nodes = [n0, n1, n01];
        Self::Model::new(&nodes, &vals)
    }
}

//
// 2D planar model construction
//

impl<F : Float> PlanarSimplex<F> {
    #[inline]
    fn midpoints(&self) -> [Loc<F, 2>; 3] {
        let [ref n0, ref n1, ref n2] = &self.0;
        let n01 = midpoint(n0, n1);
        let n12 = midpoint(n1, n2);
        let n20 = midpoint(n2, n0);
        [n01, n12, n20]
    }
}

impl<F : Float> P2LocalModel<F, 2, 3, 3, 3> {
    #[inline]
    pub fn new(
        &[n0, n1, n2, n01, n12, n20] : &[Loc<F, 2>; 6],
        &[v0, v1, v2, v01, v12, v20] : &[F; 6],
    ) -> Self {
        let p = move |x : &Loc<F,2>, v :F| {
            let Loc([c, d, e, f, g, h]) = x.p2powers_full();
            [c, d, e, f, g, h, v]
        };
        let [a0, a1, a2, a11, a12, a22] = linsolve([
            p(&n0, v0),
            p(&n1, v1),
            p(&n2, v2),
            p(&n01, v01),
            p(&n12, v12),
            p(&n20, v20),
        ]);
        P2LocalModel {
            a0 : a0,
            a1 : [a1, a2].into(),
            a2 : [a11, a12, a22].into(),
            //node_values : [v0, v1, v2].into(),
            //edge_values: [v01, v12, v20].into(),
        }
    }
}

impl<F : Float> P2Model<F,2> for PlanarSimplex<F> {
    type Model = P2LocalModel<F, 2, 3, 3, 3>;

    #[inline]
    fn p2_model<G : Fn(&Loc<F, 2>) -> F>(&self, g : G) -> Self::Model {
        let midpoints = self.midpoints();
        let [ref n0, ref n1, ref n2] = self.0;
        let [ref n01, ref n12, ref n20] = midpoints;
        let vals = [g(n0), g(n1), g(n2), g(n01), g(n12), g(n20)];
        let nodes = [*n0, *n1, *n2, *n01, *n12, *n20];
        Self::Model::new(&nodes, &vals)
    }
}

macro_rules! impl_local_model {
    ($n:literal, $e:literal, $v:literal, $q:literal) => {
        impl<F : Float> LocalModel<Loc<F, $n>, F> for P2LocalModel<F, $n, $e, $v, $q> {
            #[inline]
            fn value(&self, x : &Loc<F,$n>) -> F {
                self.a0 + x.dot(&self.a1) + x.p2powers().dot(&self.a2)
            }

            #[inline]
            fn differential(&self, x : &Loc<F,$n>) -> Loc<F,$n> {
                self.a1 + x.p2powers_diff().map(|di| di.dot(&self.a2))
            }
        }
    }
}

impl_local_model!(1, 1, 2, 0);
impl_local_model!(2, 3, 3, 3);


//
// Minimisation
//

#[replace_float_literals(F::cast_from(literal))]
impl<F : Float> P2LocalModel<F, 1, 1, 2, 0> {
    #[inline]
    fn minimise_edge(&self, x0 : Loc<F, 1>, x1 : Loc<F,1>) -> (Loc<F,1>, F) {
        let &P2LocalModel{
            a1 : Loc([a1]),
            a2 : Loc([a11]),
            //node_values : Loc([v0, v1]),
            ..
         } = self;
        // We do this in cases, first trying for an interior solution, then edges.
        // For interior solution, first check determinant; no point trying if non-positive
        if a11 > 0.0 {
            // An interior solution x[1] has to satisfy
            //   2a₁₁*x[1] + a₁ =0
            // This gives
            let t = -a1/(2.0*a11);
            let (Loc([t0]), Loc([t1])) = (x0, x1);
            if (t0 <= t && t <= t1) || (t1 <= t && t <= t0) {
                let x = [t].into();
                let v = self.value(&x);
                return (x, v)
            }
        }

        let v0 = self.value(&x0);
        let v1 = self.value(&x1);
        if v0 < v1 { (x0, v0) } else { (x1, v1) }
    }
}

impl<'a, F : Float> RealLocalModel<RealInterval<F>,Loc<F,1>,F>
for P2LocalModel<F, 1, 1, 2, 0> {
    #[inline]
    fn minimise(&self, &Simplex([x0, x1]) : &RealInterval<F>) -> (Loc<F,1>, F) {
        self.minimise_edge(x0, x1)
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F : Float> P2LocalModel<F, 2, 3, 3, 3> {
    /// Minimise the 2D model on the edge {x0 + t(x1 - x0) | t ∈ [0, 1] }
    #[inline]
    fn minimise_edge(&self, x0 : &Loc<F,2>, x1 : &Loc<F,2>/*, v0 : F, v1 : F*/) -> (Loc<F,2>, F) {
        let &P2LocalModel {
            a0,
            a1 : Loc([a1, a2]),
            a2 : Loc([a11, a12, a22]),
            ..
         } = self;
        let &Loc([x00, x01]) = x0;
        let d@Loc([d0, d1]) = x1 - x0;
        let b0 = a0 + a1*x00 + a2*x01 + a11*x00*x00 + a12*x00*x01 + a22*x01*x01;
        let b1 = a1*d0 + a2*d1 + 2.0*a11*d0*x00 + a12*(d0*x01 + d1*x00) + 2.0*a22*d1*x01;
        let b11 = a11*d0*d0 + a12*d0*d1 + a22*d1*d1;
        let edge_1d_model = P2LocalModel {
            a0 : b0,
            a1 : Loc([b1]),
            a2 : Loc([b11]),
            //node_values : Loc([v0, v1]),
        };
        let (Loc([t]), v) = edge_1d_model.minimise_edge(0.0.into(), 1.0.into());
        (x0 + d*t, v)
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<'a, F : Float> RealLocalModel<PlanarSimplex<F>,Loc<F,2>,F>
for P2LocalModel<F, 2, 3, 3, 3> {
    #[inline]
    fn minimise(&self, el : &PlanarSimplex<F>) -> (Loc<F,2>, F) {
        let &P2LocalModel {
            a1 : Loc([a1, a2]),
            a2 : Loc([a11, a12, a22]),
            //node_values : Loc([v0, v1, v2]),
            ..
         } = self;

        // We do this in cases, first trying for an interior solution, then edges.
        // For interior solution, first check determinant; no point trying if non-positive
        let r = 2.0*(a11*a22-a12*a12);
        if r > 0.0 {
            // An interior solution (x[1], x[2]) has to satisfy
            //   2a₁₁*x[1] + 2a₁₂*x[2]+a₁ =0 and 2a₂₂*x[1] + 2a₁₂*x[1]+a₂=0
            // This gives
            let x = [(a22*a1-a12*a2)/r, (a12*a1-a11*a2)/r].into();
            if el.contains(&x) {
                return (x, self.value(&x))
            }
        }

        let &[ref x0, ref x1, ref x2] = &el.0;
        let mut min_edge = self.minimise_edge(x0, x1);
        let more_edge = [self.minimise_edge(x1, x2),
                         self.minimise_edge(x2, x0)];
        
        for edge in more_edge {
            if edge.1 < min_edge.1 {
                min_edge = edge;
            }
        }

        min_edge
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<'a, F : Float> RealLocalModel<Cube<F, 2>,Loc<F,2>,F>
for P2LocalModel<F, 2, 3, 3, 3> {
    #[inline]
    fn minimise(&self, el : &Cube<F, 2>) -> (Loc<F,2>, F) {
        let &P2LocalModel {
            a1 : Loc([a1, a2]),
            a2 : Loc([a11, a12, a22]),
            //node_values : Loc([v0, v1, v2]),
            ..
         } = self;

        // We do this in cases, first trying for an interior solution, then edges.
        // For interior solution, first check determinant; no point trying if non-positive
        let r = 2.0*(a11*a22-a12*a12);
        if r > 0.0 {
            // An interior solution (x[1], x[2]) has to satisfy
            //   2a₁₁*x[1] + 2a₁₂*x[2]+a₁ =0 and 2a₂₂*x[1] + 2a₁₂*x[1]+a₂=0
            // This gives
            let x = [(a22*a1-a12*a2)/r, (a12*a1-a11*a2)/r].into();
            if el.contains(&x) {
                return (x, self.value(&x))
            }
        }

        let [x0, x1, x2, x3] = el.corners();
        let mut min_edge = self.minimise_edge(&x0, &x1);
        let more_edge = [self.minimise_edge(&x1, &x2),
                         self.minimise_edge(&x2, &x3),
                         self.minimise_edge(&x3, &x0)];

        for edge in more_edge {
            if edge.1 < min_edge.1 {
                min_edge = edge;
            }
        }

        min_edge
    }
}

// impl<F : Float> P2Model<F, 1> for Cube<F, 1> {
//     type Model = CubeP2LocalModel<F,1,1,1>;

//     fn p2_model<G : FnMut(&Loc<F, 1>) -> F>(&self, fun : &G) -> Self::Model {
//         CubeP2LocalModel([Simplex(self.corners()).p2_model(fun)])
//     }
// }

// impl<F : Float> P2Model<F, 2> for Cube<F, 2> {
//     type Model = CubeP2LocalModel<F,2,3,2>;

//     fn p2_model<G : FnMut(&Loc<F, 2>) -> F>(&self, fun : &G) -> Self::Model {
//         let [a, b, c, d] = self.corners();
//         CubeP2LocalModel([Simplex([a, b, c]).p2_model(fun),
//                           Simplex([b, c, d]).p2_model(fun)])
//     }
// }


#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn p2_model_1d_test() {
        let vertices = [Loc([0.0]), Loc([1.0])];
        let domain = Simplex(vertices);
        // A simple quadratic function for which the approximation is exact on reals,
        // and appears exact on f64 as well.
        let f = |&Loc([x]) : &Loc<f64, 1>| x*x + x + 1.0;
        let model = domain.p2_model(f);
        let xs = [Loc([0.5]), Loc([0.25])];

        for x in vertices.iter().chain(xs.iter()) {
            assert_eq!(model.value(&x), f(&x));
        }

        assert_eq!(model.minimise(&domain), (Loc([0.0]), 1.0));
    }

    #[test]
    fn p2_model_2d_test() {
        let vertices = [Loc([0.0, 0.0]), Loc([1.0, 0.0]), Loc([1.0, 1.0])];
        let domain = Simplex(vertices);
        // A simple quadratic function for which the approximation is exact on reals,
        // and appears exact on f64 as well.
        let f = |&Loc([x, y]) : &Loc<f64, 2>| - (x*x + x*y + x - 2.0 * y + 1.0);
        let model = domain.p2_model(f);
        let xs = [Loc([0.5, 0.5]), Loc([0.25, 0.25])];

        for x in vertices.iter().chain(xs.iter()) {
            assert_eq!(model.value(&x), f(&x));
        }

        assert_eq!(model.minimise(&domain), (Loc([1.0, 0.0]), -3.0));
    }
}