/*
    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 super::{ODAlmanacSnafu, ODError, ODTrajSnafu, TrackingDevice};
use crate::Spacecraft;
use crate::md::prelude::{Interpolatable, Traj};
use crate::od::msr::MeasurementType;
use crate::od::msr::measurement::Measurement;
use crate::time::Epoch;
use anise::errors::AlmanacResult;
use anise::frames::Frame;
use anise::prelude::{Aberration, Almanac, Orbit};
use hifitime::TimeUnits;
use indexmap::IndexSet;
use log::debug;
use rand_pcg::Pcg64Mcg;
use snafu::ResultExt;

use super::GroundStation;

impl TrackingDevice<Spacecraft> for GroundStation {
    fn measurement_types(&self) -> &IndexSet<MeasurementType> {
        &self.measurement_types
    }

    /// Perform a measurement from the ground station to the receiver (rx).
    fn measure(
        &mut self,
        epoch: Epoch,
        traj: &Traj<Spacecraft>,
        rng: Option<&mut Pcg64Mcg>,
        almanac: &Almanac,
    ) -> Result<Option<Measurement>, ODError> {
        match self.integration_time {
            Some(integration_time) => {
                // TODO: This should support measurement alignment
                // If out of traj bounds, return None, else the whole strand is rejected.
                let rx_0 = match traj.at(epoch - integration_time).context(ODTrajSnafu {
                    details: format!(
                        "fetching state {epoch} at start of ground station integration time {integration_time}"
                    ),
                }) {
                    Ok(rx) => rx,
                    Err(_) => return Ok(None),
                };

                let rx_1 = match traj.at(epoch).context(ODTrajSnafu {
                    details: format!(
                        "fetching state {epoch} at end of ground station integration time"
                    ),
                }) {
                    Ok(rx) => rx,
                    Err(_) => return Ok(None),
                };

                let obstructing_body =
                    if self.location.frame.ephemeris_id != rx_0.frame().ephemeris_id {
                        Some(rx_0.frame())
                    } else {
                        None
                    };

                let ab_corr = if self.light_time_correction {
                    Aberration::LT
                } else {
                    Aberration::NONE
                };

                let aer_t0 = almanac
                    .azimuth_elevation_range_sez_from_location(
                        rx_0.orbit,
                        self.location.clone(),
                        obstructing_body,
                        ab_corr,
                    )
                    .context(ODAlmanacSnafu {
                        action: "computing AER",
                    })?;

                let aer_t1 = almanac
                    .azimuth_elevation_range_sez_from_location(
                        rx_1.orbit,
                        self.location.clone(),
                        obstructing_body,
                        ab_corr,
                    )
                    .context(ODAlmanacSnafu {
                        action: "computing AER",
                    })?;

                if aer_t0.elevation_above_mask_deg() < 0.0
                    || aer_t1.elevation_above_mask_deg() < 0.0
                {
                    debug!(
                        "{} {} obstructed by terrain ({:.3} - {:.3} deg) -- no measurement",
                        self.name,
                        aer_t0.epoch,
                        aer_t0.elevation_above_mask_deg(),
                        aer_t1.elevation_above_mask_deg()
                    );
                    return Ok(None);
                } else if aer_t0.is_obstructed() || aer_t1.is_obstructed() {
                    debug!(
                        "{} {} obstruction at t0={}, t1={} -- no measurement",
                        self.name,
                        aer_t0.epoch,
                        aer_t0.is_obstructed(),
                        aer_t1.is_obstructed()
                    );
                    return Ok(None);
                }

                // Noises are computed at the midpoint of the integration time.
                let noises = self.noises(epoch - integration_time * 0.5, rng)?;

                let mut msr = Measurement::new(self.name.clone(), epoch + noises[0].seconds());

                for (ii, msr_type) in self.measurement_types.iter().enumerate() {
                    let msr_value = msr_type.compute_two_way(aer_t0, aer_t1, noises[ii + 1])?;
                    msr.push(*msr_type, msr_value);
                }

                Ok(Some(msr))
            }
            None => self.measure_instantaneous(
                traj.at(epoch).context(ODTrajSnafu {
                    details: "fetching state for instantaneous measurement".to_string(),
                })?,
                rng,
                almanac,
            ),
        }
    }

    fn name(&self) -> String {
        self.name.clone()
    }

    fn location(&self, epoch: Epoch, frame: Frame, almanac: &Almanac) -> AlmanacResult<Orbit> {
        almanac.transform_to(self.to_orbit(epoch, almanac).unwrap(), frame, None)
    }

    fn measure_instantaneous(
        &mut self,
        rx: Spacecraft,
        rng: Option<&mut Pcg64Mcg>,
        almanac: &Almanac,
    ) -> Result<Option<Measurement>, ODError> {
        let obstructing_body = if self.location.frame.ephemeris_id != rx.frame().ephemeris_id {
            Some(rx.frame())
        } else {
            None
        };

        let ab_corr = if self.light_time_correction {
            Aberration::LT
        } else {
            Aberration::NONE
        };

        let aer = almanac
            .azimuth_elevation_range_sez_from_location(
                rx.orbit,
                self.location.clone(),
                obstructing_body,
                ab_corr,
            )
            .context(ODAlmanacSnafu {
                action: "computing AER",
            })?;

        if aer.elevation_above_mask_deg() >= 0.0 && !aer.is_obstructed() {
            // Only update the noises if the measurement is valid.
            let noises = self.noises(rx.orbit.epoch, rng)?;

            let mut msr = Measurement::new(self.name.clone(), rx.orbit.epoch + noises[0].seconds());

            for (ii, msr_type) in self.measurement_types.iter().enumerate() {
                let msr_value = msr_type.compute_one_way(aer, noises[ii + 1])?;
                msr.push(*msr_type, msr_value);
            }

            Ok(Some(msr))
        } else {
            debug!(
                "{} {} object at {:.3} deg -- no measurement",
                self.name,
                rx.orbit.epoch,
                aer.elevation_above_mask_deg(),
            );
            Ok(None)
        }
    }

    /// Returns the measurement noise of this ground station.
    ///
    /// # Methodology
    /// Noises are modeled using a [StochasticNoise] process, defined by the sigma on the turn-on bias and on the steady state noise.
    /// The measurement noise is computed assuming that all measurements are independent variables, i.e. the measurement matrix is
    /// a diagonal matrix. The first item in the diagonal is the range noise (in km), set to the square of the steady state sigma. The
    /// second item is the Doppler noise (in km/s), set to the square of the steady state sigma of that Gauss Markov process.
    fn measurement_covar(&self, msr_type: MeasurementType, epoch: Epoch) -> Result<f64, ODError> {
        let stochastics = self.stochastic_noises.as_ref().unwrap();

        Ok(stochastics
            .get(&msr_type)
            .ok_or(ODError::NoiseNotConfigured {
                kind: format!("{msr_type:?}"),
            })?
            .covariance(epoch))
    }

    fn measurement_bias(&self, msr_type: MeasurementType, _epoch: Epoch) -> Result<f64, ODError> {
        let stochastics = self.stochastic_noises.as_ref().unwrap();

        if let Some(gm) = stochastics
            .get(&msr_type)
            .ok_or(ODError::NoiseNotConfigured {
                kind: format!("{msr_type:?}"),
            })?
            .bias
        {
            Ok(gm.constant.unwrap_or(0.0))
        } else {
            Ok(0.0)
        }
    }
}
