/*
 * Copyright (c) 2017-2020 Frank Fischer <frank-fischer@shadow-soft.de>
 *
 * This program is free software: you can redistribute it and/or
 * modify it under the terms of the GNU 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
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see  <http://www.gnu.org/licenses/>
 */

//! This module provides vectors that can be indexed by nodes, edges or arcs.

use crate::traits::refs::{GraphSizeRef, IndexGraphRef};

use std::marker::PhantomData;
use std::ops::{Index, IndexMut};

/// A indexer maps items to numeric indices.
pub trait Indexer {
    /// The type of the index items.
    type Idx;
    /// The an iterator over all allowed items.
    type Iter: Iterator<Item = Self::Idx>;

    /// Return the number of items.
    ///
    /// This is the maximal allowed index.
    fn num_items(&self) -> usize;

    /// Return the index associated with an item.
    fn index(&self, i: Self::Idx) -> usize;

    /// Return an iterator over all items.
    fn iter(&self) -> Self::Iter;
}

/// An indexer for a graph.
pub trait GraphIndexer<G>: Indexer {
    /// Create a new indexer for the given graph.
    fn new(g: G) -> Self;
}

/// A vector that can be indexed by nodes or edges.
///
/// In fact, the vector can be indexed by some arbitrary item that can be mapped
/// to a numeric index using an `Indexer`.
pub struct ItemVec<I, T> {
    indexer: I,
    data: Vec<T>,
}

impl<I, T> ItemVec<I, T>
where
    I: Indexer,
{
    pub fn from_indexer(indexer: I, value: T) -> Self
    where
        T: Clone,
    {
        ItemVec {
            data: vec![value; indexer.num_items()],
            indexer,
        }
    }

    pub fn from_indexer_with<F>(indexer: I, f: F) -> Self
    where
        F: Fn(I::Idx) -> T,
    {
        ItemVec {
            data: indexer.iter().map(f).collect(),
            indexer,
        }
    }

    pub fn from_indexer_and_vec(indexer: I, values: Vec<T>) -> Self {
        assert_eq!(values.len(), indexer.num_items());
        ItemVec { data: values, indexer }
    }

    pub fn new2<G>(g: G, value: T) -> Self
    where
        I: GraphIndexer<G>,
        T: Clone,
    {
        Self::from_indexer(I::new(g), value)
    }

    pub fn new<G>(g: G, value: T) -> Self
    where
        I: GraphIndexer<G>,
        T: Clone,
    {
        Self::from_indexer(I::new(g), value)
    }

    pub fn new_with<G, F>(g: G, f: F) -> Self
    where
        I: GraphIndexer<G>,
        F: Fn(I::Idx) -> T,
    {
        Self::from_indexer_with(I::new(g), f)
    }

    pub fn new_from_vec<G>(g: G, values: Vec<T>) -> Self
    where
        I: GraphIndexer<G>,
    {
        Self::from_indexer_and_vec(I::new(g), values)
    }

    pub fn iter(&self) -> std::iter::Zip<I::Iter, std::slice::Iter<T>> {
        self.indexer.iter().zip(self.data.iter())
    }

    pub fn iter_mut(&mut self) -> std::iter::Zip<I::Iter, std::slice::IterMut<T>> {
        self.indexer.iter().zip(self.data.iter_mut())
    }

    pub fn map<F>(&self, f: F) -> Self
    where
        I: Clone,
        F: Fn((I::Idx, &T)) -> T,
    {
        Self::from_indexer_and_vec(self.indexer.clone(), self.iter().map(f).collect())
    }

    pub fn reset(&mut self, value: T)
    where
        T: Clone,
    {
        for u in self.indexer.iter() {
            self.data[self.indexer.index(u)] = value.clone();
        }
    }
}

impl<I, T> Index<I::Idx> for ItemVec<I, T>
where
    I: Indexer,
{
    type Output = T;
    fn index(&self, i: I::Idx) -> &T {
        &self.data[self.indexer.index(i)]
    }
}

impl<I, T> IndexMut<I::Idx> for ItemVec<I, T>
where
    I: Indexer,
{
    fn index_mut(&mut self, i: I::Idx) -> &mut T {
        &mut self.data[self.indexer.index(i)]
    }
}

/// An indexer for nodes of an `IndexGraph`.
#[derive(Copy)]
pub struct NodeIndexer<'a, G>(pub G, PhantomData<&'a ()>);

impl<'a, G> Clone for NodeIndexer<'a, G>
where
    G: Clone,
{
    fn clone(&self) -> Self {
        NodeIndexer(self.0.clone(), PhantomData)
    }
}

impl<'a, G> Indexer for NodeIndexer<'a, G>
where
    G: IndexGraphRef<'a>,
{
    type Idx = G::Node;
    type Iter = G::NodeIter;

    fn num_items(&self) -> usize {
        self.0.num_nodes()
    }

    fn index(&self, u: G::Node) -> usize {
        self.0.node_id(u)
    }

    fn iter(&self) -> Self::Iter {
        GraphSizeRef::nodes(&self.0)
    }
}

impl<'a, G> GraphIndexer<G> for NodeIndexer<'a, G>
where
    G: IndexGraphRef<'a>,
{
    fn new(g: G) -> Self {
        NodeIndexer(g, PhantomData)
    }
}

/// An indexer for edges of an `IndexGraph`.
#[derive(Copy)]
pub struct EdgeIndexer<'a, G>(pub G, PhantomData<&'a ()>);

impl<'a, G> Clone for EdgeIndexer<'a, G>
where
    G: Clone,
{
    fn clone(&self) -> Self {
        EdgeIndexer(self.0.clone(), PhantomData)
    }
}

impl<'a, G> Indexer for EdgeIndexer<'a, G>
where
    G: IndexGraphRef<'a>,
{
    type Idx = G::Edge;
    type Iter = G::EdgeIter;

    fn num_items(&self) -> usize {
        self.0.num_edges()
    }

    fn index(&self, e: G::Edge) -> usize {
        self.0.edge_id(e)
    }

    fn iter(&self) -> Self::Iter {
        GraphSizeRef::edges(&self.0)
    }
}

impl<'a, G> GraphIndexer<G> for EdgeIndexer<'a, G>
where
    G: IndexGraphRef<'a>,
{
    fn new(g: G) -> Self {
        EdgeIndexer(g, PhantomData)
    }
}

/// A vector containing a value for each node.
///
/// # Example
///
/// Create a vector with all elements set to some value.
///
/// ```
/// # use rs_graph::LinkedListGraph;
/// # use rs_graph::traits::*;
/// # use rs_graph::classes::peterson;
/// #
/// let g = peterson::<LinkedListGraph>();
/// let weights = rs_graph::vec::NodeVec::new(&g, 0);
/// assert!(g.nodes().all(|u| weights[u] == 0));
/// ```
///
/// Create a vector with an initialization function.
///
/// ```
/// # use rs_graph::LinkedListGraph;
/// # use rs_graph::traits::*;
/// # use rs_graph::classes::peterson;
/// #
/// let g = peterson::<LinkedListGraph>();
/// let weights = rs_graph::vec::NodeVec::new_with(&g, |u| g.node_id(u));
/// assert!(g.nodes().all(|u| weights[u] == g.node_id(u)));
/// ```
pub type NodeVec<'a, G, T> = ItemVec<NodeIndexer<'a, G>, T>;

/// A vector containing a value for each edge.
///
/// # Example
///
/// Create a vector with all elements set to some value.
///
/// ```
/// # use rs_graph::LinkedListGraph;
/// # use rs_graph::traits::*;
/// # use rs_graph::classes::peterson;
/// #
/// let g = peterson::<LinkedListGraph>();
/// let weights = rs_graph::EdgeVec::new(&g, 0);
/// assert!(g.edges().all(|e| weights[e] == 0));
/// ```
///
/// Create a vector with all values initialized by an initializing function.
///
/// ```
/// # use rs_graph::LinkedListGraph;
/// # use rs_graph::traits::*;
/// # use rs_graph::classes::peterson;
/// #
/// let g = peterson::<LinkedListGraph>();
/// let weights = rs_graph::vec::EdgeVec::new_with(&g, |e| g.edge_id(e));
/// assert!(g.edges().all(|e| weights[e] == g.edge_id(e)));
/// ```
pub type EdgeVec<'a, G, T> = ItemVec<EdgeIndexer<'a, G>, T>;

#[cfg(test)]
mod tests {
    use super::*;
    use crate::classes::peterson;
    use crate::traits::*;
    use crate::LinkedListGraph;

    #[test]
    fn test_iter() {
        let g = peterson::<LinkedListGraph>();
        let mut weights = EdgeVec::new_with(&g, |e| g.edge_id(e));

        for (i, (_, &x)) in weights.iter().enumerate() {
            assert_eq!(i, x);
        }

        for (_, x) in weights.iter_mut() {
            *x = *x + 1;
        }

        for (i, (_, &x)) in weights.iter().enumerate() {
            assert_eq!(i + 1, x);
        }
    }

    #[test]
    fn test_map() {
        let g = peterson::<LinkedListGraph>();
        let weights = EdgeVec::new_with(&g, |e| g.edge_id(e));

        let new_weights = weights.map(|(_, i)| i + 1);
        for e in g.edges() {
            assert_eq!(weights[e] + 1, new_weights[e]);
        }

        let new_weights = new_weights.map(|(_, i)| i * 2);
        for e in g.edges() {
            assert_eq!(2 * (weights[e] + 1), new_weights[e]);
        }
    }
}