/*
    Nyx, blazing fast astrodynamics
    Copyright (C) 2018-onwards Christopher Rabotin <christopher.rabotin@gmail.com>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero 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 Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

use hifitime::TimeUnits;
use snafu::{ResultExt, ensure};

use crate::dynamics::guidance::{LocalFrame, Maneuver, MnvrRepr};
use crate::linalg::SVector;
use crate::md::objective::Objective;
use crate::md::{GuidanceSnafu, NotFiniteSnafu, TargetingError, prelude::*};
pub use crate::md::{Variable, Vary};
use crate::polyfit::CommonPolynomial;
use std::fmt;
use std::time::Duration;

/// Defines a targeter solution
#[derive(Clone, Debug)]
pub struct TargeterSolution<const V: usize, const O: usize> {
    /// The corrected spacecraft state at the correction epoch
    pub corrected_state: Spacecraft,
    /// The state at which the objectives are achieved
    pub achieved_state: Spacecraft,
    /// The correction vector applied
    pub correction: SVector<f64, V>,
    /// The kind of correction (position or velocity)
    pub variables: [Variable; V],
    /// The errors achieved
    pub achieved_errors: SVector<f64, O>,
    /// The objectives set in the targeter
    pub achieved_objectives: [Objective; O],
    /// The number of iterations required
    pub iterations: usize,
    /// Computation duration
    pub computation_dur: Duration,
}

impl<const V: usize, const O: usize> TargeterSolution<V, O> {
    /// Returns whether this solution is a finite burn solution or not
    pub fn is_finite_burn(&self) -> bool {
        for var in &self.variables {
            if var.component.is_finite_burn() {
                return true;
            }
        }
        false
    }

    /// Returns a maneuver if targeter solution was a finite burn maneuver
    pub fn to_mnvr(&self) -> Result<Maneuver, TargetingError> {
        ensure!(self.is_finite_burn(), NotFiniteSnafu);

        let correction_epoch = self.corrected_state.epoch();
        let achievement_epoch = self.achieved_state.epoch();
        let mut mnvr = Maneuver {
            start: correction_epoch,
            end: achievement_epoch,
            thrust_prct: 1.0,
            representation: MnvrRepr::Angles {
                azimuth: CommonPolynomial::Quadratic {
                    a: 0.0,
                    b: 0.0,
                    c: 0.0,
                },
                elevation: CommonPolynomial::Quadratic {
                    a: 0.0,
                    b: 0.0,
                    c: 0.0,
                },
            },
            frame: LocalFrame::RCN,
        };

        for (i, var) in self.variables.iter().enumerate() {
            let corr = self.correction[i];

            // Modify the maneuver, but do not change the epochs of the maneuver unless the change is greater than one millisecond
            match var.component {
                Vary::Duration => {
                    if corr.abs() > 1e-3 {
                        // Check that we are within the bounds
                        let init_duration_s = (correction_epoch - achievement_epoch).to_seconds();
                        let acceptable_corr = var.apply_bounds(init_duration_s).seconds();
                        mnvr.end = mnvr.start + acceptable_corr;
                    }
                }
                Vary::EndEpoch => {
                    if corr.abs() > 1e-3 {
                        // Check that we are within the bounds
                        let total_end_corr =
                            (mnvr.end + corr.seconds() - achievement_epoch).to_seconds();
                        let acceptable_corr = var.apply_bounds(total_end_corr).seconds();
                        mnvr.end += acceptable_corr;
                    }
                }
                Vary::StartEpoch => {
                    if corr.abs() > 1e-3 {
                        // Check that we are within the bounds
                        let total_start_corr =
                            (mnvr.start + corr.seconds() - correction_epoch).to_seconds();
                        let acceptable_corr = var.apply_bounds(total_start_corr).seconds();
                        mnvr.end += acceptable_corr;

                        mnvr.start += corr.seconds()
                    }
                }
                Vary::MnvrAlpha | Vary::MnvrAlphaDot | Vary::MnvrAlphaDDot => {
                    match mnvr.representation {
                        MnvrRepr::Angles { azimuth, elevation } => {
                            let azimuth = azimuth
                                .add_val_in_order(corr, var.component.vec_index())
                                .unwrap();
                            mnvr.representation = MnvrRepr::Angles { azimuth, elevation };
                        }
                        _ => unreachable!(),
                    };
                }
                Vary::MnvrDelta | Vary::MnvrDeltaDot | Vary::MnvrDeltaDDot => {
                    match mnvr.representation {
                        MnvrRepr::Angles { azimuth, elevation } => {
                            let elevation = elevation
                                .add_val_in_order(corr, var.component.vec_index())
                                .unwrap();
                            mnvr.representation = MnvrRepr::Angles { azimuth, elevation };
                        }
                        _ => unreachable!(),
                    };
                }
                Vary::ThrustX | Vary::ThrustY | Vary::ThrustZ => {
                    let mut vector = mnvr.direction();
                    vector[var.component.vec_index()] += corr;
                    var.ensure_bounds(&mut vector[var.component.vec_index()]);
                    mnvr.set_direction(vector).context(GuidanceSnafu)?;
                }
                Vary::ThrustRateX | Vary::ThrustRateY | Vary::ThrustRateZ => {
                    let mut vector = mnvr.rate();
                    let idx = (var.component.vec_index() - 1) % 3;
                    vector[idx] += corr;
                    var.ensure_bounds(&mut vector[idx]);
                    mnvr.set_rate(vector).context(GuidanceSnafu)?;
                }
                Vary::ThrustAccelX | Vary::ThrustAccelY | Vary::ThrustAccelZ => {
                    let mut vector = mnvr.accel();
                    let idx = (var.component.vec_index() - 1) % 3;
                    vector[idx] += corr;
                    var.ensure_bounds(&mut vector[idx]);
                    mnvr.set_accel(vector).context(GuidanceSnafu)?;
                }
                Vary::ThrustLevel => {
                    mnvr.thrust_prct += corr;
                    var.ensure_bounds(&mut mnvr.thrust_prct);
                }
                _ => unreachable!(),
            }
        }

        Ok(mnvr)
    }
}

impl<const V: usize, const O: usize> fmt::Display for TargeterSolution<V, O> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut objmsg = String::from("");
        for (i, obj) in self.achieved_objectives.iter().enumerate() {
            objmsg.push_str(&format!(
                "\n\t\t{:?} = {:.3} (wanted {:.3} ± {:.1e})",
                obj.parameter,
                obj.desired_value + self.achieved_errors[i],
                obj.desired_value,
                obj.tolerance
            ));
        }

        let mut corrmsg = format!("Correction @ {}:", self.corrected_state.epoch());
        let mut is_only_position = true;
        let mut is_only_velocity = true;
        for (i, var) in self.variables.iter().enumerate() {
            let (unit, mulp) = match var.component {
                Vary::PositionX | Vary::PositionY | Vary::PositionZ => {
                    is_only_velocity = false;
                    ("m", 1e3)
                }
                Vary::VelocityX | Vary::VelocityY | Vary::VelocityZ => {
                    is_only_position = false;
                    ("m/s", 1e3)
                }
                _ => {
                    is_only_position = false;
                    is_only_velocity = false;
                    ("", 1.0)
                }
            };
            corrmsg.push_str(&format!(
                "\n\t\t{:?} = {:.3} {}",
                var.component,
                self.correction[i] * mulp,
                unit
            ));
        }

        if is_only_position {
            corrmsg.push_str(&format!(
                "\n\t\t|Δr| = {:.3} m",
                self.correction.norm() * 1e3
            ));
        } else if is_only_velocity {
            corrmsg.push_str(&format!(
                "\n\t\t|Δv| = {:.3} m/s",
                self.correction.norm() * 1e3
            ));
        } else if self.is_finite_burn() {
            let mnvr = self.to_mnvr().unwrap();
            corrmsg.push_str(&format!("\n\t\t{mnvr}\n"));
        }

        writeln!(
            f,
            "Targeter solution correcting {:?} (converged in {:.3} seconds, {} iterations):\n\t{}\n\tAchieved @ {}:{}\n\tCorrected state:\n\t\t{}\n\t\t{:x}\n\tAchieved state:\n\t\t{}\n\t\t{:x}",
            self.variables
                .iter()
                .map(|v| format!("{:?}", v.component))
                .collect::<Vec<String>>(),
            self.computation_dur.as_secs_f64(),
            self.iterations,
            corrmsg,
            self.achieved_state.epoch(),
            objmsg,
            self.corrected_state,
            self.corrected_state,
            self.achieved_state,
            self.achieved_state
        )
    }
}
