//! This file implements a recursive structure that implements the
//! recursion scheme traits, representing trees.
//!
//! The tree is backed by a vector.

use crate::{
    algebra::Algebra,
    catana::{Ana, Cata},
    coalgebra::Coalgebra,
    functor::Functor,
    hylo::Hylo,
};

use std::{collections::VecDeque, mem::MaybeUninit, ops::Deref};

/// The evaluation strategy for the tree.
#[derive(Default, Debug, Copy, Clone)]
pub enum TEStrategy {
    #[default]
    /// This strategy uses an arena, and uses an `Option<T>` to store
    /// the data.
    ///
    /// # Comparison:
    ///
    /// Since it is an arena, it saves allocations, compared to the
    /// variant [`DepthFirst`][TEStrategy::DepthFirst].  But it needs
    /// indices to operate, so uses more memory.
    ///
    /// On the other hand, it uses an option, so is slower than the
    /// variant [`UnsafeArena`][TEStrategy::UnsafeArena], but avoids
    /// unsafe code altogether.  Applications can first use this
    /// variant to make sure the algorithm works, before converting to
    /// use the unsafe variant.
    SafeArena,
    /// This strategy uses an arena, and uses an `MaybeUninit` to
    /// store the data.
    ///
    /// # Comparison:
    ///
    /// Since it is an arena, it saves allocations, compared to the
    /// variant [`DepthFirst`][TEStrategy::DepthFirst].  But it needs
    /// indices to operate, so uses more memory.
    ///
    /// On the other hand, it uses a `MaybeUninit`, so is faster than
    /// the variant [`SafeArena`][TEStrategy::SafeArena], but uses
    /// unsafe code.  Applications can first use the safe variant to
    /// make sure the algorithm works, before converting to use this
    /// variant.
    UnsafeArena,
    /// This strategy uses a plain vector.
    ///
    /// # Comparison:
    ///
    /// Since it is a plain vector, it uses more allocations, compared
    /// to other variants.  But it does not use indices, so consumes
    /// less memory.
    ///
    /// # Warning
    ///
    /// Since it uses no indices, it relies on the depth-first order
    /// of the elements to correctly find elements.  This puts a
    /// requirement on the implementation of the [`Functor`] trait.
    DepthFirst,
}

/// A tree is just a wrapper around a vector.
///
/// # Warning
///
/// The tree is supposed to be stored in topological order.  This
/// order is used in a critical way in the implementations of
/// recursion schemes.  Violations of this assumption are fatal to
/// using those trait methods.
#[derive(Clone, Debug)]
pub struct Tree<T> {
    elements: Vec<T>,
    strategy: TEStrategy,
}

impl<T> Tree<T> {
    #[inline]
    /// Construct a new tree.
    pub fn new(elements: Vec<T>, strategy: TEStrategy) -> Self {
        Self { elements, strategy }
    }

    /// Just a function for testing.
    ///
    /// # Warning
    ///
    /// This is definitely going to be removed in the future.
    pub fn nth(&self, n: usize) -> Option<&T> {
        self.elements.get(n)
    }

    #[inline]
    /// Retrieve the strategy of the tree.
    pub fn strategy(&self) -> TEStrategy {
        self.strategy
    }

    #[inline]
    /// Set the strategy of the tree.
    pub fn set_strategy(&mut self, strategy: TEStrategy) {
        self.strategy = strategy;
    }
}

// Manual implementation can avoid unnecessary requirement on the
// parameter `T`.
impl<T> Default for Tree<T> {
    fn default() -> Self {
        let elements = Vec::new();
        let strategy = TEStrategy::default();

        Self { elements, strategy }
    }
}

#[derive(Debug, Copy, Clone)]
/// A thin wrapper around `usize`, to index vectors.
///
/// By means of the [*newtype
/// pattern*](https://doc.rust-lang.org/rust-by-example/generics/new_types.html)
/// in Rust, it is supposed to be treated as a simple `usize` in the
/// compiled codes.
pub struct TreeIndex(usize);

impl TreeIndex {
    /// Wrap an index in this type.
    pub fn new(index: usize) -> Self {
        Self(index)
    }
}

impl Deref for TreeIndex {
    type Target = usize;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T, F, G, A> Cata<T, F, A> for Tree<G>
where
    F: Functor<T>,
    G: Functor<TreeIndex, Target<T> = F>,
    A: Algebra<T, F>,
{
    fn cata(self, mut alg: A) -> T {
        // First deal with the safe case
        match self.strategy {
            TEStrategy::SafeArena => {
                let mut results: Vec<Option<T>> = std::iter::repeat_with(Default::default)
                    .take(self.elements.len())
                    .collect();

                for (index, node) in self.elements.into_iter().enumerate().rev() {
                    let algebra_result = {
                        let node = node.fmap::<T>(|index| {
                            std::mem::replace(&mut results[*index], None).unwrap()
                        });

                        alg(node)
                    };

                    // Artificially use this value to satisfy the compiler.
                    let _ = std::mem::replace(&mut results[index], Some(algebra_result));
                }

                std::mem::replace(&mut results[0], None).unwrap()
            }
            TEStrategy::UnsafeArena => {
                let mut results: Vec<MaybeUninit<T>> = std::iter::repeat_with(MaybeUninit::uninit)
                    .take(self.elements.len())
                    .collect();

                for (index, node) in self.elements.into_iter().enumerate().rev() {
                    let algebra_result = {
                        let node = node.fmap::<T>(|index| unsafe {
                            std::mem::replace(&mut results[*index], MaybeUninit::uninit())
                                .assume_init()
                        });

                        alg(node)
                    };

                    results[index].write(algebra_result);
                }

                unsafe { std::mem::replace(&mut results[0], MaybeUninit::uninit()).assume_init() }
            }
            TEStrategy::DepthFirst => {
                let mut results_stack: Vec<T> = Vec::new();

                for node in self.elements.into_iter().rev() {
                    // Replace each node data with the value from the
                    // results stack.
                    let mapped_node = node.fmap(|_| results_stack.pop().unwrap());

                    results_stack.push(alg(mapped_node));
                }

                results_stack.pop().unwrap()
            }
        }
    }
}

impl<T, F, G, C> Ana<T, F, C> for Tree<G>
where
    F: Functor<T, Target<TreeIndex> = G>,
    G: Functor<TreeIndex>,
    C: Coalgebra<T, F>,
{
    /// An anamorphism takes a single, flat, collapsed value and a
    /// co-algebra for a recursive structure, and returns that
    /// recursive structure.
    ///
    /// # Descriptions
    ///
    /// This always generates a tree which uses the default strategy.
    /// If one wants to use a different strategy, set the strategy
    /// after generating the tree.
    ///
    /// # See also
    ///
    /// To use the depth first strategy to generate the tree, use the
    /// wrapper struct [`DFTree`].
    fn ana(value: T, mut coalg: C) -> Self {
        let mut queue = VecDeque::new();

        queue.push_back(value);

        let mut elements = vec![];

        let strategy = TEStrategy::default();

        while let Some(value) = queue.pop_back() {
            let expanded_layer = coalg(value);

            let mapped_layer = expanded_layer.fmap::<TreeIndex>(|value| {
                queue.push_back(value);

                TreeIndex(elements.len() + queue.len())
            });

            elements.push(mapped_layer);
        }

        Self { elements, strategy }
    }
}

/// To generate a tree with the strategy
/// [`DepthFirst`][TEStrategy::DepthFirst], we use a wrapper struct
/// which implements [`Ana`] in the desired manner.
#[derive(Debug, Clone)]
pub struct DFTree<T>(Tree<T>);

impl<T> DFTree<T> {
    #[inline]
    /// Convert to the underlying tree.
    pub fn to_tree(self) -> Tree<T> {
        self.0
    }

    #[inline]
    /// Wrap a tree.
    pub fn new(tree: Tree<T>) -> Self {
        Self(tree)
    }
}

impl<T, F, G, C> Ana<T, F, C> for DFTree<G>
where
    F: Functor<T, Target<TreeIndex> = G>,
    G: Functor<TreeIndex>,
    C: Coalgebra<T, F>,
{
    /// An anamorphism takes a single, flat, collapsed value and a
    /// co-algebra for a recursive structure, and returns that
    /// recursive structure.
    ///
    /// # Descriptions
    ///
    /// This always generates a tree which uses the depth first
    /// strategy.  If one wants to use a different strategy, set the
    /// strategy after generating the tree.
    ///
    /// # See also
    ///
    /// To use the default strategy to generate the tree, use the
    /// original struct [`Tree`].
    fn ana(value: T, mut coalg: C) -> Self {
        let mut stack = Vec::new();

        stack.push(value);

        let mut elements = vec![];

        let strategy = TEStrategy::DepthFirst;

        while let Some(value) = stack.pop() {
            let expanded_layer = coalg(value);

            let mut local_stack = Vec::new();

            let mapped_layer = expanded_layer.fmap::<TreeIndex>(|value| {
                local_stack.push(value);

                // The index is of no meaning here, since we rely on
                // the depth-first order.
                TreeIndex(0)
            });

            stack.extend(local_stack.into_iter().rev());

            elements.push(mapped_layer);
        }

        Self::new(Tree::new(elements, strategy))
    }
}

impl<T, U, F, G, H, A, C> Hylo<T, TreeIndex, U, F, G, H, A, C> for Tree<G>
where
    F: Functor<T>,
    G: Functor<TreeIndex, Target<T> = F>,
    H: Functor<U, Target<TreeIndex> = G>,
    A: Algebra<T, F>,
    C: Coalgebra<U, H>,
{
    fn hylo(value: U, mut alg: A, mut coalg: C) -> T {
        // The hylomorphism ignores the tree.  Maybe I will add
        // different implementations later on.

        let mut result_stack: Vec<T> = Vec::new();
        let mut value_node_stack: Vec<Result<U, G>> = vec![Ok(value)];

        while let Some(value_or_node) = value_node_stack.pop() {
            match value_or_node {
                Ok(value) => {
                    let node = coalg(value);

                    let mut local_values: Vec<U> = Vec::new();

                    let mapped_node = node.fmap(|node_value| {
                        local_values.push(node_value);
                        TreeIndex::new(0)
                    });

                    value_node_stack.push(Err(mapped_node));
                    value_node_stack.extend(local_values.into_iter().rev().map(Ok));
                }
                Err(node) => {
                    let mapped_node = node.fmap(|_| result_stack.pop().unwrap());

                    result_stack.push(alg(mapped_node));
                }
            }
        }

        result_stack.pop().unwrap()
    }
}

// REVIEW: Para, Apo, Histo, and Futu await us.