pub mod traits; pub use traits::*;
pub mod relation; pub use relation::*;

pub use crate::noun::*;

use std::collections::{HashMap, HashSet};
use std::fmt::Display;
use itertools::Itertools;
use strum::IntoDiscriminant;

use enum_map::EnumMap;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Var(pub u64);
impl Display for Var {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "v{}", self.0)
    }
}

#[derive(Debug, Default)]
pub struct Unification {
    assignment: HashMap<Var, Noun>,
}
impl Unification {
    pub fn lookup(&self, v: &Var) -> Option<Noun> {
        self.assignment.get(v).cloned()
    }
    pub fn add(&mut self, v: Var, x: impl Into<Noun>) -> bool {
        let y = x.into();
        if let Some(old) = self.assignment.get(&v) {
            *old == y
        } else {
            self.assignment.insert(v, y);
            true
        }
    }
    pub fn unify_hole<A, B>(&mut self, h1: &Hole<A>, x2: &B) -> bool
    where Noun: From<A>, Noun: From<B>, A: Copy, B: Copy {
        match h1 {
            Hole::Val(x1) => Noun::from(*x1) == Noun::from(*x2),
            Hole::Var(v1) => self.add(*v1, *x2),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Hole<T> {
    Var(Var),
    Val(T),
}
impl<T: std::fmt::Debug> Display for Hole<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Var(v) => v.fmt(f),
            Self::Val(x) => write!(f, "{:?}", x),
        }
    }
}
impl<A: std::fmt::Debug> DisplayNoun for Hole<A> {
    fn display_noun(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self)
    }
}
impl<T: Clone + TryFrom<Noun>> Hole<T> {
    pub fn substitute(&self, u: &Unification) -> Self {
        match self {
            x@Self::Var(v) => if let Some(y) = u.lookup(v) && let Ok(z) = T::try_from(y) {
                Hole::Val(z)
            } else {
                x.clone()
            }
            x => x.clone(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct Rule {
    pub head: RelationHole,
    pub body: Vec<RelationHole>,
}
impl Display for Rule {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{} :- ", self.head)?;
        #[allow(unstable_name_collisions)]
        for mb in self.body.iter().map(Some).intersperse(None) {
            if let Some(b) = mb {
                write!(f, "{}", b)?;
            } else {
                write!(f, ", ")?;
            }
        }
        std::fmt::Result::Ok(())
    }
}
impl Rule {
    pub fn unify_body(&self, facts: &[&Relation]) -> Option<Unification> {
        if self.body.len() != facts.len() { return None; }
        let mut u = Unification::default();
        for (r, f) in self.body.iter().zip(facts) {
            u = r.unify(f, u)?;
        }
        Some(u)
    }
    pub fn canonize(&mut self, mut next: u64) -> u64 {
        let mut remap = HashMap::new();
        let mut f = |v| {
            if let Some(new) = remap.get(&v) {
                *new
            } else {
                let new = Var(next);
                next += 1;
                remap.insert(v, new);
                new
            }
        };
        self.head = self.head.map_variables(&mut f);
        for b in self.body.iter_mut() {
            *b = b.map_variables(&mut f);
        }
        next
    }
}

#[derive(Debug, Default, Clone)]
pub struct Facts {
    facts: EnumMap<RelationDiscriminants, HashSet<Relation>>,
}
impl Display for Facts {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for (_, s) in self.facts.iter() {
            for r in s.iter() {
                write!(f, "{}.\n", r)?;
            }
        }
        std::fmt::Result::Ok(())
    }
}
impl Facts {
    pub fn add(&mut self, f: Relation) -> bool {
        self.facts[f.discriminant()].insert(f)
    }
    pub fn with_sort(&self, s: RelationDiscriminants) -> &HashSet<Relation> {
        &self.facts[s]
    }
    pub fn with_sort_mut(&mut self, s: RelationDiscriminants) -> &mut HashSet<Relation> {
        &mut self.facts[s]
    }
    pub fn contains(&self, f: &Relation) -> bool {
        self.facts[f.discriminant()].contains(f)
    }
    pub fn absorb(&mut self, other: &mut Facts) {
        for (rd, fs) in self.facts.iter_mut() {
            fs.extend(other.with_sort_mut(rd).drain());
        }
    }
}

#[derive(Debug, Default, Clone)]
pub struct Database {
    facts: Facts,
    rules: Vec<Rule>,
}
impl Display for Database {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.facts)?;
        for r in self.rules.iter() {
            write!(f, "{}.\n", r)?;
        }
        std::fmt::Result::Ok(())
    }
}
impl Database {
    pub fn is_known(&self, r: &Relation) -> bool { self.facts.contains(r) }
    pub fn query<'a>(&'a mut self) -> Query<'a> {
        Query {
            db: self,
            facts: Facts::default(),
            rules: Vec::new(),
        }
    }
}

#[must_use]
#[derive(Debug)]
pub struct Query<'a> {
    db: &'a mut Database,
    facts: Facts,
    rules: Vec<Rule>,
}
#[derive(Debug, Clone, Copy)]
enum QueryFactsToConsider<'a> {
    LastIteration(&'a Facts),
    Query,
    DB,
}
impl<'a> Query<'a> {
    fn apply_rule(&self, r: &Rule, new_facts: &mut Facts, required_facts: QueryFactsToConsider) -> bool {
        // we don't need to consider every possible fact against each body term (naive evaluation)
        // instead, we can limit ourselves to the subset where at least one of the facts is new (semi-naive)
        let mut found = Vec::new();
        // identify which slots of this rule body possibly have new facts to consider
        let changed_slots = r.body.iter()
            .enumerate()
            .filter_map(|(i, r)| {
                match required_facts {
                    // if we are only considering facts from the provided set
                    QueryFactsToConsider::LastIteration(rf) =>
                        if rf.with_sort(r.discriminant()).is_empty() { None } else { Some(i) },
                    // if we are only considering facts added by this query
                    QueryFactsToConsider::Query =>
                        if self.facts.with_sort(r.discriminant()).is_empty() { None } else { Some(i) },
                    // if we are considering all facts, simply consider every index
                    QueryFactsToConsider::DB => Some(i),
                }
            });
        let empty = HashSet::new();
        for cs in changed_slots {
            let attempts = r.body.iter()
                .enumerate()
                .map(|(i, r)| {
                    let d = r.discriminant();
                    match required_facts {
                        // if this is an index that changed, limit it to only the provided set of facts
                        QueryFactsToConsider::LastIteration(rf) if i == cs =>
                             rf.with_sort(d).iter().chain(empty.iter()),
                        // (or the set of facts from this query)
                        QueryFactsToConsider::Query if i == cs =>
                             self.facts.with_sort(d).iter().chain(empty.iter()),
                        // otherwise, consider all facts - a new fact might be used alongside an unchanged fact!
                        _ => self.facts.with_sort(d).iter().chain(self.db.facts.with_sort(d).iter()),
                    }
                })
                .multi_cartesian_product();
            for fs in attempts {
                if let Some(u) = r.unify_body(&fs) {
                    found.push(u);
                }
            }
        }
        let mut changed = false;
        for u in found {
            if let Some(new) = r.head.substitute(&u).concretize()
                && !self.facts.contains(&new)
                && !self.db.facts.contains(&new)
            {
                if new_facts.add(new) { changed = true; }
            }
        }
        changed
    }
    fn apply_new_once(&mut self, new_facts: &mut Facts, required_facts: QueryFactsToConsider) -> bool {
        let mut changed = false;
        match &required_facts {
            QueryFactsToConsider::DB => {
                // if we are considering only DB facts, we only need to consider new rules
                for r in self.rules.iter() {
                    if self.apply_rule(r, new_facts, required_facts) { changed = true }
                }
            },
            _ => { // otherwise, we should consider all rules
                for r in self.rules.iter().chain(self.db.rules.iter()) {
                    if self.apply_rule(r, new_facts, required_facts) { changed = true }
                }
            }
        }
        changed
    }
    pub fn fact(mut self, r: Relation) -> Self {
        if !self.db.is_known(&r) {
            self.facts.add(r);
        }
        self
    }
    pub fn rule(mut self, mut r: Rule) -> Self {
        r.canonize(0);
        self.rules.push(r);
        self
    }
    fn run(&mut self) {
        let mut new_facts = Facts::default();
        // compare the new rules with the old facts
        self.apply_new_once(&mut new_facts, QueryFactsToConsider::DB);
        // compare the new+old rules with the new facts
        self.apply_new_once(&mut new_facts, QueryFactsToConsider::Query);
        let mut prev_facts = new_facts.clone();
        self.facts.absorb(&mut new_facts);
        while self.apply_new_once(&mut new_facts, QueryFactsToConsider::LastIteration(&prev_facts)) {
            prev_facts = new_facts.clone();
            self.facts.absorb(&mut new_facts);
        }
    }
    pub fn ask(mut self) -> Facts {
        self.run();
        self.facts
    }
    pub fn learn(mut self) {
        self.run();
        self.db.facts.absorb(&mut self.facts);
        self.db.rules.extend_from_slice(&mut self.rules);
    }
}