//! Patterns describe expectations to match against (yielding
//! Captures), or, symmetrically, templates to populate from
//! Substitutions.
use super::MetaVar;
use crate::ground::{Capture, Fact, Variable};
use std::convert::TryFrom;

/// A match pattern (or template...) is a slice of constant values and
/// references to metavariables.  When a reference appears multiple
/// times, it must match against (be populated with) the same ground
/// variable.
#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum Element {
    /// Constants must match the literal ground variable exactly.
    Constant(Variable),
    /// References may match any ground variable, but references with
    /// the same metavariable must match the same ground variable.
    Reference(MetaVar),
}

/// Instantiates a template with an environment mapping `MetaVar`s to
/// `Variable`s.
///
/// The template may refer to `MetaVar`s that are not in `environment`;
/// mappings for these variables will be generated by calling the
/// `backfill` argument.  That argument is `FnMut` because we expect
/// it to keep track of the `Variable`s that were created on demand.
#[must_use]
pub fn instantiate_template<VarGenerator, H: std::hash::BuildHasher>(
    template: &[Element],
    environment: &mut std::collections::HashMap<MetaVar, Variable, H>,
    backfill: &mut VarGenerator,
) -> Fact
where
    VarGenerator: FnMut(&MetaVar) -> Variable,
{
    template
        .iter()
        .map(|elt| match elt {
            Element::Constant(var) => *var,
            Element::Reference(mv) => *environment
                .entry(mv.clone())
                .or_insert_with(|| backfill(&mv)),
        })
        .collect::<Vec<_>>()
        .into()
}

/// Once instantiated, a pattern takes a Fact (which must match the
/// number of elements passed to `new`), and attempts to make it fit
/// the elements.  On success, the result is a capture.
pub struct Pattern {
    input_len: usize,
    output: Vec<MetaVar>,
    fun: Box<dyn Fn(&Fact) -> Option<Capture>>,
}

impl Pattern {
    /// Constructs a new pattern that attempts to match Facts against
    /// the `pattern` elements.
    ///
    /// # Errors
    ///
    /// Does not currently fail.
    pub fn new(pattern: &[Element]) -> Result<Self, &'static str> {
        Ok(make_pattern(pattern))
    }

    /// The input `Fact`s passed to `try_match` must have the same
    /// number of Variables as `input_len`, the number of Elements in
    /// the instantiation pattern.
    #[must_use]
    pub fn input_len(&self) -> usize {
        self.input_len
    }

    #[must_use]
    pub fn output(&self) -> &[MetaVar] {
        &self.output
    }

    #[inline]
    #[must_use]
    pub fn try_match(&self, fact: &Fact) -> Option<Capture> {
        (self.fun)(fact)
    }
}

/// A `Template` specifies a conversion from `Capture`s to `Fact`s.
pub struct Template {
    input: Vec<MetaVar>,
    output_len: usize,
    fun: Box<dyn Fn(&Capture) -> Fact>,
}

impl Template {
    /// Returns a `Template` from `Capture`s matching `inp` to `Fact`s
    /// matching `out`.
    ///
    /// # Errors
    ///
    /// Returns `Err` if `out` refers to a `MetaVar` absent from `inp`.
    pub fn new(inp: &[MetaVar], out: &[Element]) -> Result<Template, &'static str> {
        make_template(inp, out)
    }

    #[must_use]
    pub fn input(&self) -> &[MetaVar] {
        &self.input
    }

    #[must_use]
    pub fn output_len(&self) -> usize {
        self.output_len
    }

    #[inline]
    #[must_use]
    pub fn apply(&self, input: &Capture) -> Fact {
        (self.fun)(input)
    }
}

fn make_pattern(pattern: &[Element]) -> Pattern {
    // Represent indices as u8 or u32 (when u8 would be padded anyway)
    // for density.
    let mut constants = Vec::<(u32, Variable)>::new();
    let mut match_variables = Vec::<MetaVar>::new();

    for (index, elt) in pattern.iter().enumerate() {
        match elt {
            Element::Constant(var) => constants.push((
                u32::try_from(index).expect("Wide patterns not supported."),
                *var,
            )),
            Element::Reference(mv) => match_variables.push(mv.clone()),
        }
    }

    match_variables.sort();
    match_variables.dedup();

    // The first item in the tuple is the source index,
    // the last is the destination index.
    let mut match_indices = Vec::<(u8, u8)>::new();
    for (src_index, elt) in pattern.iter().enumerate() {
        if let Element::Reference(mv) = elt {
            let dst_index = match_variables.binary_search(mv).expect("must be found");
            match_indices.push((
                u8::try_from(src_index).expect("Wide source capture not supported"),
                u8::try_from(dst_index).expect("Wide destination capture not supported"),
            ));
        }
    }

    let num = pattern.len();
    let match_size = match_variables.len();
    let matcher = move |fact: &Fact| {
        let vars = fact.vars();

        assert_eq!(vars.len(), num);
        for (index, expected) in constants.iter().copied() {
            if vars[index as usize] != expected {
                return None;
            }
        }

        let mut ret = Vec::<Variable>::with_capacity(match_size);
        ret.resize(match_size, Variable::uninit());

        for (in_index, ret_index) in match_indices.iter().cloned() {
            let actual = vars[in_index as usize];
            let prev = ret[ret_index as usize];

            if prev != Variable::uninit() && actual != prev {
                return None;
            }

            ret[ret_index as usize] = actual;
        }

        Some(Capture::from_vec(ret))
    };

    Pattern {
        input_len: pattern.len(),
        output: match_variables,
        fun: Box::new(matcher),
    }
}

fn make_template(inp: &[MetaVar], pattern: &[Element]) -> Result<Template, &'static str> {
    let find_index = |needle: &MetaVar| {
        for (index, haystack) in inp.iter().enumerate() {
            if haystack == needle {
                return Ok(index);
            }
        }

        Err("Needle not found in the input haystack")
    };

    // Plop constants in place, and copy that base template for each
    // new instantiation.
    let mut base = Vec::<Variable>::with_capacity(pattern.len());
    base.resize(pattern.len(), Variable::uninit());

    // The `indices` vector tells us where to copy variables from the
    // source `Capture` to the output `Fact`.  The first index in each
    // pair is the source index, and the second the destination index.
    let mut substitutions = Vec::<(u8, u8)>::new();

    for (index, elt) in pattern.iter().enumerate() {
        match elt {
            Element::Reference(mv) => {
                let src_index = find_index(mv)?;
                substitutions.push((
                    u8::try_from(src_index).expect("Wide source captures not supported"),
                    u8::try_from(index).expect("Wide destination predicates not supported"),
                ));
            }
            Element::Constant(var) => base[index] = *var,
        }
    }

    let expected_input_len = inp.len();
    let instantiate = move |capture: &Capture| {
        let vars = capture.vars();
        let mut result = base.clone();

        assert_eq!(vars.len(), expected_input_len);
        for (src_index, dst_index) in substitutions.iter().copied() {
            result[dst_index as usize] = vars[src_index as usize];
        }

        result.into()
    };

    Ok(Template {
        input: inp.into(),
        output_len: pattern.len(),
        fun: Box::new(instantiate),
    })
}

#[test]
fn test_instantiate() {
    use std::collections::HashMap;

    let x = MetaVar::new("x");
    let y = MetaVar::new("y");
    let z = MetaVar::new("z");

    let template = [
        Element::Reference(x.clone()),
        Element::Reference(y.clone()),
        Element::Reference(z.clone()),
        Element::Reference(x.clone()),
        Element::Reference(z.clone()),
    ];

    let mut env: HashMap<MetaVar, Variable> =
        vec![(x.clone(), Variable::new(2)), (y.clone(), Variable::new(3))]
            .into_iter()
            .collect();

    let mut created = Vec::<Variable>::new();

    let mut backfill = |_mv: &MetaVar| {
        let var = Variable::new(10);
        created.push(var);
        var
    };

    let instance = instantiate_template(&template, &mut env, &mut backfill);
    assert_eq!(
        instance,
        [2, 3, 10, 2, 10]
            .iter()
            .map(|i| Variable::new(*i))
            .collect::<Vec<_>>()
            .into()
    );
    assert_eq!(created, vec![Variable::new(10)]);
}

#[test]
fn test_instantiate_with_constant() {
    use std::collections::HashMap;

    let x = MetaVar::new("x");
    let y = MetaVar::new("y");
    let z = MetaVar::new("z");

    let template = [
        Element::Constant(Variable::new(1)),
        Element::Reference(x.clone()),
        Element::Reference(y.clone()),
        Element::Reference(z.clone()),
        Element::Reference(x.clone()),
        Element::Reference(z.clone()),
    ];

    let mut env: HashMap<MetaVar, Variable> =
        vec![(x.clone(), Variable::new(2)), (y.clone(), Variable::new(3))]
            .into_iter()
            .collect();

    let mut created = Vec::<Variable>::new();

    let mut backfill = |_mv: &MetaVar| {
        let var = Variable::new(10);
        created.push(var);
        var
    };

    let instance = instantiate_template(&template, &mut env, &mut backfill);
    assert_eq!(
        instance,
        [1, 2, 3, 10, 2, 10]
            .iter()
            .map(|i| Variable::new(*i))
            .collect::<Vec<_>>()
            .into()
    );
    assert_eq!(created, vec![Variable::new(10)]);
}

#[test]
fn test_pattern_match_happy_path() {
    use super::Projection;

    let x = MetaVar::new("x");
    let y = MetaVar::new("y");

    let pattern = Pattern::new(&[
        Element::Reference(x.clone()),
        Element::Reference(y.clone()),
        Element::Reference(x.clone()),
    ])
    .expect("ok");

    assert_eq!(pattern.input_len(), 3);
    assert_eq!(pattern.output().len(), 2);

    let extract_x = Projection::new(pattern.output(), &[x.clone()]).expect("ok");
    let extract_y = Projection::new(pattern.output(), &[y.clone()]).expect("ok");

    let args: Fact = [2, 3, 2]
        .iter()
        .map(|i| Variable::new(*i))
        .collect::<Vec<_>>()
        .into();
    let extracted = pattern.try_match(&args).expect("matches");

    assert_eq!(extract_x.apply(&extracted), [Variable::new(2)].into());
    assert_eq!(extract_y.apply(&extracted), [Variable::new(3)].into());
}

#[test]
fn test_pattern_match_with_constant_happy_path() {
    use super::Projection;

    let x = MetaVar::new("x");
    let y = MetaVar::new("y");

    let pattern = Pattern::new(&[
        Element::Constant(Variable::new(1)),
        Element::Reference(x.clone()),
        Element::Reference(y.clone()),
    ])
    .expect("ok");

    assert_eq!(pattern.input_len(), 3);
    assert_eq!(pattern.output().len(), 2);

    let extract_x = Projection::new(pattern.output(), &[x.clone()]).expect("ok");
    let extract_y = Projection::new(pattern.output(), &[y.clone()]).expect("ok");

    let args: Fact = [1, 2, 3]
        .iter()
        .map(|i| Variable::new(*i))
        .collect::<Vec<_>>()
        .into();
    let extracted = pattern.try_match(&args).expect("matches");

    assert_eq!(extract_x.apply(&extracted), [Variable::new(2)].into());
    assert_eq!(extract_y.apply(&extracted), [Variable::new(3)].into());
}

#[test]
fn test_pattern_match_mismatch() {
    let x = MetaVar::new("x");
    let y = MetaVar::new("y");

    let pattern = Pattern::new(&[
        Element::Reference(x.clone()),
        Element::Reference(y.clone()),
        Element::Reference(x.clone()),
    ])
    .expect("ok");

    let args2: Fact = [2, 3, 4]
        .iter()
        .map(|i| Variable::new(*i))
        .collect::<Vec<_>>()
        .into();
    assert_eq!(pattern.try_match(&args2), None);
}

#[test]
fn test_pattern_match_mismatch_constant() {
    let x = MetaVar::new("x");
    let y = MetaVar::new("y");

    let pattern = Pattern::new(&[
        Element::Constant(Variable::new(1)),
        Element::Reference(x.clone()),
        Element::Reference(y.clone()),
    ])
    .expect("ok");

    let args: Fact = [2, 2, 3]
        .iter()
        .map(|i| Variable::new(*i))
        .collect::<Vec<_>>()
        .into();
    assert_eq!(pattern.try_match(&args), None);
}

#[test]
fn test_template_happy_path() {
    let x = MetaVar::new("x");
    let y = MetaVar::new("y");

    let template = Template::new(
        &[x.clone(), y.clone()],
        &[
            Element::Reference(x.clone()),
            Element::Reference(y.clone()),
            Element::Reference(x.clone()),
        ],
    )
    .expect("ok");

    assert_eq!(template.input(), &[x.clone(), y.clone()]);
    assert_eq!(template.output_len(), 3);

    assert_eq!(
        template.apply(&[Variable::new(2), Variable::new(3)].into()),
        [2, 3, 2]
            .iter()
            .map(|i| Variable::new(*i))
            .collect::<Vec<_>>()
            .into()
    );

    assert_eq!(
        template.apply(&[Variable::new(5), Variable::new(1)].into()),
        [5, 1, 5]
            .iter()
            .map(|i| Variable::new(*i))
            .collect::<Vec<_>>()
            .into()
    );
}

#[test]
fn test_template_with_constant_happy_path() {
    let x = MetaVar::new("x");
    let y = MetaVar::new("y");

    let template = Template::new(
        &[x.clone(), y.clone()],
        &[
            Element::Constant(Variable::new(1)),
            Element::Reference(x.clone()),
            Element::Reference(y.clone()),
            Element::Reference(x.clone()),
        ],
    )
    .expect("ok");

    assert_eq!(template.input(), &[x.clone(), y.clone()]);
    assert_eq!(template.output_len(), 4);

    assert_eq!(
        template.apply(&[Variable::new(2), Variable::new(3)].into()),
        [1, 2, 3, 2]
            .iter()
            .map(|i| Variable::new(*i))
            .collect::<Vec<_>>()
            .into()
    );

    assert_eq!(
        template.apply(&[Variable::new(5), Variable::new(1)].into()),
        [1, 5, 1, 5]
            .iter()
            .map(|i| Variable::new(*i))
            .collect::<Vec<_>>()
            .into()
    );
}

#[test]
fn test_template_missing_variable() {
    let x = MetaVar::new("x");
    let y = MetaVar::new("y");

    assert!(Template::new(
        &[x.clone(), x.clone()],
        &[
            Element::Reference(x.clone()),
            Element::Reference(y.clone()),
            Element::Reference(x.clone()),
        ],
    )
    .is_err());
}
