lexington 0.3.0

use std::fmt::Debug;
use lexington::{Any,Lexer,Match,Matcher,Scanner,Token,Within};
use lexington::{ShiftReduceParser,ShiftReduceRule};
use Expr::*;
use Proposition::*;
use Partial::*;

// ===================================================================
// Formula
// ===================================================================

// A simple trait which captures the essence of logical negation for a
// given proposition.  There is an expectation that `¬¬p == p`.
trait LogicalNegation {
    fn not(self) -> Self;
}

// ===================================================================
// Expression
// ===================================================================

#[derive(Clone,Debug,Eq,PartialEq,PartialOrd,Ord)]
enum Expr {
    // An integer coefficient
    Int(isize),
    // A variable
    Var(String)
}

// ===================================================================
// Proposition
// ===================================================================

#[derive(Clone,Debug,Eq,PartialEq,PartialOrd,Ord)]
enum Proposition {
    True,
    False,
    /// lhs == rhs
    Equality(bool,Expr,Expr),
    /// lhs <= rhs
    Inequality(bool,Expr,Expr)        
}

impl Proposition {
    pub fn equals(lhs: Expr, rhs: Expr) -> Proposition {
        // Apply some optimisation
        match (&lhs,&rhs) {
            (l,r) if l == r => True,
            (Int(x),Int(y)) if x == y => True,
            (Int(x),Int(y)) if x != y => False,            
            _ => Equality(true,lhs,rhs)
        }
    }
    pub fn not_equals(lhs: Expr, rhs: Expr) -> Proposition {
        // Apply some optimisation
        match (&lhs,&rhs) {
            (l,r) if l == r => False,
            (Int(x),Int(y)) if x != y => True,
            (Int(x),Int(y)) if x == y => False,            
            _ => Equality(false,lhs,rhs)
        }
    }
    pub fn less_than(lhs: Expr, rhs: Expr) -> Proposition {
        // Apply some optimisation
        match (&lhs,&rhs) {
            (l,r) if l == r => False,
            (Int(x),Int(y)) if x < y => True,
            (Int(x),Int(y)) if x >= y => False,            
            _ => Inequality(false,rhs,lhs)
        }                
    }
    pub fn less_than_or_equals(lhs: Expr, rhs: Expr) -> Proposition {
        match (&lhs,&rhs) {
            (l,r) if l == r => True,
            (Int(x),Int(y)) if x <= y => True,
            (Int(x),Int(y)) if x > y => False,            
            _ => Inequality(true,lhs,rhs)
        }                        
    }
    pub fn greater_than(lhs: Expr, rhs: Expr) -> Proposition {
        Self::less_than(rhs,lhs)        
    }
    pub fn greater_than_or_equals(lhs: Expr, rhs: Expr) -> Proposition {
        Self::less_than_or_equals(rhs,lhs)
    }    
}

impl LogicalNegation for Proposition {
    fn not(self) -> Self {
	match self {
            Proposition::True => Proposition::False,
            Proposition::False => Proposition::True,
	    Proposition::Equality(b,l,r) => Self::Equality(!b,l,r),
	    Proposition::Inequality(b,l,r) => Self::Inequality(!b,l,r),            
	}
    }
}

// ===================================================================
// Formula
// ===================================================================

// A logical formula over arbitrary propositions `P` held in
// _Disjunctive Normal Form_.
#[derive(Clone,Debug,PartialEq)]
struct Formula {
    // Terms stored in the form (p2 ∧ p3 ∧ p4) ∨ (p1 ∧ ¬p3), etc.
    terms: Vec<Vec<Proposition>>
}

impl Formula {
    pub fn from_prop(prop: Proposition) -> Self {
	Self{terms: vec![vec![prop]]}
    }

    pub fn from_var<T:Into<String>>(negated: bool, v: T) -> Self {
        let p = if negated {
            Proposition::not_equals(Var(v.into()),Int(0))
        } else {
            Proposition::equals(Var(v.into()),Int(0))
        };
        Formula::from_prop(p)
    }

    // Compute the logical disjunction of two formulas.
    pub fn or(mut self, other: Self) -> Self {
	self.or_into(other);
	self
    }

    // Compute the logical intersection of two formulas.
    pub fn and(mut self, other: Self) -> Self {
	self.and_into(other);
	self	
    }
    
    // Compute the logical disjunction of two formulas.
    pub fn or_into(&mut self, other: Self) {
	self.terms.extend(other.terms);
        self.terms.sort();            
    }

    // Compute the logical intersection of two formulas.
    pub fn and_into(&mut self, other: Self) {
	if other.terms.len() == 1 {
	    let term = &other.terms[0];
	    // x ∧ (y ∨ z) => (x ∧ y) ∨ (x ∧ z)
	    for t in &mut self.terms {
		t.extend(term.clone());
                t.sort()
	    }
            self.terms.sort()
	} else {
            // Handle other case.  This could be more efficient no
            // doubt.
            let mut nterms = Vec::new();
            for term in &other.terms {
                let mut nts = self.terms.clone();
	        for t in &mut nts {
		    t.extend(term.clone());
                    t.sort()
	        }                
                nterms.extend(nts);
            }
            // Assign over
	    self.terms.clear();
            self.terms.extend(nterms);
            self.terms.sort()            
	}
    }

    // Compute the logical negation of a formula.
    pub fn not(&mut self) {
	todo!()
    }
}

impl From<Proposition> for Formula {
    fn from(prop: Proposition) -> Self {
	Self{terms: vec![vec![prop]]}
    }
}

// ===================================================================
// Partial Formula / Expression
// ===================================================================

/// A partial expression respresents an expression which is
/// potentially in the process of being parsed (i.e. may not have
/// beeN- fully parsed yet).
#[derive(Clone,Debug,PartialEq)]
enum Partial {
    Empty,
    // Partial Formulas
    OpenConjunct(Formula),
    OpenDisjunct(Formula),
    // Partial Propositions
    OpenEquality(bool,Vec<Expr>),
    OpenInequality(bool,bool,Vec<Expr>),
    // Completed Formula
    ClosedForm(Formula),
    ClosedExpr(Expr)    
}

impl Partial {
    // Evolve partial expression into a logical or.  Note that, in
    // some cases, this evolution should be prevented so as to avoid
    // ambiguous expressions such as `x ∨ y ∧ z`.
    fn into_or(self) -> Result<(bool,Self),()> {
	let r = match self {
            OpenDisjunct(_) => self,
            ClosedForm(f) => { OpenDisjunct(f) }
            ClosedExpr(Var(v)) => { OpenDisjunct(Formula::from_var(false,v)) }
	    // Rule does not apply (or is ambiguous)
            _ => { return Ok((false,self)); }
	};
	// Rule applied
	Ok((true,r))
    }

    // Evolve partial expression into a logical and.  Note that, in
    // some cases, this evolution should be prevented so as to avoid
    // ambiguous expressions such as `x ∧ y ∨ z`.
    fn into_and(self) -> Result<(bool,Self),()> {
	let r = match self {
            OpenConjunct(_) => self,
            ClosedForm(f) => { OpenConjunct(f) }
            ClosedExpr(Var(v)) => { OpenConjunct(Formula::from_var(false,v)) }
	    // Rule does not apply (or is ambiguous)
            _ => { return Ok((false,self)); }
	};
	// Rule applied
	Ok((true,r))
    }

    // Evolve partial expression into an equality.  Note that, in some
    // cases, this evolution should be prevented so as to avoid
    // ambiguous expressions such as `x != y == z`.
    fn into_equality(self, sign: bool) -> Result<(bool,Self),()> {
	match self {
            ClosedExpr(e) => Ok((true,OpenEquality(sign,vec![e]))),
	    // Rule does not apply (or is ambiguous)
            _ => Ok((false,self))
	}
    }
    
    // Evolve partial expression into an inequality.  Note that, in
    // some cases, this evolution should be prevented so as to avoid
    // ambiguous expressions such as `x <= y == z`.
    fn into_inequality(self, sign: bool, strict: bool) -> Result<(bool,Self),()> {
	match self {
            ClosedExpr(e) => Ok((true,OpenInequality(sign,strict,vec![e]))),
	    // Rule does not apply (or is ambiguous)
            _ => Ok((false,self))
	}
    }
    
    
    fn close(self) -> Self {
	match self {
            OpenConjunct(t)|OpenDisjunct(t) => ClosedForm(t),
            OpenEquality(p,mut es) if es.len()== 2 => {
                let rhs = es.pop().unwrap();
                let lhs = es.pop().unwrap();                
                let prop = Proposition::Equality(p,lhs,rhs);
                ClosedForm(Formula::from_prop(prop))
            },
            OpenInequality(lt,st,mut es) if es.len()== 2 => {
                let rhs = es.pop().unwrap();
                let lhs = es.pop().unwrap();                
                let prop = match (lt,st) {
                    // Greater-than-or-equals
                    (false,false) => Proposition::greater_than_or_equals(lhs,rhs),
                    // Less-than-or-equals
                    (true,false) => Proposition::less_than_or_equals(lhs,rhs),
                    // Greater-than
                    (false,true) => Proposition::greater_than(lhs,rhs),
                    // Less-than                    
                    (true,true) => Proposition::less_than(lhs,rhs),                    
                };
                ClosedForm(Formula::from_prop(prop))
            },
            _ => self
	}
    }

    fn reduce(mut self, other: Partial) -> Result<Self,()> {
        match (&mut self,&other) {
            (Empty,_) => Ok(other),
            (OpenConjunct(p),_) => { p.and_into(other.as_formula()); Ok(self) }
            (OpenDisjunct(p),_) => { p.or_into(other.as_formula()); Ok(self) }
            (OpenEquality(_,es),_) if es.len()==1 => { es.push(other.as_expr()); Ok(self) }
            (OpenInequality(_,_,es),_) if es.len()==1 => { es.push(other.as_expr()); Ok(self) }            
            (_,_) => todo!("Reducing {self:?} <= {other:?})")
        }
    }
    
    fn as_formula(self) -> Formula {
	match self {
	    ClosedForm(t) => t,
            ClosedExpr(Var(x)) => Formula::from_var(false,x),
	    _ => { panic!("cannot unwrap incomplete expression: {self:?}"); }                
	}
    }
    
    fn as_expr(self) -> Expr {
	match self {
	    ClosedExpr(e) => e,
	    _ => { panic!("cannot unwrap incomplete expression: {self:?}"); }                
	}
    }    
}

// ===================================================================
// Lexer
// ===================================================================

#[derive(Copy,Clone,Debug,PartialEq)]    
enum Kind {
    // Meta
    Eof,
    LeftBrace,
    RightBrace,
    WhiteSpace,
    // Formula
    LogicalAnd,
    LogicalOr,
    LogicalNot,
    // Proposition
    True,
    False,
    EqualsEquals,
    NotEquals,
    LessThan,
    LessEquals,
    GreaterThan,
    GreaterEquals,    
    // Expression
    Identifier,
    Number
}

/// Construct a very simple lexer for S-expressions, and scan them to
/// produce a list of zero or more tokens.
fn lex(input: &str) -> Vec<Token<Kind>> {
    // [ \n\t]+
    let whitespace = Any([' ','\n','\t']).one_or_more();    
    // [0..9]+
    let number = Within('0'..='9').one_or_more();
    // [a..zA..Z_]([0..9a..zA..Z_]*)
    let identifier_start = Within('a'..='z')
        .or(Within('A'..='Z')).or('_');
    let identifier_rest = Within('0'..='9').or(Within('a'..='z'))
        .or(Within('A'..='Z')).or('_').zero_or_more();
    let identifier = identifier_start.then(identifier_rest);
    // Construct scanner
    let scanner = Match(number,Kind::Number)
        .and_match('⊤',Kind::True)
        .and_match('⊥',Kind::False)
        .and_match('=',Kind::EqualsEquals)
        .and_match('≠',Kind::NotEquals)
        .and_match('≤',Kind::LessEquals)
        .and_match('<',Kind::LessThan)
        .and_match('≥',Kind::GreaterEquals)
        .and_match('>',Kind::GreaterThan)            
        .and_match('∨',Kind::LogicalOr)
        .and_match('∧',Kind::LogicalAnd)
        .and_match('(',Kind::LeftBrace)
        .and_match(')',Kind::RightBrace)
        .and_match(identifier,Kind::Identifier)        
        .and_match(whitespace, Kind::WhiteSpace)
        .eof(Kind::Eof);
    // Construct the lexer.
    Lexer::new(input,scanner).collect()
}

// ===================================================================
// Parser
// ===================================================================

fn parse<'a>(input: &'a str) -> Result<Formula,()> {
    // Lex the input
    let tokens = lex(input);
    // Rule for parsing strings into numbers
    let st = ShiftReduceParser::new()
        .apply(|l,r:Partial| l.reduce(r))
        .terminate(Kind::True,|_| ClosedForm(Formula::from(True)))
        .terminate(Kind::False,|_| ClosedForm(Formula::from(False)))
        .terminate(Kind::Number,|tok| ClosedExpr(Int(input[tok.range()].parse().unwrap())))        
        .terminate(Kind::Identifier,|tok| {
            let name = input[tok.range()].to_string();
            ClosedExpr(Expr::from(Var(name)))
        })        
        .open(Kind::LeftBrace, Empty)
        .update_with(Kind::LogicalAnd,|e| e.into_and())
        .update_with(Kind::LogicalOr,|e| e.into_or())
        .update_with(Kind::EqualsEquals,|e| e.into_equality(true))
        .update_with(Kind::NotEquals,|e| e.into_equality(false))
        .update_with(Kind::LessEquals,|e| e.into_inequality(true,false))
        .update_with(Kind::LessThan,|e| e.into_inequality(true,true))
        .update_with(Kind::GreaterEquals,|e| e.into_inequality(false,false))
        .update_with(Kind::GreaterThan,|e| e.into_inequality(false,true))		
        .close_with(Kind::RightBrace, |e| e.close())
        .close_with(Kind::Eof, |e| e.close())
        .skip(Kind::WhiteSpace)        
        .first(Empty)
        .parse(tokens)?;
    // Extract the term
    Ok(st.as_formula())
}

fn check_ok<T:Into<Formula>>(input: &str, expecting: T) {
    let expected = expecting.into();
    // Lex into tokens and construct parser
    //let mut parser = Parser::new(lex(input));    
    // Parse tokens into S-expression
    let actual = parse(input).unwrap();
    // Check what we got
    assert_eq!(actual,expected);
}

fn check_eq(lhs: &str, rhs: &str) {
    // Lex into tokens and construct parser
    //let mut parser = Parser::new(lex(input));    
    // Parse tokens into S-expression
    let l = parse(lhs).unwrap();
    let r = parse(rhs).unwrap();    
    // Check what we got
    assert_eq!(l,r);
}

#[test]
fn logic_01() {
    check_ok("⊥", False);
    check_ok("⊤", True);    
}

// =============================================================================
// Conjunction
// =============================================================================

#[test]
fn logic_and_01() {
    check_ok("⊤ ∧ ⊥",Formula::from(True).and(Formula::from(False)));
}

// =============================================================================
// Disjunction
// =============================================================================

#[test]
fn logic_or_01() {
    check_ok("⊤ ∨ ⊤",Formula::from(True).or(Formula::from(True)));
}

#[test]
fn logic_or_02() {
    let a = Formula::from_var(false,"a");    
    let x = Formula::from_var(false,"x");
    let y = Formula::from_var(false,"y");
    let a_x = a.clone().and(x);
    let a_y = a.and(y);
    check_ok("(a ∧ x) ∨ (a ∧ y)",a_x.or(a_y));
}

#[test]
fn logic_or_03() {
    let a = Formula::from_var(false,"a");
    let b = Formula::from_var(false,"b");    
    let x = Formula::from_var(false,"x");
    let y = Formula::from_var(false,"y");
    let a_x = a.clone().and(x.clone());
    let a_y = a.and(y.clone());
    let b_x = b.and(x);
    check_ok("(a ∧ x) ∨ (a ∧ y) ∨ (b ∧ x)",a_x.or(a_y).or(b_x));
}

#[test]
fn logic_or_04() {
    let a = Formula::from_var(false,"a");
    let b = Formula::from_var(false,"b");    
    let x = Formula::from_var(false,"x");
    let y = Formula::from_var(false,"y");
    let a_x = a.clone().and(x.clone());
    let a_y = a.and(y.clone());
    let b_x = b.clone().and(x);
    let b_y = b.and(y);    
    check_ok("(a ∧ x) ∨ (a ∧ y) ∨ (b ∧ x) ∨ (b ∧ y)",a_x.or(a_y).or(b_x).or(b_y));
}

// =============================================================================
// Equality
// =============================================================================

#[test]
fn logic_eq_01() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    check_ok("x=0",Proposition::equals(x,z))
}

#[test]
fn logic_eq_02() {
    let x = Expr::Var("x".to_string());
    let c = Expr::Int(123);
    check_ok("123=x",Proposition::equals(c,x))
}

#[test]
fn logic_eq_03() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    check_ok("x≠0",Proposition::not_equals(x,z))
}

#[test]
fn logic_eq_04() {
    let x = Expr::Var("x".to_string());
    let c = Expr::Int(123);
    check_ok("123≠x",Proposition::not_equals(c,x))
}

#[test]
fn logic_eq_05() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    check_ok("⊤ ∨ (x=0)",Formula::from(True).or(Proposition::equals(x,z).into()))
}

#[test]
fn logic_eq_06() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    let f : Formula = Proposition::equals(x,z).into();
    check_ok("(x=0) ∨ ⊥",f.or(False.into()))
}

#[test]
fn logic_eq_07() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    check_ok("⊤ ∧ (x=0)",Formula::from(True).and(Proposition::equals(x,z).into()))
}

#[test]
fn logic_eq_08() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    let f : Formula = Proposition::equals(x,z).into();
    check_ok("(x=0) ∧ ⊥",f.and(False.into()))
}

// =============================================================================
// Inequality
// =============================================================================

#[test]
fn logic_ieq_01() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    check_ok("x≤0",Proposition::less_than_or_equals(x,z))
}

#[test]
fn logic_ieq_02() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    check_ok("x>0",Proposition::greater_than(x,z))
}

#[test]
fn logic_ieq_03() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    check_ok("x≥0",Proposition::greater_than_or_equals(x,z))
}

#[test]
fn logic_ieq_04() {
    let x = Expr::Var("x".to_string());
    let z = Expr::Int(0);
    check_ok("x<0",Proposition::less_than(x,z))
}

// =============================================================================
// Equivalences
// =============================================================================

// #[test]
// fn logic_equiv_01() {
//     check_eq("⊤ ∧ ⊥","⊥");
// }

// #[test]
// fn logic_equiv_02() {
//     check_eq("⊤ ∨ ⊥","⊤");
// }

#[test]
fn logic_equiv_03() {
    check_eq("0 < 0","⊥");
}

#[test]
fn logic_equiv_04() {
    check_eq("0 < 1","⊤");
}

#[test]
fn logic_equiv_05() {
    check_eq("0 ≤ 0","⊤");
}

#[test]
fn logic_equiv_06() {
    check_eq("1 ≤ 0","⊥");
}

#[test]
fn logic_equiv_07() {
    check_eq("x ∧ (y ∨ z)","(x∧y) ∨ (x∧z)");
}

#[test]
fn logic_equiv_08() {
    check_eq("(a ∨ b) ∧ (x ∨ y)","(a ∧ (x ∨ y)) ∨ (b ∧ (x ∨ y))");
    check_eq("(a ∨ b) ∧ (x ∨ y)","(a ∧ x) ∨ (b ∧ x) ∨ (a ∧ y) ∨ (b ∧ y)");    
}