/*
    Nyx, blazing fast astrodynamics
    Copyright (C) 2023 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 std::collections::{HashMap, HashSet};
use std::error::Error;
use std::fmt::{Debug, Display};
use std::fs::File;
use std::ops::RangeBounds;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use crate::cosmic::Cosm;
use crate::io::watermark::pq_writer;
use crate::io::{ConfigError, ConfigRepr, ExportCfg};
use crate::linalg::allocator::Allocator;
use crate::linalg::{DefaultAllocator, DimName};
use crate::md::trajectory::Interpolatable;
use crate::od::{Measurement, TrackingDeviceSim};
use crate::State;
use arrow::array::{Array, Float64Builder, StringBuilder};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use hifitime::prelude::{Duration, Epoch};
use parquet::arrow::ArrowWriter;

/// Tracking arc contains the tracking data generated by the tracking devices defined in this structure.
/// This structure is shared between both simulated and real tracking arcs.
#[derive(Clone, Default, Debug)]
pub struct TrackingArc<Msr>
where
    Msr: Measurement,
    DefaultAllocator: Allocator<f64, Msr::MeasurementSize>,
{
    /// The YAML configuration to set up these devices
    pub device_cfg: String,
    /// A chronological list of the measurements to the devices used to generate these measurements. If the name of the device does not appear in the list of devices, then the measurement will be ignored.
    pub measurements: Vec<(String, Msr)>,
}

impl<Msr> Display for TrackingArc<Msr>
where
    Msr: Measurement,
    DefaultAllocator: Allocator<f64, Msr::MeasurementSize>,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{} measurements from {:?}",
            self.measurements.len(),
            self.device_names()
        )
    }
}

impl<Msr> TrackingArc<Msr>
where
    Msr: Measurement,
    DefaultAllocator: Allocator<f64, Msr::MeasurementSize>,
{
    /// Store this tracking arc to a parquet file.
    pub fn to_parquet_simple<P: AsRef<Path> + Debug>(
        &self,
        path: P,
    ) -> Result<PathBuf, Box<dyn Error>> {
        self.to_parquet(path, ExportCfg::default())
    }

    /// Store this tracking arc to a parquet file, with optional metadata and a timestamp appended to the filename.
    pub fn to_parquet<P: AsRef<Path> + Debug>(
        &self,
        path: P,
        cfg: ExportCfg,
    ) -> Result<PathBuf, Box<dyn Error>> {
        let path_buf = cfg.actual_path(path);

        if cfg.step.is_some() {
            warn!("The `step` parameter in the export is not supported for tracking arcs.");
        }

        if cfg.fields.is_some() {
            warn!("The `fields` parameter in the export is not supported for tracking arcs.");
        }

        // Build the schema
        let mut hdrs = vec![
            Field::new("Epoch:Gregorian UTC", DataType::Utf8, false),
            Field::new("Epoch:Gregorian TAI", DataType::Utf8, false),
            Field::new("Epoch:TAI (s)", DataType::Float64, false),
            Field::new("Tracking device", DataType::Utf8, false),
        ];

        let mut msr_fields = Msr::fields();

        hdrs.append(&mut msr_fields);

        // Build the schema
        let schema = Arc::new(Schema::new(hdrs));
        let mut record: Vec<Arc<dyn Array>> = Vec::new();

        // Build the measurement iterator

        let measurements =
            if cfg.start_epoch.is_some() || cfg.end_epoch.is_some() || cfg.step.is_some() {
                let start = cfg
                    .start_epoch
                    .unwrap_or_else(|| self.measurements.first().unwrap().1.epoch());
                let end = cfg
                    .end_epoch
                    .unwrap_or_else(|| self.measurements.last().unwrap().1.epoch());

                info!("Exporting measurements from {start} to {end}.");

                self.measurements
                    .iter()
                    .filter(|msr| msr.1.epoch() >= start && msr.1.epoch() <= end)
                    .cloned()
                    .collect()
            } else {
                self.measurements.to_vec()
            };

        // Build all of the records

        // Epochs
        let mut utc_epoch = StringBuilder::new();
        let mut tai_epoch = StringBuilder::new();
        let mut tai_s = Float64Builder::new();
        for m in &measurements {
            utc_epoch.append_value(format!("{}", m.1.epoch()));
            tai_epoch.append_value(format!("{:x}", m.1.epoch()));
            tai_s.append_value(m.1.epoch().to_tai_seconds());
        }
        record.push(Arc::new(utc_epoch.finish()));
        record.push(Arc::new(tai_epoch.finish()));
        record.push(Arc::new(tai_s.finish()));

        // Device names
        let mut device_names = StringBuilder::new();
        for m in &measurements {
            device_names.append_value(m.0.clone());
        }
        record.push(Arc::new(device_names.finish()));

        // Measurement data
        for obs_no in 0..Msr::MeasurementSize::USIZE {
            let mut data_builder = Float64Builder::new();

            for m in &measurements {
                data_builder.append_value(m.1.observation()[obs_no]);
            }
            record.push(Arc::new(data_builder.finish()));
        }

        // Serialize all of the devices and add that to the parquet file too.
        let mut metadata = HashMap::new();
        metadata.insert("devices".to_string(), self.device_cfg.clone());
        metadata.insert("Purpose".to_string(), "Tracking Arc Data".to_string());
        if let Some(add_meta) = cfg.metadata {
            for (k, v) in add_meta {
                metadata.insert(k, v);
            }
        }

        let props = pq_writer(Some(metadata));

        let file = File::create(&path_buf)?;

        let mut writer = ArrowWriter::try_new(file, schema.clone(), props).unwrap();

        let batch = RecordBatch::try_new(schema, record)?;
        writer.write(&batch)?;
        writer.close()?;

        info!("Serialized {self} to {}", path_buf.display());

        // Return the path this was written to
        Ok(path_buf)
    }

    /// Returns the set of devices from which measurements were taken. This accounts for the availability of measurements, so if a device was not available, it will not appear in this set.
    pub fn device_names(&self) -> HashSet<&String> {
        let mut set = HashSet::new();
        self.measurements.iter().for_each(|(name, _msr)| {
            set.insert(name);
        });
        set
    }

    /// Returns the minimum duration between two subsequent measurements.
    /// This is important to correctly set up the propagator and not miss any measurement.
    pub fn min_duration_sep(&self) -> Option<Duration> {
        let mut windows = self.measurements.windows(2);
        let first_window = windows.next()?;
        let mut min_interval = first_window[1].1.epoch() - first_window[0].1.epoch();
        for window in windows {
            let interval = window[1].1.epoch() - window[0].1.epoch();
            if interval != Duration::ZERO && interval < min_interval {
                min_interval = interval;
            }
        }

        Some(min_interval)
    }

    /// If this tracking arc has devices that can be used to generate simulated measurements,
    /// then this function can be used to rebuild said measurement devices
    pub fn rebuild_devices<MsrIn, D>(
        &self,
        cosm: Arc<Cosm>,
    ) -> Result<HashMap<String, D>, ConfigError>
    where
        MsrIn: Interpolatable,
        D: TrackingDeviceSim<MsrIn, Msr>,
        DefaultAllocator: Allocator<f64, <MsrIn as State>::Size>
            + Allocator<f64, <MsrIn as State>::Size, <MsrIn as State>::Size>
            + Allocator<f64, <MsrIn as State>::VecLength>,
    {
        let devices_repr = D::IntermediateRepr::loads_many(&self.device_cfg)?;

        let mut devices = HashMap::new();

        for serde in devices_repr {
            let device = D::from_config(serde, cosm.clone())?;
            if !self.device_names().contains(&device.name()) {
                warn!(
                    "{} from arc config does not appear in measurements -- ignored",
                    device.name()
                );
                continue;
            }
            devices.insert(device.name(), device);
        }

        Ok(devices)
    }

    /// Returns a new tracking arc that only contains measurements that fall within the given epoch range.
    pub fn filter_by_epoch<R: RangeBounds<Epoch>>(&self, bound: R) -> Self {
        let mut measurements = Vec::new();
        for (name, msr) in &self.measurements {
            if bound.contains(&msr.epoch()) {
                measurements.push((name.clone(), *msr));
            }
        }

        Self {
            measurements,
            device_cfg: self.device_cfg.clone(),
        }
    }

    /// Returns a new tracking arc that only contains measurements that fall within the given offset from the first epoch
    pub fn filter_by_offset<R: RangeBounds<Duration>>(&self, bound: R) -> Self {
        if self.measurements.is_empty() {
            return Self {
                device_cfg: self.device_cfg.clone(),
                measurements: Vec::new(),
            };
        }
        let ref_epoch = self.measurements[0].1.epoch();
        let mut measurements = Vec::new();
        for (name, msr) in &self.measurements {
            if bound.contains(&(msr.epoch() - ref_epoch)) {
                measurements.push((name.clone(), *msr));
            }
        }

        Self {
            measurements,
            device_cfg: self.device_cfg.clone(),
        }
    }
}
