//! This file provides a default implementation of the
//! [`Chain`][crate::Chain] trait.
//!
//! The reason for using a trait is that I might want to experiment
//! with different implementation ideas in the future, and this
//! modular design makes that easy.

use super::*;
use crate::atom::{Atom, DefaultAtom};
use crate::item::{
    default::DefaultForest, generate_fragment, genins::index_out_of_bounds_conversion, Forest,
    ForestLabel, ForestLabelError,
};
use grammar::{GrammarLabel, GrammarLabelType, START_NONTERMINAL, TNT};
#[allow(unused_imports)]
use graph::{
    labelled::DLGBuilder, Builder, DLGraph, Graph, LabelExtGraph, LabelGraph, ParentsGraph,
};
use std::fmt::Display;

use graph_macro::Graph;

use std::{
    borrow::Borrow,
    collections::{HashMap as Map, HashSet, TryReserveError},
};

/// The errors related to taking derivatives by chain rule.
#[non_exhaustive]
#[derive(Debug)]
pub enum Error {
    /// General error for indices out of bounds.
    IndexOutOfBounds(usize, usize),
    /// The forest encounters a duplicate node, for some reason.
    DuplicateNode(usize),
    /// The chain rule machine encounters a duplicate edge, for some
    /// reason.
    DuplicateEdge(usize, usize),
    /// A node has no labels while it is required to have one.
    NodeNoLabel(usize),
    /// Reserving memory fails.
    CannotReserve(TryReserveError),
    /// Cannot create label for cloning nodes.
    CannotClone(ForestLabelError),
    /// Cannot close the virtual node.
    ///
    /// The format is (nt, t, node, pos).
    CannotClose(usize, usize, usize, usize),
    /// Cannot find a suitable node to plant the new forest fragment.
    CannotPlant,
    /// Cannot split a packed node.
    SplitPack(usize),
    /// A cloned node should have exactly one parent.
    InvalidClone(usize),
    /// An invalid situation happens.
    Invalid,
}

impl Display for Error {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::IndexOutOfBounds(index, bound) => write!(f, "index {index} out of bound {bound}"),
            Self::DuplicateNode(n) => write!(f, "the forest has a node {n} with a duplicate label"),
            Self::DuplicateEdge(source, target) => write!(
                f,
                "the forest has a duplicate edge from {source} to {target}"
            ),
            Self::NodeNoLabel(n) => write!(f, "node {n} has no labels while it should have one"),
            Self::CannotReserve(e) => write!(f, "cannot reserve memory: {e}"),
            Self::CannotClone(fe) => write!(f, "cannot clone due to {fe}"),
            Self::CannotClose(nt, t, node, pos) => write!(
                f,
                "cannot close the virtual node {node} \
                 with the expansion from nt {nt} by t \
                 {t} at pos {pos}"
            ),
            Self::CannotPlant => write!(f, "cannot plant a new forest fragment"),
            Self::SplitPack(n) => write!(f, "cannot split the packed node {n}"),
            Self::InvalidClone(n) => {
                write!(f, "the cloned node {n} should have exactly one parent")
            }
            Self::Invalid => write!(f, "invalid"),
        }
    }
}

impl std::error::Error for Error {}

impl From<GError> for Error {
    fn from(value: GError) -> Self {
        match value {
            GError::IndexOutOfBounds(index, bound) => Self::IndexOutOfBounds(index, bound),
            GError::DuplicatedNode(n) => Self::DuplicateNode(n),
            GError::DuplicatedEdge(source, target) => Self::DuplicateEdge(source, target),
            _ => Self::Invalid,
        }
    }
}

impl From<ForestError> for Error {
    fn from(e: ForestError) -> Self {
        match e {
            ForestError::IndexOutOfBounds(index, bound) => Error::IndexOutOfBounds(index, bound),
            ForestError::DuplicatedNode(n) => Error::DuplicateNode(n),
            ForestError::InvalidGraphError(ge) => ge.into(),
            ForestError::NodeNoLabel(n) => Error::NodeNoLabel(n),
            ForestError::LabelConversion(ce) => Error::CannotClone(ce),
            ForestError::SplitPack(n) => Error::SplitPack(n),
            ForestError::InvalidClone(n) => Error::SplitPack(n),
        }
    }
}

impl From<TryReserveError> for Error {
    fn from(value: TryReserveError) -> Self {
        Self::CannotReserve(value)
    }
}

/// The type of an index into an element in [`DerIter`].
#[derive(Debug, Copy, Clone)]
enum DerIterIndex {
    Single(usize),
    Map(usize),
}

impl Default for DerIterIndex {
    fn default() -> Self {
        Self::Map(0)
    }
}

/// A complex type used for storing values of edges with two layers.
type SecondTypeValue = (PaVi, bool, Vec<(Roi, usize)>);

/// An iterator of TwoLayers.
#[derive(Debug, Default)]
pub struct DerIter {
    /// Stores edges of only one layer.
    singles: Vec<(Roi, usize)>,
    /// Stores edges with two layers.  They are grouped by their
    /// labels of the second layer.
    ///
    /// The values are tuples (forest_source, accepting, edges), where
    /// the edges are the grouped edges of the first layer and the
    /// destination.
    seconds: Map<usize, SecondTypeValue>,
    /// We want to iterate the elements of the map, for which purpose
    /// we need an array.  Since hashmaps provide no arrays, we keep
    /// an array of keys for iteration purposes.
    second_array: Vec<usize>,
    /// The index of the current element, either in `second_array` or
    /// in `singles` .
    index: DerIterIndex,
}

impl DerIter {
    fn add_second_layer(
        &mut self,
        label: usize,
        forest_source: PaVi,
        accepting: bool,
        edges: Vec<(Roi, usize)>,
    ) {
        if let Some((_, _, vec)) = self.seconds.get_mut(&label) {
            vec.extend(edges);
        } else {
            self.seconds
                .insert(label, (forest_source, accepting, edges));

            self.second_array.push(label);
        }
    }
}

impl Iterator for DerIter {
    type Item = TwoLayers;

    fn next(&mut self) -> Option<Self::Item> {
        // We iterate through two layered edges first.
        match self.index {
            DerIterIndex::Map(index) => {
                if let Some(key) = self.second_array.get(index) {
                    if let Some((forest_source, accepting, edges)) = self.seconds.remove(key) {
                        self.index = DerIterIndex::Map(index + 1);

                        Some(TwoLayers::Two(*key, forest_source, accepting, edges))
                    } else {
                        // this should not happen
                        dbg!("a key does not exist in the hashmap: something is wrong when taking derivatives");
                        None
                    }
                } else {
                    // If the zero-th element is present, we will
                    // advance the index to one; if it is not present,
                    // we will stop iteration.  In each case we can
                    // safely set the index to one.
                    self.index = DerIterIndex::Single(1);

                    self.singles
                        .first()
                        .map(|(edge, to)| TwoLayers::One(*edge, *to))
                }
            }
            DerIterIndex::Single(index) => self.singles.get(index).map(|(edge, to)| {
                self.index = DerIterIndex::Single(index + 1);
                TwoLayers::One(*edge, *to)
            }),
        }
    }
}

/// A default implementation for the [`Chain`] trait.
#[derive(Debug, Clone, Default, Graph)]
pub struct DefaultChain {
    #[graph]
    graph: DLGraph<Edge>,
    atom: DefaultAtom,
    current: usize,
    history: Vec<usize>,
    forest: DefaultForest<ForestLabel<GrammarLabel>>,
    accepting_vec: Vec<bool>,
    accepting_sources: Vec<PaVi>,
}

impl DefaultChain {
    /// Return the current node.
    #[inline]
    pub fn current(&self) -> usize {
        self.current
    }

    /// Return the complete slice of histories.
    #[inline]
    pub fn history(&self) -> &[usize] {
        self.history.as_slice()
    }

    /// Return a reference to the associated forest.
    #[inline]
    pub fn forest(&self) -> &DefaultForest<ForestLabel<GrammarLabel>> {
        &self.forest
    }

    /// Print the rule positions of the labels.
    pub fn print_rule_positions(&self) -> Result<(), Box<dyn std::error::Error>> {
        let mut labels = HashSet::<usize>::default();

        for node in 0..self.graph.nodes_len() {
            labels.extend(self.graph.labels_of(node)?.map(|(label, _)| label.label));
        }

        for label in labels.into_iter() {
            println!("{}", self.atom.rule_pos_string(label)?);
        }

        Ok(())
    }
}

impl LabelGraph<Edge> for DefaultChain {
    type Iter<'a> = <DLGraph<Edge> as LabelGraph<Edge>>::Iter<'a>
    where
        Self: 'a;

    type LabelIter<'a> = <DLGraph<Edge> as LabelGraph<Edge>>::LabelIter<'a>
    where
        Self: 'a,
        Edge: 'a;

    type EdgeLabelIter<'a> = <DLGraph<Edge> as LabelGraph<Edge>>::EdgeLabelIter<'a>
    where
        Self: 'a,
        Edge: 'a;

    #[inline]
    fn edge_label(&self, source: usize, target: usize) -> Result<Self::EdgeLabelIter<'_>, GError> {
        self.graph.edge_label(source, target)
    }

    #[inline]
    fn find_children_with_label(
        &self,
        node_id: usize,
        label: &Edge,
    ) -> Result<<Self as LabelGraph<Edge>>::Iter<'_>, GError> {
        self.graph.find_children_with_label(node_id, label)
    }

    #[inline]
    fn labels_of(&self, node_id: usize) -> Result<Self::LabelIter<'_>, GError> {
        self.graph.labels_of(node_id)
    }

    #[inline]
    fn has_edge_label(&self, node_id: usize, label: &Edge, target: usize) -> Result<bool, GError> {
        self.graph.has_edge_label(node_id, label, target)
    }
}

impl LabelExtGraph<Edge> for DefaultChain {
    #[inline]
    fn extend(&mut self, edges: impl IntoIterator<Item = (Edge, usize)>) -> Result<usize, GError> {
        let new = self.graph.extend(edges)?;
        let accepting_len = self.accepting_vec.len();

        // Update the acceptace of the new node, if any.
        if self.accepting_vec.get(new).is_none() {
            // assert it can only grow by one node at a time.
            #[cfg(debug_assertions)]
            assert_eq!(new, accepting_len);

            let mut updated = false;

            for (label, child_iter) in self.graph.labels_of(new)? {
                let old_accepting = {
                    let mut result = false;
                    for child in child_iter {
                        if *self
                            .accepting_vec
                            .get(child)
                            .ok_or(GError::IndexOutOfBounds(child, accepting_len))?
                        {
                            result = true;
                            break;
                        }
                    }

                    result
                };

                if !old_accepting {
                    self.accepting_vec.push(false);
                    updated = true;

                    break;
                }

                if label.is_accepting() {
                    self.accepting_vec.push(true);
                    updated = true;

                    break;
                }
            }

            if !updated {
                self.accepting_vec.push(false);
            }
        }

        Ok(new)
    }
}

impl Chain for DefaultChain {
    type Error = Error;

    type Atom = DefaultAtom;

    fn unit(atom: Self::Atom) -> Result<Self, Self::Error> {
        let mut builder: DLGBuilder<Edge> = Default::default();

        let root = builder.add_vertex();
        let first = builder.add_vertex();

        let empty_state = atom.empty();

        let initial_nullable = atom
            .is_nullable(START_NONTERMINAL)
            .map_err(|_| Error::IndexOutOfBounds(0, atom.non_num()))?;

        builder.add_edge(
            first,
            root,
            Edge::new(empty_state, Default::default(), initial_nullable),
        )?;

        let graph = builder.build();

        let forest =
            DefaultForest::new_leaf(GrammarLabel::new(GrammarLabelType::TNT(TNT::Non(0)), 0));

        #[cfg(debug_assertions)]
        assert_eq!(forest.root(), Some(0));

        let current = 1;

        let history = Vec::new();

        let accepting_vec = vec![true, initial_nullable];

        let accepting_sources = Vec::new();

        Ok(Self {
            graph,
            atom,
            current,
            history,
            forest,
            accepting_vec,
            accepting_sources,
        })
    }

    fn atom(&self) -> &Self::Atom {
        self.atom.borrow()
    }

    fn epsilon(&self) -> Result<bool, Self::Error> {
        self.accepting_vec
            .get(self.current)
            .copied()
            .ok_or(Error::IndexOutOfBounds(
                self.current,
                self.accepting_vec.len(),
            ))
    }

    fn update_history(&mut self, new: usize) {
        debug_assert!(new < self.graph.nodes_len());

        self.history.push(self.current);

        self.current = new;
    }

    fn update_epsilon(
        &mut self,
        node_id: usize,
        edges: impl Iterator<Item = (Roi, usize)>,
    ) -> Result<(), Self::Error> {
        for (roi, _) in edges {
            if let Some(target) = roi.imaginary_part() {
                if matches!(self.accepting_vec.get(target).copied(), Some(true)) {
                    let accepting_vec_len = self.accepting_vec.len();

                    *self
                        .accepting_vec
                        .get_mut(node_id)
                        .ok_or(Error::IndexOutOfBounds(node_id, accepting_vec_len))? = true;
                }
            }
        }

        Ok(())
    }

    type DerIter = DerIter;

    fn derive(
        &mut self,
        t: usize,
        pos: usize,
        no_item: bool,
    ) -> Result<Self::DerIter, Self::Error> {
        use TNT::*;

        /// The function `generate_edges` takes too many arguments, so
        /// a type is used to group them together.
        ///
        /// The format is as follows:
        ///
        /// `(graph, atom, forest, accepting_vec, forest_source, new_source)`
        type GeneratorInput<'a> = (
            &'a DLGraph<Edge>,
            &'a DefaultAtom,
            &'a DefaultForest<ForestLabel<GrammarLabel>>,
            &'a [bool],
            PaVi,
            PaVi,
        );

        /// A helper function to generate edges to join.
        ///
        /// It first checks if the base edge is accepting.  If yes,
        /// then pull in the children of the target.
        ///
        /// Then check if the label of the base edge has children.  If
        /// no, then do not add this base edge itself.
        ///
        /// The generated edges will be pushed to `output` directly,
        /// to save some allocations.
        fn generate_edges(
            input: GeneratorInput<'_>,
            mut child_iter: impl Iterator<Item = usize> + ExactSizeIterator + Clone,
            atom_child_iter: impl Iterator<Item = usize> + Clone,
            mut output: impl AsMut<Vec<(Roi, usize)>>,
        ) -> Result<bool, Error> {
            let mut accepting = false;

            let (graph, atom, _forest, accepting_vec, pavi, true_pavi) = input;

            // First check the values from iterators are all valid.
            let graph_len = graph.nodes_len();
            let atom_len = atom.nodes_len();

            for child in child_iter.clone() {
                if !graph.has_node(child) {
                    return Err(Error::IndexOutOfBounds(child, graph_len));
                }
            }

            for atom_child in atom_child_iter.clone() {
                if !atom.has_node(atom_child) {
                    return Err(Error::IndexOutOfBounds(atom_child, atom_len));
                }
            }

            // From now on the nodes are all valid, so we can just
            // call `unwrap`.

            // Then calculate the number of edges to append, to avoid
            // repeated allocations
            let mut num = 0usize;

            let child_iter_total_degree = child_iter
                .clone()
                .map(|child| graph.degree(child).unwrap())
                .sum::<usize>();

            for atom_child in atom_child_iter.clone() {
                let atom_child_accepting = atom.is_accepting(atom_child).unwrap();

                num += child_iter.len();

                if atom_child_accepting {
                    accepting = true;
                    num += child_iter_total_degree;
                }
            }

            let num = num;

            let output = output.as_mut();

            output.try_reserve(num)?;

            // now push into output

            for atom_child in atom_child_iter {
                let atom_child_accepting = atom.is_accepting(atom_child).unwrap();
                let atom_child_empty_node = atom.is_empty_node(atom_child).unwrap();

                let roi = Edge::new(atom_child, pavi, atom_child_accepting).into();

                if atom_child_empty_node {
                    output.extend(child_iter.clone().map(|child| (child.into(), child)));
                } else {
                    output.extend(child_iter.clone().map(|child| (roi, child)));
                };

                if atom_child_accepting {
                    for child in child_iter.clone() {
                        for (child_label, child_child) in graph.labels_of(child).unwrap() {
                            let mut new_label = *child_label;

                            new_label.set_true_source(true_pavi);

                            output.extend(
                                std::iter::repeat(new_label.into())
                                    .take(child_child.len())
                                    .zip(child_child),
                            );
                        }
                    }
                }
            }

            accepting = accepting && child_iter.any(|child| *accepting_vec.get(child).unwrap());

            Ok(accepting)
        }

        let mut der_iter = DerIter::default();

        let mut accepting_pavi: HashSet<PaVi> = HashSet::new();

        for (label, child_iter) in self.graph.labels_of(self.current)? {
            for (atom_label, atom_child_iter) in self.atom.labels_of(label.label())? {
                if atom_label.is_left_p() {
                    // We do not consider left-linearly expanded
                    // children in the first layer.
                    continue;
                }

                let atom_moved = atom_label.get_moved();

                match *atom_label.get_value() {
                    Some(Ter(ter)) if ter == t => {
                        let new_pavi = if !no_item {
                            // prepare forest fragment

                            let fragment =
                                generate_fragment([atom_moved.into(), Ter(ter).into()], pos)?;

                            self.forest.insert_item(
                                *label,
                                t,
                                fragment,
                                atom_child_iter.clone(),
                                &self.atom,
                            )?
                        } else {
                            PaVi::default()
                        };

                        let generator_input = (
                            &self.graph,
                            &self.atom,
                            &self.forest,
                            self.accepting_vec.as_slice(),
                            new_pavi,
                            new_pavi,
                        );

                        let accepting = generate_edges(
                            generator_input,
                            child_iter.clone(),
                            atom_child_iter.clone(),
                            &mut der_iter.singles,
                        )?;

                        if accepting {
                            accepting_pavi.insert(new_pavi);
                        }
                    }
                    Some(Non(non)) => {
                        if !self
                            .atom
                            .is_first_of(non, t)
                            .map_err(index_out_of_bounds_conversion)?
                        {
                            continue;
                        }

                        let virtual_node = self
                            .atom
                            .atom(non, t)
                            .map_err(index_out_of_bounds_conversion)?;

                        if let Some(virtual_node) = virtual_node {
                            let accepting = self
                                .atom
                                .is_accepting(virtual_node)
                                .map_err(index_out_of_bounds_conversion)?;

                            let first_segment_pavi: PaVi;
                            let virtual_pavi: PaVi;

                            if !no_item {
                                let first_fragment =
                                    generate_fragment([atom_moved.into(), Non(non).into()], pos)?;

                                first_segment_pavi = self.forest.insert_item(
                                    *label,
                                    t,
                                    first_fragment,
                                    atom_child_iter.clone(),
                                    &self.atom,
                                )?;

                                let virtual_fragment =
                                    DefaultForest::new_leaf(GrammarLabel::new(Ter(t), pos));

                                // NOTE: We only need the PaVi from the
                                // first segment, so we pass an empty
                                // iterator, in which case the passed
                                // label is only used for the PaVi.
                                virtual_pavi = self.forest.insert_item(
                                    Edge::new(0, first_segment_pavi, accepting),
                                    t,
                                    virtual_fragment,
                                    std::iter::empty(),
                                    &self.atom,
                                )?;
                            } else {
                                first_segment_pavi = PaVi::default();
                                virtual_pavi = PaVi::default();
                            }

                            let mut new_edges = Vec::new();

                            let generator_input = (
                                &self.graph,
                                &self.atom,
                                &self.forest,
                                self.accepting_vec.as_slice(),
                                first_segment_pavi,
                                virtual_pavi,
                            );

                            let virtual_accepting = generate_edges(
                                generator_input,
                                child_iter.clone(),
                                atom_child_iter.clone(),
                                &mut new_edges,
                            )?;

                            if virtual_accepting {
                                accepting_pavi.insert(first_segment_pavi);
                            }

                            if accepting {
                                accepting_pavi.insert(virtual_pavi);

                                for (roi, target) in new_edges.iter() {
                                    match roi {
                                        Roi::Re(edge) => {
                                            let mut edge = *edge;

                                            edge.set_true_source(virtual_pavi);

                                            der_iter.singles.push((Roi::Re(edge), *target));
                                        }
                                        Roi::Im(_) => der_iter.singles.push((*roi, *target)),
                                    }
                                }
                            }

                            if !self.atom.is_empty_node(virtual_node).unwrap() {
                                der_iter.add_second_layer(
                                    virtual_node,
                                    virtual_pavi,
                                    accepting,
                                    new_edges,
                                );

                                // account for atom_children without
                                // children.

                                for atom_child in atom_child_iter {
                                    // this has been checked in
                                    // `generate_edges`
                                    if self.atom.is_empty_node(atom_child).unwrap() {
                                        der_iter.singles.extend(child_iter.clone().map(|child| {
                                            (
                                                Edge::new(virtual_node, virtual_pavi, accepting)
                                                    .into(),
                                                child,
                                            )
                                        }));
                                    }
                                }
                            } else {
                                for atom_child in atom_child_iter {
                                    // this has been checked in
                                    // `generate_edges`
                                    if self.atom.is_empty_node(atom_child).unwrap() {
                                        let mut extra_len = 0usize;

                                        for child in child_iter.clone() {
                                            extra_len += self
                                                .graph
                                                .labels_of(child)
                                                .unwrap()
                                                .map(|(_, child_child_iter)| child_child_iter.len())
                                                .sum::<usize>();
                                        }

                                        der_iter.singles.try_reserve(extra_len)?;

                                        for child in child_iter.clone() {
                                            for (child_label, child_child_iter) in
                                                self.graph.labels_of(child).unwrap()
                                            {
                                                let mut child_label = *child_label;

                                                child_label.set_true_source(virtual_pavi);

                                                der_iter.singles.extend(child_child_iter.map(
                                                    |child_child| (child_label.into(), child_child),
                                                ));
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                    _ => {
                        continue;
                    }
                }
            }
        }

        self.accepting_sources.clear();
        self.accepting_sources.extend(accepting_pavi);

        Ok(der_iter)
    }

    type Item = DefaultForest<ForestLabel<GrammarLabel>>;

    fn end_of_input(&mut self, pos: usize, ter: usize) -> Result<Self::Item, Error> {
        let mut result = Default::default();

        let root = if let Some(root) = self.forest.root() {
            root
        } else {
            return Ok(result);
        };

        // The 1-th node is a dummy node that should be removed.
        assert_ne!(root, 1);

        self.forest.remove_node(1)?;

        let root_degree = self.forest.degree(root)?;

        let root_label = self
            .forest
            .vertex_label(root)?
            .ok_or(Error::NodeNoLabel(root))?;

        dbg!(root_degree, root_label);

        // REVIEW: Why nodes?
        let mut nodes = Vec::with_capacity(root_degree);

        if root_label.is_packed() {
            let mut all_completed_clones = Vec::with_capacity(root_degree);

            for clone in self.forest.children_of(root)? {
                let mut all_completed = true;

                // The last child does not count.
                let clone_child_degree_minus_one = std::cmp::max(self.forest.degree(clone)?, 1) - 1;

                // The loop to determine whether or not it is all
                // completed, except for the last child.
                'completion_det_loop: for clone_child in self
                    .forest
                    .children_of(clone)?
                    .take(clone_child_degree_minus_one)
                {
                    let clone_child_label = self
                        .forest
                        .vertex_label(clone_child)?
                        .ok_or(Error::NodeNoLabel(clone_child))?;

                    if clone_child_label.label().end().is_none() {
                        all_completed = false;

                        break 'completion_det_loop;
                    }
                }

                if all_completed {
                    all_completed_clones.push(clone);
                }
            }

            if all_completed_clones.is_empty() {
                // Nothing to do
                return Ok(result);
            }

            for clone in all_completed_clones {
                nodes.push(
                    self.forest
                        .reduction(clone, pos, ter, self.atom.borrow(), true)?,
                );
            }
        } else if root_label.clone_index().is_some() {
            panic!(
                "If the root of the forest is cloned, \
                 the root should be changed to the packed node."
            );
        } else {
            // The last child does not count.
            let root_degree_minus_one = std::cmp::max(root_degree, 1) - 1;

            dbg!(root_degree_minus_one);

            // The loop to determine whether or not it is all
            // completed, except for the last child.
            for root_child in self.forest.children_of(root)?.take(root_degree_minus_one) {
                let root_child_label = self
                    .forest
                    .vertex_label(root_child)?
                    .ok_or(Error::NodeNoLabel(root_child))?;

                if root_child_label.label().end().is_none() {
                    return Ok(result);
                }
            }

            nodes.push(
                self.forest
                    .reduction(root, pos, ter, self.atom.borrow(), true)?,
            );
        }

        dbg!(&nodes);

        // self.forest
        //     .print_viz("forest before extraction.gv")
        //     .unwrap();

        result = self.forest.extract(pos)?;

        // result.print_viz("extracted forest.gv").unwrap();

        Ok(result)
    }
}

#[cfg(test)]
mod test_der_iter {
    use super::*;

    #[test]
    fn test() -> Result<(), Box<dyn std::error::Error>> {
        let mut der_iter = DerIter::default();

        let parent = Default::default();

        der_iter
            .singles
            .push((Edge::new(0, parent, true).into(), 0));

        der_iter
            .singles
            .push((Edge::new(1, parent, true).into(), 0));

        der_iter
            .singles
            .push((Edge::new(2, parent, true).into(), 0));

        der_iter.add_second_layer(
            3,
            parent,
            true,
            vec![(Edge::new(4, parent, true).into(), 1)],
        );

        der_iter.add_second_layer(
            6,
            parent,
            true,
            vec![(Edge::new(5, parent, true).into(), 1)],
        );

        // add an entry with a repeated label
        der_iter.add_second_layer(
            3,
            parent,
            true,
            vec![(Edge::new(7, parent, true).into(), 2)],
        );

        assert_eq!(
            der_iter.next(),
            Some(TwoLayers::Two(
                3,
                parent,
                true,
                vec![
                    (Edge::new(4, parent, true).into(), 1),
                    (Edge::new(7, parent, true).into(), 2)
                ]
            ))
        );

        assert_eq!(
            der_iter.next(),
            Some(TwoLayers::Two(
                6,
                parent,
                true,
                vec![(Edge::new(5, parent, true).into(), 1)]
            ))
        );

        assert_eq!(
            der_iter.next(),
            Some(TwoLayers::One(Edge::new(0, parent, true).into(), 0))
        );

        assert_eq!(
            der_iter.next(),
            Some(TwoLayers::One(Edge::new(1, parent, true).into(), 0))
        );

        assert_eq!(
            der_iter.next(),
            Some(TwoLayers::One(Edge::new(2, parent, true).into(), 0))
        );

        assert_eq!(der_iter.next(), None);
        assert_eq!(der_iter.next(), None);

        Ok(())
    }
}

#[cfg(test)]
mod test_chain {
    use super::*;
    use grammar::test_grammar_helper::*;

    #[test]
    fn base_test() -> Result<(), Box<dyn std::error::Error>> {
        let mut no_item = false;

        for arg in std::env::args() {
            if arg == "no_item" {
                no_item = true;
                break;
            }
        }

        let grammar = new_notes_grammar()?;

        let atom = DefaultAtom::from_grammar(grammar)?;

        let mut chain = DefaultChain::unit(atom)?;
        chain.chain(3, 00, no_item)?;
        chain.chain(1, 01, no_item)?;
        chain.chain(2, 02, no_item)?;
        chain.chain(2, 03, no_item)?;
        chain.chain(2, 04, no_item)?;
        chain.chain(0, 05, no_item)?;
        chain.chain(5, 06, no_item)?;
        chain.chain(1, 07, no_item)?;
        chain.chain(6, 08, no_item)?;
        chain.chain(6, 09, no_item)?;
        chain.chain(6, 10, no_item)?;
        chain.chain(0, 11, no_item)?;
        chain.chain(0, 12, no_item)?;

        if !no_item {
            let _output = chain.end_of_input(13, 0)?;
            chain.forest.print_viz("forest.gv")?;

            chain.graph.print_viz("chain.gv")?;
            chain.atom.print_nfa("nfa.gv")?;
            item::default::print_labels(&chain.atom, &chain.forest)?;
        }

        for label in chain.labels_of(chain.current())?.map(|(label, _)| label) {
            dbg!(label);
        }

        assert!(matches!(chain.epsilon(), Ok(true)));

        #[cfg(feature = "test-print-viz")]
        {
            chain.graph.print_viz("chain.gv")?;
            chain.atom.print_nfa("nfa.gv")?;

            chain.forest.print_viz("forest.gv")?;

            item::default::print_labels(&chain.atom, &chain.forest)?;
        }

        Ok(())
    }

    #[test]
    fn test_assumption() -> Result<(), Box<dyn std::error::Error>> {
        use grammar::Grammar;

        dbg!();
        let grammar_str = std::fs::read_to_string(
            "/Users/durand/Desktop/Centre/A propos de programmes/Rust/rep/grammar/abnf grammars/test.abnf",
        )
        .unwrap();
        dbg!();
        let grammar: Grammar = grammar_str.parse()?;

        let atom = DefaultAtom::from_grammar(grammar)?;
        let mut chain = DefaultChain::unit(atom)?;

        let no_item = false;

        let input: &[usize] = &[3, 0, 2, 2, 2, 1, 1, 1, 0, 1];

        for (pos, t) in input.iter().copied().enumerate().take(2) {
            chain.chain(t, pos, no_item)?;
            chain.forest().print_viz(&format!("forest {pos}.gv"))?;
            dbg!(pos, t);
        }

        item::default::print_labels(&chain.atom, &chain.forest)?;

        Ok(())
    }

    #[test]
    fn test_ambiguity() -> Result<(), Box<dyn std::error::Error>> {
        if std::fs::metadata("debug.log").is_ok() {
            std::fs::remove_file("debug.log")?;
        }

        let mut no_item = false;

        for arg in std::env::args() {
            if arg == "no_item" {
                no_item = true;
                break;
            }
        }

        let grammar = new_paren_grammar()?;

        let atom = DefaultAtom::from_grammar(grammar)?;

        let mut chain = DefaultChain::unit(atom)?;
        chain.chain(0, 0, no_item)?;

        if !no_item {
            chain.forest.print_viz("forest0.gv")?;
        }

        chain.chain(2, 1, no_item)?;

        if !no_item {
            chain.forest.print_viz("forest1.gv")?;
        }

        chain.chain(2, 2, no_item)?;

        if !no_item {
            chain.forest.print_viz("forest2.gv")?;
        }

        chain.chain(2, 3, no_item)?;

        if !no_item {
            chain.forest.print_viz("forest3.gv")?;
        }

        chain.chain(1, 4, no_item)?;

        if !no_item {
            let _output = chain.end_of_input(5, 1)?;
            chain.forest.print_viz("forest.gv")?;
            // chain.forest.print_closed_viz("closed.gv")?;
        }

        chain.graph.print_viz("chain.gv")?;
        chain.atom.print_nfa("nfa.gv")?;
        item::default::print_labels(&chain.atom, &chain.forest)?;

        for label in chain.labels_of(chain.current())?.map(|(label, _)| label) {
            dbg!(label);
        }

        dbg!(chain.current(), chain.history());

        #[cfg(feature = "test-print-viz")]
        {
            chain.graph.print_viz("chain.gv")?;
            chain.atom.print_nfa("nfa.gv")?;

            chain.forest.print_viz("forest.gv")?;

            chain.forest.print_closed_viz("closed.gv")?;

            item::default::print_labels(&chain.atom, &chain.forest)?;
        }

        assert!(matches!(chain.epsilon(), Ok(true)));

        Ok(())
    }

    #[test]
    fn test_speed() -> Result<(), Box<dyn std::error::Error>> {
        let mut no_item = false;

        for arg in std::env::args() {
            if arg == "no_item" {
                no_item = true;
                break;
            }
        }

        let grammar = new_notes_grammar_no_regexp()?;

        println!("grammar: {grammar}");

        let atom = DefaultAtom::from_grammar(grammar)?;

        let mut chain = DefaultChain::unit(atom)?;

        let input_template = vec![3, 1, 2, 2, 2, 0, 5, 1, 6, 6, 6, 0, 0];

        let repeat_times = {
            let mut result = 1;

            for arg in std::env::args() {
                let parse_as_digit: Result<usize, _> = arg.parse();

                // just use the first number in the arguments
                if let Ok(parse_result) = parse_as_digit {
                    result = parse_result;

                    break;
                }
            }

            result
        };

        println!("repeating {repeat_times} times");

        let input = input_template.repeat(repeat_times);

        let start = std::time::Instant::now();

        for (index, t) in input.iter().copied().enumerate() {
            chain.chain(t, index, no_item)?;
        }

        let elapsed = start.elapsed();

        assert!(matches!(chain.epsilon(), Ok(true)));

        dbg!(elapsed);
        dbg!(chain.current());

        assert_eq!(input.len(), chain.history().len());

        let history_file_name = "output/history";
        if std::fs::metadata(history_file_name).is_ok() {
            std::fs::remove_file(history_file_name)?;
        }

        let mut history_file = std::fs::OpenOptions::new()
            .create(true)
            .write(true)
            .open(history_file_name)?;

        use std::fmt::Write;
        use std::io::Write as IOWrite;

        let mut log_string = String::new();

        writeln!(&mut log_string, "index: terminal, history")?;

        for (index, t) in input.iter().copied().enumerate().take(input.len() - 1) {
            writeln!(
                &mut log_string,
                "{index}: {t}, {}",
                chain.history().get(index).unwrap()
            )?;
        }

        println!("Successfully logged to output/history");

        history_file.write_all(log_string.as_bytes())?;

        for label in chain.labels_of(chain.current())?.map(|(label, _)| label) {
            dbg!(label);
        }

        #[cfg(feature = "test-print-viz")]
        {
            chain.graph.print_viz("chain.gv")?;
            chain.atom.print_nfa("nfa.gv")?;
            item::default::print_labels(&chain.atom, &chain.forest)?;

            chain.forest.print_viz("forest.gv")?;

            chain.forest.print_closed_viz("closed.gv")?;
        }

        Ok(())
    }
}