trs-dataframe 0.10.2

use data_value::DataValue;
use halfbrown::HashMap;
use ndarray::ArrayView1;
use serde::{Deserialize, Serialize};

use crate::{error::Error, Key};

/// [`KeyIndex`] is used to store the keys for the [`super::ColumnFrame`]
/// The keys are stored in the order they are added - the order is preserved
/// The keys are stored in the [`Vec`] and the indexes are stored in the [`HashMap`]
/// The indexes are used to access the data in the [`super::ColumnFrame`] by the column [`Key`]
/// NOTE: The keys are unique - if the key is already present, it will be removed
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct KeyIndex {
    pub keys: Vec<Key>,
    indexes: HashMap<String, usize>,
    pub alias: HashMap<String, String>,
}

impl KeyIndex {
    pub fn new(keys: Vec<Key>) -> Self {
        let mut indexes = HashMap::with_capacity(keys.len());
        let mut removed = 0;
        let mut actual_keys = Vec::with_capacity(keys.len());
        for (idx, key) in keys.into_iter().enumerate() {
            if indexes.contains_key(key.name()) {
                removed += 1;
            } else {
                indexes.insert(key.name().to_string(), idx.saturating_sub(removed));
                actual_keys.push(key)
            }
        }
        Self {
            keys: actual_keys,
            indexes,
            alias: HashMap::new(),
        }
    }

    pub fn len(&self) -> usize {
        self.keys.len()
    }
    pub fn is_empty(&self) -> bool {
        self.keys.is_empty()
    }

    pub fn get_column_index(&self, key: &Key) -> Option<usize> {
        self.get_column_index_by_name(key.name())
    }

    pub fn get_column_index_by_name(&self, key: &str) -> Option<usize> {
        if let Some(f) = self.indexes.get(key) {
            Some(*f)
        } else {
            self.alias
                .get(key)
                .and_then(|alias| self.indexes.get(alias).copied())
        }
    }

    pub fn get_keys(&self) -> &[Key] {
        &self.keys
    }
    pub fn get_keys_mut(&mut self) -> &mut [Key] {
        &mut self.keys
    }

    pub fn get_key(&self, idx: usize) -> Option<Key> {
        self.keys.get(idx).cloned()
    }
    pub fn get_complement_keys(&self, keys: &[Key]) -> Vec<Key> {
        self.keys
            .iter()
            .filter(|key| !keys.contains(key))
            .cloned()
            .collect()
    }

    pub fn select(&self, keys: &[Key]) -> KeyIndex {
        let mut new_keys = Vec::with_capacity(keys.len());
        let mut new_indexes = HashMap::with_capacity(keys.len());

        for key in keys.iter() {
            if let Some(idx) = self.indexes.get(key.name()) {
                new_indexes.insert(key.name().to_string(), *idx);
                new_keys.push(key.to_owned());
            } else if let Some(alias_key) = self.alias.get(key.name()) {
                if let Some(idx) = self.indexes.get(alias_key) {
                    new_indexes.insert(key.name().to_string(), *idx);
                    new_keys.push(key.to_owned());
                }
            }
        }
        Self {
            keys: new_keys,
            indexes: new_indexes,
            alias: HashMap::new(),
        }
    }

    pub fn indexes(&self) -> Vec<usize> {
        self.indexes.values().copied().collect()
    }

    pub fn store_key(&mut self, key: Key) {
        if self.indexes.contains_key(key.name()) {
            return;
        }
        self.keys.push(key.clone());
        self.indexes
            .insert(key.name().to_string(), self.keys.len() - 1);
    }

    pub fn remove_key(&mut self, key: &Key) -> Option<(Key, usize)> {
        let idx = self.indexes.remove(key.name())?;
        let current = self.keys.remove(idx);

        Some((current, idx))
    }

    pub fn get_as_candidate(&self, row: ArrayView1<DataValue>) -> HashMap<Key, DataValue> {
        let mut result = HashMap::with_capacity(self.keys.len());
        for (key, idx) in self.indexes.iter() {
            result.insert(key.into(), row[*idx].clone());
        }
        result
    }

    pub fn to_vec_row(&self, candidate: HashMap<Key, DataValue>) -> Vec<DataValue> {
        self.keys
            .iter()
            .map(|key| candidate.get(key).cloned().unwrap_or_default())
            .collect()
    }

    pub fn check_order_of_indexes(&self, other: &Self) -> Result<(), Error> {
        for (self_key, other_key) in self.keys.iter().zip(other.keys.iter()) {
            if self_key != other_key {
                return Err(Error::IndexOutOfOrder(
                    self.keys.clone(),
                    other.keys.clone(),
                ));
            }
        }
        Ok(())
    }

    pub fn rename_key(&mut self, key: &str, new_key: Key) -> Result<(), Error> {
        if let Some(idx) = self.indexes.remove(key) {
            self.indexes.insert(new_key.to_string(), idx);
            self.keys[idx] = new_key;
            Ok(())
        } else {
            Err(Error::NotFound(key.into()))
        }
    }

    pub fn add_alias(&mut self, key: &str, alias: &str) -> Result<(), Error> {
        if !self.indexes.contains_key(key) {
            return Err(Error::NotFound(key.into()));
        }
        self.alias.insert(alias.to_string(), key.to_string());
        Ok(())
    }
}

impl From<Vec<Key>> for KeyIndex {
    fn from(keys: Vec<Key>) -> Self {
        Self::new(keys)
    }
}

#[cfg(test)]
mod test {
    use crate::DataType;

    use super::*;
    use rstest::*;

    #[rstest]
    fn test_alias() {
        let key = Key::new("a", DataType::U32);
        let mut key_index = KeyIndex::new(vec![key.clone()]);
        assert_eq!(key_index.add_alias(key.name(), "alias"), Ok(()));
        assert!(key_index.add_alias("c", "alias").is_err());
        assert_eq!(key_index.alias.get("alias"), Some(&key.name().to_string()));
        assert_eq!(key_index.get_column_index(&key), Some(0));
        assert_eq!(
            key_index.get_column_index(&Key::new("alias", DataType::U32)),
            Some(0)
        );
    }

    #[rstest]
    fn test_rename() {
        let key = Key::new("a", DataType::U32);
        let mut key_index = KeyIndex::new(vec![key.clone(), Key::new("b", DataType::U32)]);
        assert_eq!(key_index.rename_key("a", "new_key".into()), Ok(()));
        assert_eq!(
            key_index.get_column_index(&Key::new("new_key", DataType::U32)),
            Some(0)
        );
        assert!(key_index.rename_key("c", "alias".into()).is_err());
    }

    #[rstest]
    #[case(
        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)]
    )]
    fn test_key_index_new(#[case] keys: Vec<Key>, #[case] expected: Vec<Key>) {
        let key_index = KeyIndex::new(keys);
        assert_eq!(key_index.keys, expected);
        assert_eq!(key_index.get_keys(), expected);
    }

    #[rstest]
    #[case(
        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
        Ok(())
    )]
    #[case(
        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
        vec![Key::new("b", DataType::U32), Key::new("a", DataType::U32)],
        Err(Error::IndexOutOfOrder(
            vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
            vec![Key::new("b", DataType::U32), Key::new("a", DataType::U32)]
        ))
    )]
    fn test_key_index_check_order_of_indexes(
        #[case] keys: Vec<Key>,
        #[case] other_keys: Vec<Key>,
        #[case] expected: Result<(), Error>,
    ) {
        let key_index: KeyIndex = keys.into();
        let other_key_index = KeyIndex::new(other_keys);
        assert_eq!(key_index.check_order_of_indexes(&other_key_index), expected);
    }

    #[rstest]
    #[case(
        KeyIndex::new(vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)]),
        (Key::new("a", DataType::U32), 0)
    )]
    #[case(
        KeyIndex::new(vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)]),
        (Key::new("b", DataType::U32), 1)
    )]
    fn test_key_index_remove_key(#[case] mut key_index: KeyIndex, #[case] expected: (Key, usize)) {
        let key = expected.0.clone();
        assert_eq!(key_index.remove_key(&key), Some(expected));
    }
}