//! This module implements the generation and insertion of item
//! derivation forests.
//!
//! This is used for the chain-rule machine to conveniently produce
//! item derivations into a forest.  This forest can serve as a rough
//! approximation of the parse forests, and can be executed in other
//! semirings later on.

use super::*;
use crate::{
    atom::{Atom, DefaultAtom},
    default::Error,
    item::default::DefaultForest,
    Edge,
};
use grammar::{Error as GrammarError, GrammarLabel, GrammarLabelType, TNT};
#[allow(unused_imports)]
use graph::{builder::BuilderMut, labelled::binary::PLGBuilderMut, Graph, RedirectGraph};

use core::borrow::Borrow;

/// Convert an error telling us that an index is out of bounds.
///
/// # Panics
///
/// The function panics if the error is not of the expected kind.
pub(crate) fn index_out_of_bounds_conversion(ge: GrammarError) -> Error {
    match ge {
        GrammarError::IndexOutOfBounds(index, bound) => Error::IndexOutOfBounds(index, bound),
        _ => Error::Invalid,
    }
}

/// Determine if a label is labelled by a terminal.
fn is_labelled_by_terminal(label: GrammarLabelType) -> bool {
    matches!(label.tnt(), Some(t) if matches!(t, TNT::Ter(_)))
}

/// A helper function to generate a fragment of forest.
///
/// It simply constructs a root node and then appends
/// successive nodes as successive children of the previous
/// node.  Also the starting positions will all be set to the
/// same position.
///
/// If the input is empty, this returns an empty forest;
/// otherwise the result is not empty.
pub fn generate_fragment(
    labels: impl AsRef<[GrammarLabelType]>,
    pos: usize,
) -> Result<DefaultForest<ForestLabel<GrammarLabel>>, crate::default::Error> {
    let labels_slice = labels.as_ref();

    let labels_len = labels_slice.len();

    let last_label = if labels_len > 0 {
        labels_slice.get(labels_len - 1).copied().unwrap()
    } else {
        return Ok(Default::default());
    };

    let labels_iter = labels_slice.iter();

    let labels_iter_zipped = labels_iter
        .clone()
        .zip(labels_iter.skip(1).chain(std::iter::once(&last_label)));

    let mut mapped_iter = labels_iter_zipped.map(|(label, next_label)| {
        if is_labelled_by_terminal(*next_label) {
            GrammarLabel::new_closed(*label, pos, pos + 1)
        } else {
            GrammarLabel::new(*label, pos)
        }
    });

    let first_label = mapped_iter.next().unwrap();

    let mut result = DefaultForest::new_leaf(first_label);

    let mut index = 0;

    for label in mapped_iter {
        result.plant(index, DefaultForest::new_leaf(label), false)?;

        index = result
            .query_label(label.into())
            // REVIEW: Perhaps a LabelNoNode error?
            .ok_or(Error::Invalid)?;
    }

    Ok(result)
}

/// Generate a virtual fragment representing the left-linear null
/// closure \[nt\]^t.
pub fn virtual_generate_fragment(
    atom: impl Borrow<DefaultAtom>,
    nt: usize,
    t: usize,
    pos: usize,
) -> Result<DefaultForest<ForestLabel<GrammarLabel>>, crate::default::Error> {
    let atom = atom.borrow();

    let non_start = atom.nth_accumulator(nt).unwrap() * 2;

    let mut result = DefaultForest::default();

    for (label, child_iter) in atom.labels_of(non_start)? {
        if matches!(*label.get_value(),
                    Some(TNT::Ter(ter)) if ter == t)
        {
            for child in child_iter {
                let line: Vec<GrammarLabelType> = atom
                    .query_expansion(non_start, child)
                    .map_err(index_out_of_bounds_conversion)?
                    .iter()
                    .copied()
                    .flatten()
                    .flat_map(|(nt, rule)| [(*rule).into(), TNT::Non(*nt).into()])
                    .rev()
                    .chain(std::iter::once(TNT::Ter(t).into()))
                    .collect();

                if result.is_empty() {
                    result = generate_fragment(line, pos)?;
                } else {
                    let mut new_fragment = generate_fragment(line, pos)?;

                    new_fragment.remove_node(0)?;

                    new_fragment.set_root(1)?;

                    let cloned = result.clone_node(0, 0, false)?;

                    result.plant(cloned, new_fragment, false)?;
                }
            }
        }
    }

    Ok(result)
}

// TODO: Examine `insert_item` again.

impl DefaultForest<ForestLabel<GrammarLabel>> {
    /// Insert an item derivation forest into a recording forest.
    ///
    /// We need the help of other things just for finding the correct
    /// places to insert these item fragments.
    pub(crate) fn insert_item(
        &mut self,
        label: Edge,
        fragment: impl Borrow<DefaultForest<ForestLabel<GrammarLabel>>>,
        atom_child_iter: impl Iterator<Item = usize> + ExactSizeIterator + Clone,
        atom: &DefaultAtom,
    ) -> Result<PaVi, Error> {
        let root = if let Some(root) = self.root() {
            root
        } else {
            panic!("the forest must be non-empty when we insert items");
        };

        let pavi = label.forest_source();

        let true_source = label.true_source();

        let fragment = fragment.borrow();

        let fragment_root = if let Some(root) = fragment.root() {
            root
        } else {
            unreachable!("empty item");
        };

        let fragment_root_label = fragment
            .vertex_label(fragment_root)?
            .ok_or(Error::NodeNoLabel(fragment_root))?;

        let pos = fragment_root_label.label().start();

        if pavi != true_source {
            dbg!(label, pos);

            // Completing from true_source up to the pavi
            let true_source_node = match true_source {
                PaVi::Parent(node, edge) => {
                    let nth_child = self.nth_child(node, edge)?;
                    let nth_child_label = self
                        .vertex_label(nth_child)?
                        .ok_or(Error::NodeNoLabel(nth_child))?;

                    let nth_child_degree = self.degree(nth_child)?;
                    let nth_child_last = if nth_child_degree > 0 {
                        nth_child_degree - 1
                    } else {
                        0
                    };

                    if matches!(nth_child_label.label().label().tnt(), Some(TNT::Non(_)))
                        && !nth_child_label.is_packed()
                    {
                        let sploned = self.splone(nth_child, Some(pos), nth_child_last, false)?;

                        sploned
                    } else {
                        nth_child
                    }
                }
                PaVi::Virtual(nt, t, mut node) => {
                    let node_label_start = self
                        .vertex_label(node)?
                        .ok_or(Error::NodeNoLabel(node))?
                        .label()
                        .start();

                    let reduction_fragment = atom.generate_virtual_frags(nt, t, None);

                    // FIXME: We shall "close" this fragment as well.

                    for frag in reduction_fragment.into_iter().flatten() {
                        let mut frag = frag.clone();
                        frag.set_pos(node_label_start)?;

                        // NOTE: The function `plant_if_needed`
                        // assumes that we want to plant the fragment
                        // as the first child of the node.  This
                        // assumption holds in this case, but not in
                        // general.

                        node = self.plant_if_needed(node, frag)?;
                    }

                    node
                }
                PaVi::Empty => {
                    dbg!();
                    root
                }
            };

            let true_source_pos = self
                .vertex_label(true_source_node)?
                .ok_or(Error::NodeNoLabel(true_source_node))?
                .label()
                .start();

            let top_node = match pavi {
                PaVi::Parent(node, edge) => self.nth_child(node, edge)?,
                PaVi::Virtual(_nt, _t, _node) => {
                    dbg!(label);

                    self.print_viz("dbg forest.gv").unwrap();
                    fragment.print_viz("dbg fragment.gv").unwrap();

                    unreachable!("I do not think this is possible")
                }
                PaVi::Empty => {
                    unreachable!("The empty segment should not need to be reduced")
                }
            };

            let mut stack = vec![vec![(top_node, 0)]];
            let mut result_stack = Vec::new();

            while let Some(mut top_stack) = stack.pop() {
                let (node, _) = top_stack.pop().unwrap();

                let node_label = self.vertex_label(node)?.ok_or(Error::NodeNoLabel(node))?;

                if node_label.is_packed() {
                    for node in self.children_of(node)? {
                        let search_child = self.position_search(node, true_source_pos)?;

                        if let Some((child, index)) = search_child {
                            let mut top_stack = top_stack.clone();
                            // Fix the previous element
                            top_stack.push((node, index));

                            if child == true_source_node {
                                result_stack.push(top_stack);
                            } else {
                                top_stack.push((child, 0));

                                stack.push(top_stack);
                            }
                        }
                    }
                } else {
                    let search_child = self.position_search(node, true_source_pos)?;

                    if let Some((child, index)) = search_child {
                        top_stack.push((node, index));

                        if child == true_source_node {
                            result_stack.push(top_stack);
                        } else {
                            top_stack.push((child, 0));

                            stack.push(top_stack);
                        }
                    }
                }
            }

            // FIXME: We have to change the pavi as well, otherwise we
            // are going to use the wrong branch for planting later.

            for stack in result_stack {
                // dbg!(&stack);
                // self.print_viz("dbg forest.gv").unwrap();
                // panic!("test");

                let mut first_time = true;

                for (node, index) in stack.into_iter().rev() {
                    if matches!(
                        self.vertex_label(node)?
                            .ok_or(Error::NodeNoLabel(node))?
                            .label()
                            .label()
                            .tnt(),
                        Some(TNT::Non(_))
                    ) {
                        let sploned = self.splone(node, Some(pos), index, false)?;

                        if !first_time {
                            let last_index = self.degree(sploned)? - 1;

                            let last_child = self.nth_child(sploned, last_index)?;

                            let last_child_label = self
                                .vertex_label(last_child)?
                                .ok_or(Error::NodeNoLabel(last_child))?
                                .label();

                            if last_child_label.end() != Some(pos) {
                                let closed_label = GrammarLabel::new_closed(
                                    last_child_label.label(),
                                    last_child_label.start(),
                                    pos,
                                );

                                let closed_label_id = self
                                    .query_label(closed_label.into())
                                    .expect("last child closed label not found");

                                let mut builder = PLGBuilderMut::from_graph_mut(&mut self.graph);

                                builder.redirect(sploned, last_index, closed_label_id)?;
                            }
                        }

                        first_time = false;
                    }
                }
            }
        }

        // Ensure the last node in the PaVi is a terminal or a
        // non-terminal node, as an extra safety guard during
        // development.
        #[cfg(debug_assertions)]
        {
            match pavi {
                PaVi::Parent(node, edge) => {
                    assert!(matches!(
                        self.vertex_label(self.nth_child(node, edge)?),
                        Ok(Some(label))
                            if label.label().label().tnt().is_some()));
                }
                PaVi::Virtual(nt, t, node) => {
                    if !matches!(
                        self.vertex_label(node),
                        Ok(Some(label))
                            if matches!(
                                label.label().label().tnt(),
                                Some(TNT::Non(_))))
                    {
                        dbg!(node, self.vertex_label(node)?, pavi);

                        self.print_viz("dbg forest.gv").unwrap();

                        panic!("assumption fails");
                    }

                    if nt >= atom.non_num() {
                        dbg!();
                        return Err(Error::IndexOutOfBounds(nt, atom.non_num()));
                    }

                    if t >= atom.ter_num() {
                        dbg!();
                        return Err(Error::IndexOutOfBounds(t, atom.ter_num()));
                    }
                }
                PaVi::Empty => {}
            }
        }

        let is_empty_segment = pavi.is_empty();

        let mut parents: Vec<Parent> = {
            let mut result = Vec::new();

            match pavi {
                PaVi::Parent(node, edge) => {
                    result.push(Parent::new(node, edge));
                }
                PaVi::Virtual(nt, t, node) => {
                    let node_label = self.vertex_label(node)?.ok_or(Error::NodeNoLabel(node))?;

                    for atom_child in atom_child_iter.clone() {
                        for rule in atom.trace(nt, t, atom_child).into_iter().flatten() {
                            let virtual_frag = atom.generate_virtual_frags(nt, t, Some(rule));

                            if let Some(frag) = virtual_frag {
                                let mut frag = (*frag.get(0).unwrap()).clone();

                                frag.set_pos(node_label.label().start())?;

                                let frag_nodes_len = frag.nodes_len();

                                assert!(frag_nodes_len > 1);

                                let last_but_one_label = frag
                                    .vertex_label(frag_nodes_len - 2)?
                                    .ok_or(Error::NodeNoLabel(frag_nodes_len - 2))?;

                                // NOTE: The function
                                // `plant_if_needed` assumes that we
                                // want to plant the fragment as the
                                // first child of the node.  This
                                // assumption holds in this case, but
                                // not in general.

                                self.plant_if_needed(node, frag)?;

                                let rule_label_pos = self
                                    .query_label(last_but_one_label)
                                    .expect("the forest was wrongly planted");

                                result.push(Parent::new(rule_label_pos, 0));
                            }
                        }
                    }
                }
                PaVi::Empty => {
                    result.push(Parent::new(root, 0));
                }
            }

            result
        };

        for parent in parents.iter() {
            if !self.has_node(parent.node()) {
                return Err(Error::IndexOutOfBounds(parent.node(), self.nodes_len()));
            }
        }

        if !is_empty_segment {
            parents = parents
                .into_iter()
                .flat_map(|parent| {
                    self.parents_of(parent.node()).unwrap().filter(|n| {
                        matches!(
                            self.vertex_label(n.node())
                                .unwrap()
                                .unwrap()
                                .label()
                                .label()
                                .tnt(),
                            Some(TNT::Non(_))
                        )
                    })
                })
                .collect();
        }

        let mut non_empty = false;

        for atom_child in atom_child_iter {
            // dbg!(label.label(), atom_child);

            // Find reduction information.
            let reduction_info = atom
                .query_reduction(label.label(), atom_child)
                .map_err(index_out_of_bounds_conversion)?;

            let mut stack = parents.clone();
            let mut second_stack = Vec::new();

            // locate the nodes
            for reduction_nt in reduction_info.iter().copied().flatten().rev() {
                while let Some(mut node) = stack.pop() {
                    let mut node_label = self
                        .vertex_label(node.node())?
                        .ok_or_else(|| Error::NodeNoLabel(node.node()))?;

                    if matches!(
                        node_label
                            .label()
                            .label(),
                        GrammarLabelType::TNT(TNT::Non(nt)) if nt == *reduction_nt
                    ) {
                        let sploned_node =
                            self.splone(node.node(), Some(pos), node.edge(), false)?;

                        node_label = self
                            .vertex_label(sploned_node)?
                            .ok_or(Error::NodeNoLabel(sploned_node))?;

                        if node_label.clone_index().is_some() {
                            let mut parent_iter = self.parents_of(sploned_node)?;

                            #[cfg(debug_assertions)]
                            assert_eq!(parent_iter.len(), 1);

                            node = parent_iter.next().unwrap();

                            #[cfg(debug_assertions)]
                            assert!(self
                                .vertex_label(node.node())?
                                .ok_or(Error::NodeNoLabel(node.node()))?
                                .is_packed());
                        } else {
                            node = Parent::new(sploned_node, node.edge());
                        }

                        let parents_iter = self.parents_of(node.node())?;

                        for parent in parents_iter {
                            let parent_label = self
                                .vertex_label(parent.node())?
                                .ok_or_else(|| Error::NodeNoLabel(parent.node()))?
                                .label();

                            if parent_label.label().rule().is_none() {
                                crate::item::default::print_labels(atom, self.borrow()).unwrap();
                                self.print_viz("dbg forest.gv").unwrap();

                                dbg!(parent, parent_label, label, node, sploned_node);

                                panic!("assumption fails");
                            }

                            second_stack.extend(self.parents_of(parent.node())?.filter(|n| {
                                matches!(self.vertex_label(n.node()),
                                         Ok(Some(label))
                                         if matches!(
                                             label.label().label().tnt(),
                                             Some(TNT::Non(_))))
                            }));
                        }
                    }
                }

                std::mem::swap(&mut stack, &mut second_stack);

                if stack.is_empty() {
                    break;
                }
            }

            if stack.is_empty() {
                dbg!(
                    label,
                    atom_child,
                    parents,
                    reduction_info,
                    atom.query_reduction(label.label(), atom_child).unwrap(),
                    is_empty_segment,
                    atom.trace(0, 3, atom_child)
                        .into_iter()
                        .flatten()
                        .collect::<Vec<_>>(),
                );

                self.print_viz("dbg forest.gv").unwrap();

                #[cfg(test)]
                crate::item::default::print_labels(atom, self.borrow()).unwrap();

                return Err(Error::CannotPlant);
            }

            if pos == 4 && matches!(true_source, PaVi::Virtual(_, _, _)) {
                dbg!(&stack, reduction_info, true_source, pavi);
                self.print_viz(&format!("pos4ib.gv")).unwrap();
            }

            for parent in stack {
                let sploned_node = self.splone(parent.node(), None, parent.edge(), false)?;

                self.plant(sploned_node, fragment, non_empty)?;

                non_empty = true;
            }
        }

        // If the iterator is empty, assert the fragment has length
        // one, and do not plant anything.
        if !non_empty {
            assert_eq!(fragment.nodes_len(), 1);
        }

        let result = if fragment.nodes_len() == 2 {
            let root_label = fragment_root_label;
            let leaf_label = fragment
                .vertex_label(1 - fragment_root)?
                .ok_or(Error::NodeNoLabel(1 - fragment_root))?;

            // it has been planted, so should be safe.
            let node = self
                .query_label(root_label)
                .expect("root label was not found");

            let edge = {
                let mut result = None;

                for (index, child) in self.children_of(node)?.enumerate() {
                    if matches!(self.vertex_label(child)?, Some(child_label) if child_label == leaf_label)
                    {
                        result = Some(index);
                        break;
                    }
                }

                if let Some(result) = result {
                    result
                } else {
                    unreachable!("the forest is wrongly planted");
                }
            };
            // dbg!(node, edge, root_label, leaf_label);
            PaVi::Parent(node, edge)
        } else {
            assert_eq!(
                fragment.nodes_len(),
                1,
                "a virtual fragment should consist of a single terminal node."
            );

            let root_label = fragment_root_label;

            let pavi_parent = pavi.parent().expect(
                "When we insert a virtual fragment, the forest_source of
            the label must be a parent.",
            );

            let nth_child = self.nth_child(pavi_parent.node(), pavi_parent.edge())?;

            let nth_child_label = self
                .vertex_label(nth_child)?
                .ok_or(Error::NodeNoLabel(nth_child))?
                .label()
                .label();

            let error_str = "When we insert a virtual fragment, the \
                             forest source of the label must point to \
                             a non-terminal node";

            let nt = match nth_child_label.tnt().expect(error_str) {
                TNT::Non(nt) => nt,
                _ => {
                    dbg!(nth_child, nth_child_label);

                    panic!("{error_str}");
                }
            };

            let t = match root_label.label().label().tnt().unwrap() {
                TNT::Ter(t) => t,
                _ => {
                    dbg!(root_label);

                    panic!("a virtual fragment should consist of a single terminal node")
                }
            };

            PaVi::Virtual(nt, t, nth_child)
        };

        Ok(result)
    }
}

#[cfg(test)]
mod genins_test {
    use super::*;
    use crate::item::default::leaf;

    use grammar::test_grammar_helper::*;

    #[test]
    fn test_generate_fragment() -> Result<(), Box<dyn std::error::Error>> {
        let grammar = new_notes_grammar()?;

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

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

        let fragment = generate_fragment([72.into(), TNT::Non(0).into()], 0)?;

        let mut test_fragment = leaf!(
            GrammarLabel::new(GrammarLabelType::from(72), 0),
            GrammarLabel
        );

        test_fragment.plant(
            0,
            leaf!(
                GrammarLabel::new(GrammarLabelType::from(TNT::Non(0)), 0),
                GrammarLabel
            ),
            false,
        )?;

        assert_eq!(fragment, test_fragment);

        // virtual fragments

        println!("nt = 0, t = 3");

        let virtual_fragment = virtual_generate_fragment(&atom, 0, 3, 0)?;

        assert_eq!(virtual_fragment.nodes_len(), 7);

        let virtual_node = virtual_fragment.vertex_label(5)?.unwrap().label();

        let test_fragment = generate_fragment(
            [
                TNT::Non(0).into(),
                2.into(),
                TNT::Non(1).into(),
                8.into(),
                TNT::Non(2).into(),
                virtual_node.label(),
                TNT::Ter(3).into(),
            ],
            0,
        )?;

        crate::item::default::print_labels(&atom, &virtual_fragment)?;

        assert_eq!(virtual_fragment, test_fragment);

        #[cfg(feature = "test-print-viz")]
        virtual_fragment.print_viz("virtual fragment (0, 3).gv")?;

        println!("nt = 3, t = 2");

        let virtual_fragment = virtual_generate_fragment(&atom, 3, 2, 1)?;

        let test_fragment =
            generate_fragment([TNT::Non(3).into(), 38.into(), TNT::Ter(2).into()], 1)?;

        crate::item::default::print_labels(&atom, &virtual_fragment)?;

        assert_eq!(virtual_fragment, test_fragment);

        #[cfg(feature = "test-print-viz")]
        virtual_fragment.print_viz("virtual fragment (3, 2).gv")?;

        // querying reductions

        assert!(matches!(atom.query_reduction(17, 9), Ok(Some(&[1]))));

        assert!(matches!(atom.query_reduction(35, 9), Ok(Some(&[1, 2]))));
        assert!(matches!(atom.query_reduction(35, 25), Ok(Some(&[2]))));

        Ok(())
    }

    #[test]
    fn test_reduction() -> Result<(), Box<dyn std::error::Error>> {
        let grammar = new_paren_grammar()?;

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

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

        // querying reductions

        println!("{:?}", atom.query_reduction(32, 17)?);

        // assert!(matches!(atom.query_reduction(35, 25), Ok(Some(&[2]))));

        Ok(())
    }
}