rust-roche 0.1.4

use std::f64::consts::TAU;
use crate::errors::RocheError;
use crate::{Vec3, Star, Etype};
use crate::{rpot_val, rpot_grad, rpot1, rpot2, drpot1, drpot2};
use crate::{pot_min, dbrent, x_l1, x_l1_1, x_l1_2, x_l2, x_l3};
use crate::{sphere_eclipse, sphere_eclipse_vector, set_earth};


///
/// RocheContext is a basic struct packaging the mass ratio, q=M2/M1,
/// the star (i.e. the Primary or Secondary), the spin of that star,
/// and the x-coordinate of the L1 point together with the functions
/// that accept these values as inputs. x_l1 is calculated on initialisation
/// with RocheContext::new(q, star, spin) and saves it being recalculated
/// over and over in loops.
/// 
pub struct RocheContext {
    pub q: f64,
    pub star: Star,
    pub spin: f64,
    pub x_l1: f64,
}

impl RocheContext {
    pub fn new(q: f64, star: Star, spin: f64) -> Result<Self, RocheError> {

        let x_l1: f64 = match star {
            Star::Primary => x_l1_1(q, spin)?,
            Star::Secondary => x_l1_2(q, spin)?,
        };
        Ok(Self { q, star, spin, x_l1 })
    }

    pub fn potential(&self, earth: &Vec3, p:&Vec3, lam: f64) -> Result<f64, RocheError> {
        Ok(rpot_val(self.q, self.star, self.spin, earth, p, lam)?)
    }


    pub fn gradient(&self, earth: &Vec3, p: &Vec3, lam: f64) -> Result<(f64, f64), RocheError> {
        let (dp, dl) = rpot_grad(self.q, self.star, self.spin, earth, p, lam)?;
        Ok((dp, dl))
    }


    pub fn potential_grad(&self, earth: &Vec3, p: &Vec3, lam: f64) -> Result<(f64, f64, f64), RocheError> {
        let f = self.potential(earth, p, lam)?;
        let (dp, dl) = rpot_grad(self.q, self.star, self.spin, earth, p, lam)?;
        Ok((f, dp, dl))
    }


    pub fn ref_sphere(&self, ffac: f64) -> Result<(f64, f64), RocheError> {
        let tref: f64;
        let rref: f64;
        let pref: f64;
        if self.star == Star::Primary {
            tref = self.x_l1;
            rref = tref * 1.0_f64.min(1.001*ffac);
            pref = rpot1(self.q, self.spin, &Vec3 { x: ffac*tref, y: 0.0, z: 0.0 })?;
            Ok((rref, pref))
        } else if self.star == Star::Secondary {
            tref = 1.0 - self.x_l1;
            rref = tref * 1.0_f64.min(1.001*ffac);
            pref = rpot2(self.q, self.spin, &Vec3 { x: 1.0 - ffac*tref, y: 0.0, z: 0.0 })?;
            Ok((rref, pref))
        } else {
            let message = format!("{:?} is not and instance of Star.", self.star);
            return Err(RocheError::ParameterError(message));
        }
    }


    pub fn fblink(&self, ffac: f64, acc: f64, earth: &Vec3, p: &Vec3) -> Result<bool, RocheError> {

        let (rref, pref) = self.ref_sphere(ffac)?;

        let cofm: Vec3 = match self.star {
            Star::Primary => Vec3::cofm1(),
            Star::Secondary => Vec3::cofm2(),
        };
        
        // First compute the multipliers cutting the reference sphere (if any)
        let mut lam1 = 0.0;
        let mut lam2 = 0.0;
        if !sphere_eclipse_vector(earth, p, &cofm, rref, &mut lam1, &mut lam2) {
            return Ok(false);
        }
        if lam1 == 0.0 {
            return Ok(true);
        }

        // Create function objects for 1D minimisation in lambda direction
        let func = |lam: f64| {
            Ok(self.potential(earth, p, lam)?)
        };

        // Now try to bracket a minimum. We just crudely compute function at regularly spaced intervals filling in the
        // gaps until the step size between the points drops below the threshold. Take every opportunity to jump out early
        // either if the potential is below the threshold or if we have bracketed a minimum.
        let mut nstep: i32 = 1;
        let mut step: f64 = lam2 - lam1;

        let mut f1: f64 = 0.0;
        let mut f2: f64 = 0.0;
        let mut flam: f64 = 1.0;
        let mut lam: f64 = lam1;

        while step > acc {

            lam = lam1 + step/2.0;

            for _ in 0..nstep {

                flam = func(lam)?;
                if flam <= pref {
                    return Ok(true);
                }

                // Calculate these as late as possible because they may often not be needed
                if nstep == 1 {
                    f1 = func(lam1)?;
                    f2 = func(lam2)?;
                }

                if flam < f1 && flam < f2 {
                    break;
                }

                lam += step;
            }
            if flam < f1 && flam < f2 {
                break;
            }
            step /= 2.0;
            nstep *= 2;
        }

        if flam < f1 && flam < f2 {

            // OK, minimum bracketted, so finally pin it down accurately
            // Possible that multiple minima could cause problems but I have
            // never seen this in practice.
            let dfunc = |lam: f64| {
                let (_dp, dl) = self.gradient(earth, p, lam)?;
                Ok(dl)
            };

            let (_xmin, flam) = dbrent(lam1, lam, lam2, |x| func(x), |x| dfunc(x), acc, true, pref)?;

            Ok(flam < pref)
        } else {
            // Not bracketted even after a detailed search, and we have not jumped 
	        // out either, so assume no eclipse
            Ok(false)
        }


    }


    pub fn face(&self, direction: Vec3, rref: f64, pref: f64, acc: f64) -> Result<(Vec3, Vec3, f64, f64), RocheError> {

        let mut pvec: Vec3;
        let mut r: f64;

        let cofm: Vec3 = match self.star {
            Star::Primary => Vec3::cofm1(),
            Star::Secondary => Vec3::cofm2(),
        };

        let rp: fn(f64, f64, &Vec3) -> Result<f64, RocheError> = match self.star {
            Star::Primary => rpot1,
            Star::Secondary => rpot2,
        };

        let drp: fn(f64, f64, &Vec3) -> Result<Vec3, RocheError> = match self.star {
            Star::Primary => drpot1,
            Star::Secondary => drpot2,
        };

        let mut tref: f64 = rp(self.q, self.spin, &(cofm + rref*direction))?;
        if tref < pref {
            let message = format!(
            "point at reference radius {} appears to be at lower potential {} than the reference potential {}", rref, tref, pref
            );
            return Err(RocheError::FaceError(message));
        }

        let mut r1: f64 = rref/2.;
        let mut r2: f64 = rref;
        tref = pref + 1.;

        const MAXSEARCH: i32 = 30;
        let mut i: i32 = 0;
        while i < MAXSEARCH && tref > pref {
            r1 = r2/2.;
            tref = rp(self.q, self.spin, &(cofm + r1*direction))?;
            if tref > pref {
                r2 = r1;
            }
            i+=1;
        }
        if tref > pref {
            let message = "could not find a radius with a potential below the reference potential; probably bad inputs.";
            return Err(RocheError::FaceError(message.to_string()));
        }

        const MAXCHOP: i32 = 100;
        let mut nchop: i32 = 0;
        while r2 - r1 > acc && nchop < MAXCHOP {
            r = (r1 + r2)/2.;
            pvec = cofm + r*direction;
            if rp(self.q, self.spin, &pvec)? < pref {
                r1 = r;
            }else {
                r2 = r;
            }
            nchop += 1;
        }
        if nchop == MAXCHOP {
            return Err(RocheError::FaceError("reached maximum number of binary chops".to_string()));
        }
        r = (r1 + r2)/2.;
        pvec = cofm + r*direction;
        let mut dvec: Vec3 = drp(self.q, self.spin, &pvec)?;
        let g = dvec.length();
        dvec /= g;
        return Ok((pvec, dvec, r, g))
    }


    pub fn ingress_egress(&self, ffac: f64, iangle: f64, delta: f64, r: &Vec3, ingress: &mut f64, egress: &mut f64) -> Result<bool, RocheError> {
        let rref: f64;
        let pref: f64;
        (rref, pref) = self.ref_sphere(ffac)?;
        let ri: f64 = iangle.to_radians();
        let (sini, cosi) = ri.sin_cos();

        let cofm: Vec3 = match self.star {
            Star::Primary => Vec3::cofm1(),
            Star::Secondary => Vec3::cofm2(),
        };

        let mut phi1: f64 = 0.0;
        let mut phi2: f64 = 0.0;
        let mut lam1: f64 = 0.0;
        let mut lam2: f64 = 0.0;
        let mut phi: f64 = 0.0;
        let mut lam: f64 = 0.0;

        if sphere_eclipse(cosi, sini, r, &cofm, rref, &mut phi1, &mut phi2, &mut lam1, &mut lam2) {
            
            let acc: f64 = 2.*(2.*TAU*(lam2 - lam1)*delta).sqrt();

            if self.pot_min(cosi, sini, r, phi1, phi2, lam1, lam2, rref, pref, acc, &mut phi, &mut lam)? {

                let mut pin: f64 = phi;
                let mut pout: f64 = phi1;
                let mut pmid: f64;

                while (pin - pout).abs() > delta {
                    pmid = (pin + pout)/2.0;
                    if self.fblink(ffac, acc, &set_earth(cosi, sini, pmid), r).unwrap() {
                        pin = pmid;
                    } else {
                        pout = pmid;
                    }
                }
                *ingress = (pin+pout)/2.0;
                *ingress = *ingress - ingress.floor();

                pin = phi;
                pout = phi2;
                while (pin-pout).abs() > delta {
                    pmid = (pin+pout)/2.;
                    if self.fblink(ffac, acc, &set_earth(cosi, sini, pmid), r).unwrap() {
                        pin = pmid;
                    } else {
                        pout = pmid;
                    }
                }
                *egress = (pin+pout)/2.0;
                *egress = *egress - egress.floor();
                if *egress < *ingress {
                    *egress += 1.0;
                }
                return Ok(true);
            } else {
                return Ok(false);
            }
        } else {
            return Ok(false);
        }

    }


    pub fn star_eclipse(&self, r: f64, ffac: f64, iangle: f64, posn: &Vec3, delta: f64, roche: bool, star: Star, eclipses: &mut Etype) -> Result<(), RocheError> {
        let ri = iangle.to_radians();
        let (sini, cosi) = ri.sin_cos();
        let cofm = match star {
            Star::Primary => Vec3::cofm1(),
            Star::Secondary => Vec3::cofm2(),
        };
        let mut lam1: f64 = 0.0;
        let mut lam2: f64 = 0.0;
        let mut ingress: f64 = 0.0;
        let mut egress: f64 = 0.0;
        // let mut eclipses = Etype::new();
        if (roche && self.ingress_egress(ffac, iangle, delta, &posn, &mut ingress, &mut egress)?) ||
            (!roche && sphere_eclipse(cosi, sini, &posn, &cofm, r, &mut ingress, &mut egress, &mut lam1, &mut lam2)) {
            eclipses.push((ingress, egress));
        }
        Ok(())
    }


    pub fn pot_min(&self, cosi: f64, sini: f64, p: &Vec3, phi1: f64, phi2: f64, lam1: f64, lam2: f64, rref: f64, pref: f64, acc: f64, phi: &mut f64, lam: &mut f64) -> Result<bool, RocheError> {
        Ok(pot_min(self.q, self.star, self.spin, cosi, sini, p, phi1, phi2, lam1, lam2, rref, pref, acc, phi, lam)?)
    }


    pub fn x_l1(&self) -> Result<f64, RocheError> {
        Ok(x_l1(self.q)?)
    }


    pub fn x_l1_asyncronous(&self) -> Result<f64, RocheError> {
        match self.star {
            Star::Primary => Ok(x_l1_1(self.q, self.spin)?),
            Star::Secondary => Ok(x_l1_2(self.q, self.spin)?),
        }
    }


    pub fn x_l2(&self) -> Result<f64, RocheError> {
        Ok(x_l2(self.q)?)
    }


    pub fn x_l3(&self) -> Result<f64, RocheError> {
        Ok(x_l3(self.q)?)
    }

}