//! Maps items to lifespans
//! It support views on a set of items at a certain point in time.

use std::collections::BTreeMap;
use std::fmt::Debug;
use std::slice;

use crate::datastructs::lifespan::Lifespan;

#[cfg(test)]
mod tests;


/// A bucket in the tree, containing at least one Span
#[derive(Debug)]
struct Bucket<U>
    where
        U: Copy + PartialOrd + Debug
{
    first_span: Lifespan<U>,
    rest: Vec<Lifespan<U>>,
}

impl<U> Bucket<U>
    where
        U: Copy + PartialOrd + Debug
{
    fn new(begin: U, end: Option<U>) -> Self {
        Bucket { first_span: Lifespan::new(begin, end), rest: vec!() }
    }
}

/// A Multiset of items T, each with it's own lifetime
#[derive(Debug)]
pub struct LifespanMap<T, U>
    where
        T: Clone + Ord + Debug,
        U: Copy + PartialOrd + Debug
{
    data: BTreeMap<T, Bucket<U>>
}

impl<T, U> LifespanMap<T, U>
    where
        T: Clone + Ord + Debug,
        U: Copy + PartialOrd + Debug
{
    pub fn new() -> Self {
        LifespanMap { data: BTreeMap::new() }
    }

    /// Insert an item without an end time
    pub fn insert(&mut self, item: T, begin: U) {
        self.data
            .entry(item)
            .and_modify(|bucket| bucket.rest.push(Lifespan::new(begin, None)))
            .or_insert(Bucket::new(begin, None))
        ;
    }

    /// Insert an item with a given end time
    pub fn insert_with_end(&mut self, item: T, begin: U, end: U) {
        debug_assert!(begin <= end);

        self.data
            .entry(item)
            .and_modify(|bucket| bucket.rest.push(Lifespan::new(begin, Some(end))))
            .or_insert(Bucket::new(begin, Some(end)))
        ;
    }

    /// Set an end time on an item with open end
    pub fn set_end(&mut self, item: &T, end: U) {
        self.data
            .get_mut(item)
            .map(
                |bucket|
                    match bucket.first_span.end
                    {
                        None => {
                            bucket.first_span.end = Some(end)
                        }
                        Some(end) => {
                            bucket.rest.iter_mut()
                                .filter(|span| !(span.begin <= end && span.end.is_none()))
                                .next()
                                .expect("no item with open end time found in collection")
                                .end = Some(end)
                        }
                    }
            );
    }

    pub fn iter(&self) -> Iter<T, U> {
        Iter::new(self.data.iter())
    }

    /// Creates an iterator that returns all items that exist just before the given timestamp,
    /// in unspecified order
    pub fn iter_before(&self, before: U) -> IterBefore<T, U> {
        IterBefore::new(before, self.data.iter())
    }

    /// Creates an iterator that returns all items that exist just after the given timestamp,
    /// in unspecified order
    pub fn iter_after(&self, before: U) -> IterAfter<T, U> {
        IterAfter::new(before, self.data.iter())
    }

    /// Creates an iterator that returns all items that exist at the given timestamp,
    /// in unspecified order
    pub fn iter_at(&self, at: U) -> IterAt<T, U> {
        IterAt::new(at, self.data.iter())
    }

    /// Find and removes all spans with an end date
    pub fn cleanup(&mut self) {
        // todo: for now we do a linear search, but this could be improved by storing the locations
        // upon setting the end time, then we wouldn't have to search all

        let mut to_remove = vec!();

        for (value, bucket) in self.data.iter_mut() {
            let mut pos: usize = bucket.rest.len();
            while pos > 0 {
                pos = pos - 1;

                if bucket.rest[pos].end.is_some() {
                    bucket.rest.remove(pos);
                }
            }

            if bucket.first_span.end.is_some() {
                if bucket.rest.len() > 0 {
                    bucket.first_span = bucket.rest.pop().unwrap();
                } else {
                    to_remove.push(value.clone());
                }
            }
        }

        for value in to_remove {
            self.data.remove(&value);
        }
    }

    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }
}

/// Iterator
pub struct Iter<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    iter: std::collections::btree_map::Iter<'a, T, Bucket<U>>,
    current: Option<(&'a T, slice::Iter<'a, Lifespan<U>>)>,
}

impl<'a, T, U> Iter<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    fn new(iter: std::collections::btree_map::Iter<'a, T, Bucket<U>>) -> Self {
        Iter { iter, current: None }
    }
}

impl<'a, T, U> Iterator for Iter<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        // first try to exhaust current bucket
        match &mut self.current {
            Some(tuple) => {
                // every time we find a match (a lifespan), we return a reference to the T
                if let Some(_span) = tuple.1.next() {
                    return Some(tuple.0);
                }
            }
            None => {}
        };

        if let Some(tuple) = self.iter.next() {
            let bucket = tuple.1;


            self.current = Some((tuple.0, bucket.rest.iter()));
            return Some(tuple.0);
        };

        None
    }
}

/// Iterator just before a given point in time
pub struct IterBefore<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    before: U,
    iter: std::collections::btree_map::Iter<'a, T, Bucket<U>>,
    current: Option<(&'a T, slice::Iter<'a, Lifespan<U>>)>,
}

impl<'a, T, U> IterBefore<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    fn new(before: U, iter: std::collections::btree_map::Iter<'a, T, Bucket<U>>) -> Self {
        IterBefore { before, iter, current: None }
    }
}

impl<'a, T, U> Iterator for IterBefore<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        // first try to exhaust current bucket
        match &mut self.current {
            Some(tuple) => {
                while let Some(span) = tuple.1.next() {
                    if span.before(self.before) {
                        return Some(tuple.0);
                    }
                }
            }
            None => {}
        };

        while let Some(tuple) = self.iter.next() {
            let bucket = tuple.1;

            if tuple.1.first_span.before(self.before) {
                self.current = Some((tuple.0, bucket.rest.iter()));
                return Some(tuple.0);
            }

            let mut iter = bucket.rest.iter();
            while let Some(span) = iter.next() {
                if span.before(self.before) {
                    self.current = Some((tuple.0, iter));
                    return Some(tuple.0);
                }
            }
        };

        None
    }
}

/// Iterator at a given point in time
pub struct IterAt<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    at: U,
    iter: std::collections::btree_map::Iter<'a, T, Bucket<U>>,
    current: Option<(&'a T, slice::Iter<'a, Lifespan<U>>)>,
}

impl<'a, T, U> IterAt<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    fn new(at: U, iter: std::collections::btree_map::Iter<'a, T, Bucket<U>>) -> Self {
        IterAt { at, iter, current: None }
    }
}

impl<'a, T, U> Iterator for IterAt<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        // first try to exhaust current bucket
        match &mut self.current {
            Some(tuple) => {
                while let Some(span) = tuple.1.next() {
                    if span.at(self.at) {
                        return Some(tuple.0);
                    }
                }
            }
            None => {}
        };

        while let Some(tuple) = self.iter.next() {
            let bucket = tuple.1;

            if tuple.1.first_span.at(self.at) {
                self.current = Some((tuple.0, bucket.rest.iter()));
                return Some(tuple.0);
            }

            let mut iter = bucket.rest.iter();
            while let Some(span) = iter.next() {
                if span.at(self.at) {
                    self.current = Some((tuple.0, iter));
                    return Some(tuple.0);
                }
            }
        };

        None
    }
}

/// Iterator just after a given point in time
pub struct IterAfter<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    after: U,
    iter: std::collections::btree_map::Iter<'a, T, Bucket<U>>,
    current: Option<(&'a T, slice::Iter<'a, Lifespan<U>>)>,
}

impl<'a, T, U> IterAfter<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    fn new(after: U, iter: std::collections::btree_map::Iter<'a, T, Bucket<U>>) -> Self {
        IterAfter { after, iter, current: None }
    }
}

impl<'a, T, U> Iterator for IterAfter<'a, T, U>
    where U: Copy + PartialOrd + Debug
{
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        // first try to exhaust current bucket
        match &mut self.current {
            Some(tuple) => {
                while let Some(span) = tuple.1.next() {
                    if span.after(self.after) {
                        return Some(tuple.0);
                    }
                }
            }
            None => {}
        };

        while let Some(tuple) = self.iter.next() {
            let bucket = tuple.1;

            if tuple.1.first_span.after(self.after) {
                self.current = Some((tuple.0, bucket.rest.iter()));
                return Some(tuple.0);
            }

            let mut iter = bucket.rest.iter();
            while let Some(span) = iter.next() {
                if span.after(self.after) {
                    self.current = Some((tuple.0, iter));
                    return Some(tuple.0);
                }
            }
        };

        None
    }
}