eorst 1.0.1

//! Composition methods for RasterDataset.
//!
//! This module contains methods for stacking, extending, and generating column names.

use crate::core_types::RasterType;
use crate::types::Dimension as UtilsDimension;
use crate::rasterdataset::RasterDataset;
use crate::selection::Stack;

use crate::data_sources::DateType;
#[cfg(feature = "use_polars")]
use crate::gdal_utils::create_rayon_pool;
#[cfg(feature = "use_polars")]
use rayon::iter::{IntoParallelIterator, ParallelIterator};

impl<R> RasterDataset<R>
where
    R: RasterType,
{
    /// Stack RasterDataset along a dimension.
    pub fn stack(
        &mut self,
        other: &RasterDataset<R>,
        dimension_to_stack: UtilsDimension,
    ) -> &mut RasterDataset<R> {
        match dimension_to_stack {
            UtilsDimension::Layer => {
                if self.metadata.date_indices != other.metadata.date_indices {
                    panic!(
                        "To stack datasets over Layers, both datasets have to the same time indics"
                    )
                }
                self.metadata
                    .layer_indices
                    .extend_from_slice(&other.metadata.layer_indices)
            }
            UtilsDimension::Time => {
                if self.metadata.layer_indices != other.metadata.layer_indices {
                    panic!("To stack datasets over Time, both have to have the same layers.")
                }
                self.metadata
                    .date_indices
                    .extend_from_slice(&other.metadata.date_indices);
            }
        }

        let axis = dimension_to_stack.get_axis();

        let max_dim_self = self.metadata.shape[axis];
        self.metadata
            .shape
            .stack(other.metadata.shape, dimension_to_stack);

        for other_layer in &other.metadata.layers {
            let mut layer = other_layer.clone();
            layer.stack_position(other_layer.to_owned(), dimension_to_stack, max_dim_self);
            self.metadata.layers.push(layer);
        }

        self
    }

    /// Extend a RasterDataset.
    pub fn extend<V>(&mut self, other: &RasterDataset<V>) -> &mut RasterDataset<R>
    where
        V: RasterType,
    {
        self.metadata.shape.extend(other.metadata.shape);

        for layer in &other.metadata.layers {
            self.metadata.layers.push(layer.to_owned());
        }
        self
    }

    /// Create names like time0_layer0, time0_layer1, time1_layer0, etc.
    pub fn column_names(&self) -> Vec<String> {
        let mut col_names = Vec::new();
        for time in &self.metadata.date_indices {
            for layer in &self.metadata.layer_indices {
                let time_str = match time {
                    DateType::Date(date) => format!("{}", date.format("%Y-%m-%d %H:%M:%S UTC%Z")),

                    DateType::Index(index) => format!("time_{}", index),
                };
                let c = format!("{time_str}--{layer}");
                col_names.push(c);
            }
        }
        col_names
    }

    /// Gets unique values across all blocks.
    #[cfg(feature = "use_polars")]
    pub(crate) fn get_unique_values(&self) -> Vec<i32> {
        let block_to_process: Vec<usize> = (0..self.blocks.len()).collect();
        let pool = create_rayon_pool(1);
        let handle = pool.install(|| {
            block_to_process.into_par_iter().map(|id| {
                let data = self.read_block::<i32>(id);

                let unique_values: std::collections::HashSet<_> = data.into_iter().collect();

                unique_values
            })
        });
        let collected: Vec<_> = handle.collect();
        let mut unique: std::collections::HashSet<_> = std::collections::HashSet::new();
        for set in collected {
            unique.extend(&set);
        }
        let mut unique: Vec<_> = unique.iter().cloned().collect();
        unique.sort();
        unique
    }

    /// An iterator over the [`crate::blocks::RasterBlock`]s of a [`RasterDataset`].
    pub fn iter(&self) -> RasterDatasetIter<'_, R> {
        let iter_index = 0;
        RasterDatasetIter {
            parent: self,
            block: &self.blocks[iter_index],
            iter_index,
        }
    }
}

/// Iterator over RasterDataset blocks.
pub struct RasterDatasetIter<'a, R>
where
    R: RasterType,
{
    /// Reference to the parent dataset
    pub parent: &'a RasterDataset<R>,
    /// Current block
    pub block: &'a crate::blocks::RasterBlock<R>,
    /// Current iteration index
    pub iter_index: usize,
}

impl<'a, R> Iterator for RasterDatasetIter<'a, R>
where
    R: RasterType,
{
    type Item = RasterDatasetIter<'a, R>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.iter_index < self.parent.blocks.len() {
            let current = RasterDatasetIter {
                parent: self.parent,
                block: self.block,
                iter_index: self.iter_index,
            };
            self.iter_index += 1;
            if self.iter_index < self.parent.blocks.len() {
                self.block = &self.parent.blocks[self.iter_index];
            }
            Some(current)
        } else {
            None
        }
    }
}