nyx-space 2.4.0

/*
    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 anise::almanac::Almanac;
use anise::analysis::AnalysisError;
use anise::structure::LocationDataSet;
use hifitime::{Duration, Epoch, TimeSeries, TimeUnits};
use log::{info, warn};
use num::integer::gcd;
use rand::SeedableRng;
use rand::rngs::SysRng;
use rand_pcg::Pcg64Mcg;

use crate::Spacecraft;
use crate::State;
use crate::io::ConfigError;
use crate::md::trajectory::Interpolatable;
use crate::od::GroundStation;
use crate::od::msr::TrackingDataArc;
use crate::od::prelude::Strand;
use crate::od::simulator::Cadence;
use crate::{linalg::DefaultAllocator, md::prelude::Traj};
use crate::{linalg::allocator::Allocator, od::TrackingDevice};
use std::collections::BTreeMap;
use std::fmt::Display;
use std::marker::PhantomData;

use super::{Handoff, TrkConfig};

#[derive(Clone)]
pub struct TrackingArcSim<MsrIn, D>
where
    D: TrackingDevice<MsrIn>,
    MsrIn: State,
    MsrIn: Interpolatable,
    DefaultAllocator: Allocator<<MsrIn as State>::Size>
        + Allocator<<MsrIn as State>::Size, <MsrIn as State>::Size>
        + Allocator<<MsrIn as State>::VecLength>,
{
    /// Map of devices from their names.
    pub devices: BTreeMap<String, D>,
    /// Receiver trajectory
    pub trajectory: Traj<MsrIn>,
    /// Configuration of each device
    pub configs: BTreeMap<String, TrkConfig>,
    /// Random number generator used for this tracking arc, ensures repeatability
    rng: Pcg64Mcg,
    /// Greatest common denominator time series that allows this arc to meet all of the conditions.
    time_series: TimeSeries,
    _msr_in: PhantomData<MsrIn>,
}

impl<MsrIn, D> TrackingArcSim<MsrIn, D>
where
    D: TrackingDevice<MsrIn>,
    MsrIn: State,
    MsrIn: Interpolatable,
    DefaultAllocator: Allocator<<MsrIn as State>::Size>
        + Allocator<<MsrIn as State>::Size, <MsrIn as State>::Size>
        + Allocator<<MsrIn as State>::VecLength>,
{
    /// Constructs a deterministic tracking simulator initialized with a specific random number generator.
    ///
    /// Evaluates the configuration of all provided devices, filtering out invalid definitions via sanity checks.
    /// Computes the greatest common denominator (GCD) of all configured sampling rates to construct a unified,
    /// high-fidelity time series encompassing the entire trajectory span. This ensures that interpolation artifacts
    /// are minimized when extracting state vectors at non-uniform station sampling intervals.
    pub fn with_rng(
        devices: BTreeMap<String, D>,
        trajectory: Traj<MsrIn>,
        configs: BTreeMap<String, TrkConfig>,
        rng: Pcg64Mcg,
    ) -> Result<Self, ConfigError> {
        // Check that each device has an associated configurations.
        // We don't care if there are more configurations than chosen devices.
        let mut sampling_rates_ns = Vec::with_capacity(devices.len());
        for name in devices.keys() {
            if let Some(cfg) = configs.get(name) {
                if let Err(e) = cfg.sanity_check() {
                    warn!("Ignoring device {name}: {e}");
                    continue;
                }
                sampling_rates_ns.push(cfg.sampling.truncated_nanoseconds());
            } else {
                warn!("Ignoring device {name}: no associated tracking configuration",);
                continue;
            }
        }

        if sampling_rates_ns.is_empty() {
            return Err(ConfigError::InvalidConfig {
                msg: "None of the devices are properly configured".to_string(),
            });
        }

        let common_sampling_rate_ns = sampling_rates_ns
            .iter()
            .fold(sampling_rates_ns[0], |a, &b| gcd(a, b));

        // The overall time series is the one going from the start to the end of the trajectory with the smallest time step
        // of all the tracking configurations.
        let time_series = TimeSeries::inclusive(
            trajectory.first().epoch(),
            trajectory.last().epoch(),
            Duration::from_truncated_nanoseconds(common_sampling_rate_ns),
        );

        let me = Self {
            devices,
            trajectory,
            configs,
            rng,
            time_series,
            _msr_in: PhantomData,
        };

        info!("{me}");

        Ok(me)
    }

    /// Build a new tracking arc simulator using the provided seed to initialize the random number generator.
    pub fn with_seed(
        devices: BTreeMap<String, D>,
        trajectory: Traj<MsrIn>,
        configs: BTreeMap<String, TrkConfig>,
        seed: u64,
    ) -> Result<Self, ConfigError> {
        let rng = Pcg64Mcg::new(seed as u128);

        Self::with_rng(devices, trajectory, configs, rng)
    }

    /// Build a new tracking arc simulator using the system entropy to seed the random number generator.
    pub fn new(
        devices: BTreeMap<String, D>,
        trajectory: Traj<MsrIn>,
        configs: BTreeMap<String, TrkConfig>,
    ) -> Result<Self, ConfigError> {
        let rng = Pcg64Mcg::try_from_rng(&mut SysRng).unwrap();

        Self::with_rng(devices, trajectory, configs, rng)
    }

    /// Generates measurements for the tracking arc using the defined strands
    ///
    /// # Warning
    /// This function will return an error if any of the devices defines as a scheduler.
    /// You must create the schedule first using `build_schedule` first.
    ///
    /// # Notes
    /// Although mutable, this function may be called several times to generate different measurements.
    ///
    /// # Algorithm
    /// For each tracking device, and for each strand within that device, sample the trajectory at the sample
    /// rate of the tracking device, adding a measurement whenever the spacecraft is visible.
    /// Build the measurements as a vector, ordered chronologically.
    ///
    pub fn generate_measurements(
        &mut self,
        almanac: &Almanac,
    ) -> Result<TrackingDataArc, ConfigError> {
        let mut measurements = Vec::new();

        for (name, device) in self.devices.iter_mut() {
            if let Some(cfg) = self.configs.get(name) {
                if cfg.scheduler.is_some() {
                    if cfg.strands.is_none() {
                        return Err(ConfigError::InvalidConfig {
                            msg: format!(
                                "schedule for {name} must be built before generating measurements"
                            ),
                        });
                    } else {
                        warn!(
                            "scheduler for {name} is ignored, using the defined tracking strands instead"
                        )
                    }
                }

                let init_msr_count = measurements.len();
                let tick = Epoch::now().unwrap();

                match cfg.strands.as_ref() {
                    Some(strands) => {
                        // Strands are defined at this point
                        'strands: for (ii, strand) in strands.iter().enumerate() {
                            // Build the time series for this strand, sampling at the correct rate
                            for epoch in
                                TimeSeries::inclusive(strand.start, strand.end, cfg.sampling)
                            {
                                match device.measure(
                                    epoch,
                                    &self.trajectory,
                                    Some(&mut self.rng),
                                    almanac,
                                ) {
                                    Ok(msr_opt) => {
                                        if let Some(msr) = msr_opt {
                                            measurements.push(msr);
                                        }
                                    }
                                    Err(e) => {
                                        if epoch != strand.end {
                                            warn!(
                                                "Skipping the remaining strand #{ii} ending on {}: {e}",
                                                strand.end
                                            );
                                        }
                                        continue 'strands;
                                    }
                                }
                            }
                        }

                        info!(
                            "Simulated {} measurements for {name} for {} tracking strands in {}",
                            measurements.len() - init_msr_count,
                            strands.len(),
                            (Epoch::now().unwrap() - tick).round(1.0_f64.milliseconds())
                        );
                    }
                    None => {
                        warn!("No tracking strands defined for {name}, skipping");
                    }
                }
            }
        }

        // Build the tracking arc.
        let mut trk_data = TrackingDataArc {
            measurements,
            source: None,
            moduli: None,
            force_reject: false,
        };

        trk_data.sort();

        Ok(trk_data)
    }
}

impl<MsrIn, D> Display for TrackingArcSim<MsrIn, D>
where
    D: TrackingDevice<MsrIn>,
    MsrIn: Interpolatable,
    DefaultAllocator: Allocator<<MsrIn as State>::Size>
        + Allocator<<MsrIn as State>::Size, <MsrIn as State>::Size>
        + Allocator<<MsrIn as State>::VecLength>,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Tracking Arc Simulator on {} with devices {:?} over {}",
            self.trajectory,
            self.devices.keys(),
            self.time_series
        )
    }
}

// Literally the same as above, but can't make it generic =(
impl TrackingArcSim<Spacecraft, GroundStation> {
    /// Builds the schedule provided the config. Requires the tracker to be a ground station.
    ///
    /// # Algorithm
    ///
    /// 1. For each tracking device:
    /// 2. Find when the vehicle trajectory has an elevation greater or equal to zero, and use that as the first start of the first tracking arc for this station
    /// 3. Find when the vehicle trajectory has an elevation less than zero (i.e. disappears below the horizon), after that initial epoch
    /// 4. Repeat 2, 3 until the end of the trajectory
    /// 5. Build each of these as "tracking strands" for this tracking device.
    /// 6. Organize all of the built tracking strands chronologically.
    /// 7. Iterate through all of the strands:
    ///    7.a. if that tracker is marked as `Greedy` and it ends after the start of the next strand, change the start date of the next strand.
    ///    7.b. if that tracker is marked as `Eager` and it ends after the start of the next strand, change the end date of the current strand.
    pub fn generate_schedule(
        &self,
        almanac: &Almanac,
    ) -> Result<BTreeMap<String, TrkConfig>, AnalysisError> {
        // Build a Location Dataset so we can use ANISE's visibility finder.
        let mut loc_dataset = LocationDataSet::default();
        let mut loc_ids = Vec::with_capacity(self.devices.len());
        for (dno, (name, device)) in self.devices.iter().enumerate() {
            let loc_id = dno as i32 + 1_000;
            loc_dataset
                .push(device.location.clone(), Some(loc_id), Some(name))
                .map_err(|e| AnalysisError::GenericAnalysisError {
                    err: format!("{e}"),
                })?;
            loc_ids.push((name.clone(), loc_id));
        }

        // Deep clone of the Almanac so we can insert the locations of these ground stations.
        let almanac = (*almanac).clone();
        let almanac = almanac.with_location_data(loc_dataset);

        let mut built_cfg = self.configs.clone();

        let traj = &self.trajectory;

        for (name, loc_id) in &loc_ids {
            if let Some(cfg) = self.configs.get(name)
                && let Some(scheduler) = cfg.scheduler
            {
                info!("Building schedule for {name}");
                if scheduler.handoff == Handoff::Overlap {
                    warn!("Overlap unsupported for {name} on identical epochs.");
                }
                built_cfg.get_mut(name).unwrap().scheduler = None;
                built_cfg.get_mut(name).unwrap().strands = Some(Vec::new());

                let visibilty_arcs = almanac.report_visibility_arcs(
                    traj,
                    *loc_id,
                    traj.first().epoch(),
                    traj.last().epoch(),
                    cfg.sampling,
                    None,
                )?;

                for arc in visibilty_arcs {
                    let strand_start = arc.rise.orbit.epoch;
                    let strand_end = arc.fall.orbit.epoch;

                    if strand_end - strand_start < cfg.sampling * i64::from(scheduler.min_samples) {
                        warn!(
                            "Dropped [{strand_start}, {strand_end}]: < {} samples @ {}",
                            scheduler.min_samples, cfg.sampling
                        );
                        continue;
                    }

                    let mut strand_range = Strand {
                        start: strand_start,
                        end: strand_end,
                    };

                    // If there is an alignment, apply it
                    if let Some(alignment) = scheduler.sample_alignment {
                        strand_range.start = strand_range.start.round(alignment);
                        strand_range.end = strand_range.end.round(alignment);
                    }

                    if let Cadence::Intermittent { on, off } = scheduler.cadence {
                        // Check that the next start time is within the allocated time
                        if let Some(prev_strand) = built_cfg[name].strands.as_ref().unwrap().last()
                            && prev_strand.end + off > strand_range.start
                        {
                            // We're turning on the tracking sooner than the schedule allows, so let's fix that.
                            strand_range.start = prev_strand.end + off;
                            // Check that we didn't eat into the whole tracking opportunity
                            if strand_range.start > strand_end {
                                // Lost this whole opportunity.
                                warn!(
                                    "Dropped {name} [{strand_start}, {strand_end}]: cadence {:?}",
                                    scheduler.cadence
                                );
                                continue;
                            }
                        }
                        // Check that we aren't tracking for longer than configured
                        if strand_range.end - strand_range.start > on {
                            strand_range.end = strand_range.start + on;
                        }
                    }

                    // We've found when the spacecraft is below the horizon, so this is a new strand.
                    built_cfg
                        .get_mut(name)
                        .unwrap()
                        .strands
                        .as_mut()
                        .unwrap()
                        .push(strand_range);
                }

                info!(
                    "Built {} tracking strands for {name}",
                    built_cfg[name].strands.as_ref().unwrap().len()
                );
            }
        }
        // Build all of the strands, remembering which tracker they come from.
        let mut cfg_as_vec = Vec::new();
        for (name, cfg) in &built_cfg {
            if let Some(strands) = &cfg.strands {
                for (ii, strand) in strands.iter().enumerate() {
                    cfg_as_vec.push((name.clone(), ii, *strand));
                }
            }
        }
        // Iterate through the strands by chronological order. Cannot use maps because we change types.
        cfg_as_vec.sort_by_key(|(_, _, strand)| strand.start);
        for (ii, (this_name, this_pos, this_strand)) in
            cfg_as_vec.iter().take(cfg_as_vec.len() - 1).enumerate()
        {
            // Grab the config
            if let Some(config) = self.configs[this_name].scheduler.as_ref() {
                // Grab the next strand, chronologically
                if let Some((next_name, next_pos, next_strand)) = cfg_as_vec.get(ii + 1) {
                    if config.handoff == Handoff::Greedy && this_strand.end >= next_strand.start {
                        // Modify the built configurations to change the start time of the next strand because the current one is greedy.
                        let next_config = built_cfg.get_mut(next_name).unwrap();
                        let new_start = this_strand.end + next_config.sampling;
                        next_config.strands.as_mut().unwrap()[*next_pos].start = new_start;
                        info!("Greedy handoff for {this_name}: {next_name} delayed to {new_start}");
                    } else if config.handoff == Handoff::Eager
                        && this_strand.end >= next_strand.start
                    {
                        let this_config = built_cfg.get_mut(this_name).unwrap();
                        let new_end = next_strand.start - this_config.sampling;
                        this_config.strands.as_mut().unwrap()[*this_pos].end = new_end;
                        info!("Eager handoff for {this_name}: {this_name} terminated at {new_end}");
                    }
                } else {
                    // Reached the end
                    break;
                }
            }
        }

        Ok(built_cfg)
    }

    /// Sets the schedule to that built in `build_schedule`
    pub fn build_schedule(&mut self, almanac: &Almanac) -> Result<(), AnalysisError> {
        self.configs = self.generate_schedule(almanac)?;

        Ok(())
    }
}