//! 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::DefaultAtom, default::Error, item::default::DefaultForest, Edge};
use grammar::{Error as GrammarError, GrammarLabel, GrammarLabelType, TNT};
use graph::Graph;

use core::borrow::Borrow;
use std::collections::HashSet as Set;

/// 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.
fn index_out_of_bounds_conversion(ge: GrammarError) -> Error {
    match ge {
        GrammarError::IndexOutOfBounds(index, bound) => Error::IndexOutOfBounds(index, bound),
        _ => Error::Invalid,
    }
}
/// 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 mut labels_iter = labels.as_ref().iter();
    let mut result: DefaultForest<ForestLabel<GrammarLabel>>;

    if let Some(first_label) = labels_iter.next() {
        result = DefaultForest::new_leaf(GrammarLabel::new(*first_label, pos));

        let mut index = 0;

        for label in labels_iter {
            result.plant(
                index,
                DefaultForest::new_leaf(GrammarLabel::new(*label, pos)),
                false,
            )?;

            // To prevent duplication of labels causing
            // panics, we query the index of the new node.
            index = result
                .query_label(GrammarLabel::new(*label, pos).into())
                // REVIEW: Perhaps a LabelNoNode error?
                .ok_or(Error::Invalid)?;
        }
    } else {
        result = Default::default();
    }

    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 {
                // #[allow(unused_must_use)]
                // {
                //     dbg!(non_start, child, atom.query_expansion(non_start, child));
                // }

                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 {
                    result.plant(0, generate_fragment(line, pos)?, false)?;
                }
            }
        }
    }

    Ok(result)
}

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 fn insert_item(
        &mut self,
        label: Edge,
        fragment: impl Borrow<DefaultForest<ForestLabel<GrammarLabel>>>,
        atom_child_iter: impl Iterator<Item = usize>,
        atom: &DefaultAtom,
    ) -> Result<PaSe, Error> {
        let fragment = fragment.borrow();

        let pase = label.forest_source();

        let _fragment_root = if let Some(root) = fragment.root() {
            root
        } else {
            return Err(Error::EmptyItem);
        };

        // Ensure the last node in the PaSe is a terminal or a
        // non-terminal node, as an extra safety guard during
        // development.
        #[cfg(debug_assertions)]
        {
            if let Some(parent) = pase.parent() {
                assert!(matches!(
                    self.vertex_label(self.nth_child(parent.node(), parent.edge())?),
                    Ok(Some(label))
                        if label.label().label().tnt().is_some()));
            } else if let Some((_, leaf)) = pase.segment() {
                assert!(matches!(
                    self.vertex_label(leaf),
                    Ok(Some(label))
                        if label.label().label().tnt().is_some()));
            } else {
                unreachable!("a pase is neither a parent nor a segment");
            }
        }

        let mut is_empty_segment = false;

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

            if let Some(parent) = pase.parent() {
                result.push(parent);
            } else {
                let (root, leaf) = pase.segment().unwrap();
                let mut seen_nodes = Set::new();

                let mut stack = if root == leaf {
                    // an empty segment means the root
                    is_empty_segment = true;

                    result.push(Parent::new(root, 0));
                    vec![]
                } else {
                    vec![root]
                };

                while let Some(top) = stack.pop() {
                    if seen_nodes.contains(&top) {
                        continue;
                    }

                    seen_nodes.insert(top);

                    for (index, child) in self.children_of(top)?.enumerate() {
                        if child == leaf {
                            result.push(Parent::new(top, index));
                        } else {
                            stack.push(child);
                        }
                    }
                }
            }

            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(node) = stack.pop() {
                    // REVIEW: A lot of labels appear here.
                    // Perhaps I need to refactor something?
                    if matches!(
                        self
                            .vertex_label(node.node())?
                            .ok_or_else(|| Error::NodeNoLabel(node.node()))?
                            .label()
                            .label(),
                        GrammarLabelType::TNT(TNT::Non(nt)) if nt == *reduction_nt
                    ) {
                        for parent in self.parents_of(node.node())? {
                            debug_assert!(matches!(
                                self.vertex_label(parent.node()),
                                Ok(Some(label)) if
                                    label
                                    .label()
                                    .label()
                                    .rule()
                                    .is_some()));

                            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()
                );

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

                #[cfg(test)]
                genins_test::print_labels(atom, self).unwrap();

                return Err(Error::CannotPlant);
            }

            for parent in stack.into_iter() {
                if parent.edge() + 1 >= self.degree(parent.node())? {
                    // dbg!(&fragment);
                    self.plant(parent.node(), fragment, non_empty)?;
                } else {
                    let nth_child = self.nth_child(parent.node(), parent.edge() + 1)?;

                    if self.is_prefix(nth_child, fragment)? {
                        // do nothing
                        continue;
                    }
                    dbg!("clone?");

                    let cloned_node = self.clone_node(nth_child)?;

                    self.plant(cloned_node, fragment, non_empty)?;
                }

                non_empty = true;
            }
        }

        // If the iterator is empty, just plant at the specified
        // child.
        if !non_empty {
            for parent in parents.into_iter() {
                let nth_child = self.nth_child(parent.node(), parent.edge())?;

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

                non_empty = true;
            }
        }

        let result = if fragment.nodes_len() == 2 {
            let root_label = fragment.vertex_label(0)?.unwrap();
            let leaf_label = fragment.vertex_label(1)?.unwrap();

            // it has been planted, so should be safe.
            let node = self.query_label(root_label).unwrap();

            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);
            PaSe::Parent(node, edge)
        } else {
            let root_label = fragment.vertex_label(0)?.unwrap();
            let leaf_label = fragment.vertex_label(fragment.nodes_len() - 1)?.unwrap();
            // dbg!(root_label, leaf_label);

            PaSe::Segment(
                self.query_label(root_label).unwrap(),
                self.query_label(leaf_label).unwrap(),
            )
        };

        Ok(result)
    }
}

#[cfg(test)]
mod genins_test {
    use super::*;
    #[allow(unused_imports)]
    use crate::{default::DefaultChain, item::default::leaf, Chain};

    use grammar::test_grammar_helper::*;

    pub fn print_labels(
        atom: impl Borrow<DefaultAtom>,
        forest: impl Borrow<DefaultForest<ForestLabel<GrammarLabel>>>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let forest = forest.borrow();
        let atom = atom.borrow();

        for node in forest.nodes() {
            let label = forest.vertex_label(node)?;

            if let Some(label) = label {
                let label = label.label.label();

                match label {
                    GrammarLabelType::TNT(tnt) => {
                        println!("{tnt} = {}", atom.name_of_tnt(tnt)?);
                    }
                    GrammarLabelType::Rule(pos) => {
                        println!("pos {pos} = {}", atom.rule_pos_string(pos)?);
                    }
                }
            } else {
                return Err(Error::NodeNoLabel(node).into());
            }
        }

        Ok(())
    }

    #[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,
        )?;

        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)?;

        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(())
    }
}