/* Copyright 2023 Bruce Merry
 *
 * This program is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program. If not, see <https://www.gnu.org/licenses/>.
 */

//! Prediction of the angle of the sun relative to solar panels
//!
//! This uses a relatively simple model that ignores all kinds of
//! effects:
//! - precession/nutation/frame bias
//! - refraction
//! - light travel time
//! - relativistic effects (aberration, deflection)
//! - polar motion
//! - variable rotation rate of the earth (specifically, dUT1)
//! - the Moon (it treats the Earth-Moon barycentre as the geocentre)
//!
//! Nevertheless it agrees with high-precision astronomy libraries to
//! better than a degree.
//!
//! The orbital parameters and the equations for applying them are taken
//! from <https://ssd.jpl.nasa.gov/planets/approx_pos.html>, table 2a.

// Lots of variables from external equations don't have snake case
#![allow(non_snake_case)]

use chrono::{DateTime, TimeZone};
use radians::{Angle, Deg64, Rad64, Unit, Wrap64};
use std::f64::consts::PI;
use std::ops::{Index, IndexMut, Mul, Neg};
use std::slice::SliceIndex;

#[derive(PartialEq, Default, Copy, Clone, Debug)]
struct Vector([f64; 3]);

impl<I: SliceIndex<[f64]>> Index<I> for Vector {
    type Output = I::Output;

    fn index(&self, index: I) -> &Self::Output {
        &self.0[index]
    }
}

impl<I: SliceIndex<[f64]>> IndexMut<I> for Vector {
    fn index_mut(&mut self, index: I) -> &mut Self::Output {
        &mut self.0[index]
    }
}

impl Mul<f64> for Vector {
    type Output = Self;

    fn mul(self, rhs: f64) -> Self::Output {
        Vector(self.0.map(|x| x * rhs))
    }
}

impl Neg for Vector {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Vector(self.0.map(|x| -x))
    }
}

impl Vector {
    fn iter(&self) -> std::slice::Iter<'_, f64> {
        self.0.iter()
    }

    /// Scale to unit length
    fn normalized(&self) -> Self {
        *self * (1.0 / dot(self, self)).sqrt()
    }
}

fn dot(a: &Vector, b: &Vector) -> f64 {
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| x * y)
        .fold(0.0, |x, y| x + y)
}

fn cross(a: &Vector, b: &Vector) -> Vector {
    Vector([
        a[1] * b[2] - a[2] * b[1],
        a[2] * b[0] - a[0] * b[2],
        a[0] * b[1] - a[1] * b[0],
    ])
}

#[derive(PartialEq, Default, Copy, Clone, Debug)]
struct Matrix([[f64; 3]; 3]);

impl Index<(usize, usize)> for Matrix {
    type Output = f64;

    fn index(&self, (r, c): (usize, usize)) -> &Self::Output {
        &self.0[r][c]
    }
}

impl IndexMut<(usize, usize)> for Matrix {
    fn index_mut(&mut self, (r, c): (usize, usize)) -> &mut Self::Output {
        &mut self.0[r][c]
    }
}

impl Mul<Vector> for Matrix {
    type Output = Vector;

    fn mul(self, rhs: Vector) -> Self::Output {
        Vector(self.0.map(|x| dot(&Vector(x), &rhs)))
    }
}

/// Solve for E in Kepler's equation M = E - e sin E.
fn kepler(M: Wrap64, e: f64) -> Wrap64 {
    let mut E = M.val() - e * M.sin();
    loop {
        let dM = M.val() - (E - e * E.sin());
        let dE = dM / (1.0 - e * E.cos());
        E += dE;
        if dE.abs() < 1e-8 {
            return Rad64::new(E).wrap();
        }
    }
}

/// Rotate around the X axis
fn Rx<U: Unit<f64>>(r: Angle<f64, U>) -> Matrix {
    let (s, c) = r.sin_cos();
    Matrix([[1.0, 0.0, 0.0], [0.0, c, s], [0.0, -s, c]])
}

/// Rotate around the Z axis
fn Rz<U: Unit<f64>>(r: Angle<f64, U>) -> Matrix {
    let (s, c) = r.sin_cos();
    Matrix([[c, s, 0.0], [-s, c, 0.0], [0.0, 0.0, 1.0]])
}

/// Convert time to a Julian Date relative to some base time (given as UNIX time)
fn timestamp_f64<Tz: TimeZone>(time: &DateTime<Tz>, epoch: f64) -> f64 {
    ((time.timestamp() as f64 - epoch) + 1e-9 * (time.timestamp_subsec_nanos() as f64)) / 86400.0
}

fn earth_rotation_angle<Tz: TimeZone>(time: &DateTime<Tz>) -> Wrap64 {
    // This ignores the difference between UTC and UT1. As such, there
    // isn't too much point worrying about the loss of precision in
    // presenting time as a single floating-point value.
    // This timestamp is relative to 2000-01-1T12:00:00 UTC, ignoring
    // leap seconds.
    let t = timestamp_f64(time, 946728000.0);
    Rad64::new((t.fract() + 0.779057273264 + 0.00273781191135448 * t).fract() * 2.0 * PI).wrap()
}

/// Direction from location to the sun, in east-north-up coordinate frame
fn sun_direction<Tz, U1, U2>(
    lat: Angle<f64, U1>,
    lon: Angle<f64, U2>,
    time: &DateTime<Tz>,
) -> Vector
where
    Tz: TimeZone,
    U1: Unit<f64>,
    U2: Unit<f64>,
{
    const J2000_EPOCH: f64 = 946727935.816;
    // Would ideally be `const` but Deg64::new is not a const function
    let OBLIQUITY = Deg64::new(23.43928);

    // Orbital elements, from NASA model
    let T = timestamp_f64(time, J2000_EPOCH) / 36525.0; // centuries
    let e = 0.01673163 - 0.00003661 * T;
    let I = Deg64::new(-0.00054346 - 0.01337178 * T);
    let L = Deg64::new(100.46691572 + 35999.37306329 * T);
    let ω_bar = Deg64::new(102.93005885 + 0.31795260 * T);
    let Ω = Deg64::new(-5.11260389 - 0.24123856 * T);

    let ω = ω_bar - Ω;
    let M = (L - ω_bar).rad().wrap();
    let E = kepler(M, e);
    let r_orbital = Vector([E.cos() - e, (1.0 - e * e).sqrt() * E.sin(), 0.0]);
    let r_eq = Rx(-OBLIQUITY) * (Rz(-Ω) * (Rx(-I) * (Rz(-ω) * r_orbital)));
    // Ignore precession-nutation and frame bias. Sign flip is to get Sun
    // relative to Earth instead of vice versa (ignoring the difference between
    // the geocentre and the Earth-Moon barycentre).
    let r_cirs = -r_eq;
    let era = earth_rotation_angle(time);
    let r_tirs = Rz(era.inner()) * r_cirs;
    let (slat, clat) = lat.sin_cos();
    let (slon, clon) = lon.sin_cos();
    let l_z = Vector([clat * clon, clat * slon, slat]);
    let l_x = cross(&Vector([0.0, 0.0, 1.0]), &l_z).normalized();
    let l_y = cross(&l_z, &l_x);
    Matrix([l_x.0, l_y.0, l_z.0]) * r_tirs.normalized() // ignores TIRS -> ITRS corrections
}

pub trait SolarPredictor<Tz: TimeZone>: Sync + Send {
    /// Compute peak energy (in W) for a solar panel at a given time
    fn power(&self, time: &DateTime<Tz>) -> f64;
}

pub struct PanelSolarPredictor<U1, U2, U3, U4>
where
    U1: Unit<f64>,
    U2: Unit<f64>,
    U3: Unit<f64>,
    U4: Unit<f64>,
{
    lat: Angle<f64, U1>,
    lon: Angle<f64, U2>,
    elevation: Angle<f64, U3>,
    azimuth: Angle<f64, U4>,
    power: f64,
}

impl<U1, U2, U3, U4> PanelSolarPredictor<U1, U2, U3, U4>
where
    U1: Unit<f64>,
    U2: Unit<f64>,
    U3: Unit<f64>,
    U4: Unit<f64>,
{
    pub fn new(
        lat: Angle<f64, U1>,
        lon: Angle<f64, U2>,
        elevation: Angle<f64, U3>,
        azimuth: Angle<f64, U4>,
        power: f64,
    ) -> Self {
        Self {
            lat,
            lon,
            elevation,
            azimuth,
            power,
        }
    }
}

impl<Tz, U1, U2, U3, U4> SolarPredictor<Tz> for PanelSolarPredictor<U1, U2, U3, U4>
where
    Tz: TimeZone,
    U1: Unit<f64> + Sync + Send,
    U2: Unit<f64> + Sync + Send,
    U3: Unit<f64> + Sync + Send,
    U4: Unit<f64> + Sync + Send,
{
    fn power(&self, time: &DateTime<Tz>) -> f64 {
        let sun_dir = sun_direction(self.lat, self.lon, time);
        if sun_dir[2] <= 0.0 {
            return 0.0; // below horizon
        }
        let (s_el, c_el) = self.elevation.sin_cos();
        let (s_az, c_az) = self.azimuth.sin_cos();
        let panel_dir = Vector([c_el * s_az, c_el * c_az, s_el]);
        self.power * dot(&sun_dir, &panel_dir).max(0.0)
    }
}

pub struct ArraySolarPredictor<T> {
    elements: Vec<T>,
}

impl<T> ArraySolarPredictor<T> {
    pub fn new(elements: Vec<T>) -> Self {
        Self { elements }
    }
}

impl<Tz: TimeZone, T: SolarPredictor<Tz>> SolarPredictor<Tz> for ArraySolarPredictor<T> {
    fn power(&self, time: &DateTime<Tz>) -> f64 {
        self.elements.iter().map(|e| e.power(time)).sum()
    }
}

#[cfg(test)]
pub(crate) mod test {
    use chrono::{DateTime, NaiveTime, TimeZone};

    use crate::sun::SolarPredictor;

    /// Simplified model (for testing).
    ///
    /// The power is linear over a given time range (every day, in UTC) and zero outside it.
    pub struct LinearSolarPredictor {
        start: NaiveTime,
        stop: NaiveTime,
        start_power: f64,
        stop_power: f64,
    }

    impl LinearSolarPredictor {
        pub fn new(start: NaiveTime, stop: NaiveTime, start_power: f64, stop_power: f64) -> Self {
            Self {
                start,
                stop,
                start_power,
                stop_power,
            }
        }
    }

    impl<Tz: TimeZone> SolarPredictor<Tz> for LinearSolarPredictor {
        fn power(&self, time: &DateTime<Tz>) -> f64 {
            let t = time.to_utc().time();
            if t < self.start || t >= self.stop {
                0.0
            } else {
                let t_ofs = (t - self.start).num_milliseconds() as f64;
                let stop_ofs = (self.stop - self.start).num_milliseconds() as f64;
                t_ofs / stop_ofs * (self.stop_power - self.start_power) + self.start_power
            }
        }
    }
}
