asklyphe/asklyphe-frontend/src/math.rs

use tracing::{debug, error};
use once_cell::sync::Lazy;
use astro_float::{BigFloat, Sign, RoundingMode, Consts};
use std::str::FromStr;
use std::sync::{Arc, Mutex};
use std::mem;

pub const PRECISION: usize = 2048;
static CONST_CACHE: Lazy<Arc<Mutex<Consts>>> = Lazy::new(|| Arc::new(Mutex::new(Consts::new().expect("Unable to allocate memory for Conts cache"))));
// static PI: Lazy<BigFloat> = Lazy::new(|| BigFloat::from_str("3.141592653589793238462643383279").unwrap());
// static E: Lazy<BigFloat> = Lazy::new(|| BigFloat::from_str("2.718281828459045235360287471352").unwrap());


#[derive(Debug)]
pub struct Calculation {
	pub equation: String,
	pub result: String,
}

pub fn calculate(query: &str) -> Option<Calculation> {
	debug!("Got query {}", query);
	let mut parser = Parser::new(Lexer::new(query));
	let mut tree = parser.parse()?;
	let res = tree.eval();

	debug!("Calculation: {}", query);
	// debug!("Tree: {:?}", tree);
	debug!("Result: {}", res);
	Some(Calculation {equation: query.to_string(), result: res.to_string()})
}

// TODO: put into own crate with dependency astro-float = "0.9.2" so I can use more than BigFloat
#[derive(Debug, Clone, PartialEq)]
enum Token {
	Op(Op),
	Atom(Atom),
/*	Number(BigFloat),
	Func(Func),*/
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum Op {
	Add,
	Subtract,
	Multiply,
	Divide,
	Exponent,
	LParen,
	RParen,
	Func(Func), // A function is an Op that takes whatever the next thing is and binds it, either the next number or whatever is in parens
}

impl Op {
	fn bp_infix(&self) -> Option<(f64, f64)> {
		match self {
			// Op::LParen => Some(0.0),
			// Op::RParen => Some(0.0),
			Op::Add => Some((1.0, 1.1)),
			Op::Subtract => Some((1.0, 1.1)),
			Op::Multiply => Some((2.0, 2.1)),
			Op::Divide => Some((2.0, 2.1)),
			Op::Exponent => Some((3.1, 3.0)),
			_ => None,
			// Op::Func(_) => 0.0, // TODO: decide if this is a good LBP
		}
	}

	fn bp_prefix(&self) -> Option<f64> {
		match self {
			Op::Func(_) => Some(6.0),
			Op::Subtract => Some(5.0),
			Op::Add => Some(5.0),
			_ => None,
		}
	}

	fn apply_to(&self, args: &mut Vec<Expr>) -> BigFloat {
		match args.len() {
			1 => match self {
				Op::Subtract => {
					let mut res = args[0].eval();
					res.set_sign(Sign::Neg);
					res
				},
				Op::Add => {
					let mut res = args[0].eval();
					res.set_sign(Sign::Pos);
					res
				}
				Op::Func(f) => match f {
					Func::Sine => args[0].eval().sin(PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::Cosine => args[0].eval().cos(PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::Tangent => args[0].eval().tan(PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::ArcSine => args[0].eval().asin(PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::ArcCosine => args[0].eval().acos(PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::ArcTangent => args[0].eval().atan(PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::Log2 => args[0].eval().log2(PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::Log10 => args[0].eval().log10(PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::LogN => args[0].eval().ln(PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::Square => args[0].eval().pow(&BigFloat::from_f64(2.0, PRECISION), PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
					Func::SquareRoot => args[0].eval().sqrt(PRECISION, RoundingMode::None),
					Func::Abs => args[0].eval().abs(),
					_ => {
						error!("Got 1 params for func {:?} which expects 2 (should not be possible)", self);
						astro_float::NAN
					},
				},
				_ => {
					error!("Got 1 params for {:?} which expects 2 (should not be possible)", self);
					astro_float::NAN
				},
			}
			2 => match self {
				Op::LParen => args[0].eval(),
				Op::RParen => args[0].eval(),
				Op::Add => args[0].eval().add(&mut args[1].eval(), PRECISION, RoundingMode::None),
				Op::Subtract => args[0].eval().sub(&mut args[1].eval(), PRECISION, RoundingMode::None),
				Op::Multiply => args[0].eval().mul(&mut args[1].eval(), PRECISION, RoundingMode::None),
				Op::Divide => args[0].eval().div(&mut args[1].eval(), PRECISION, RoundingMode::None),
				Op::Exponent => args[0].eval().pow(&mut args[1].eval(), PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
				Op::Func(Func::Log) => args[0].eval().log(&mut args[1].eval(), PRECISION, RoundingMode::None, &mut CONST_CACHE.lock().unwrap()),
				_ => {
					error!("Got 2 params for {:?} which only expects 1 (should not be possible)", self);
					astro_float::NAN
				},
			}
			_ => {
				error!("Unexpected number of params ({}) for {:?} (should not be possible)", args.len(), self);
				astro_float::NAN
			},
		}
	}
}

#[derive(Debug, Clone, PartialEq)]
enum Atom {
	Number(BigFloat),
	Const(Const),
}

/*impl Atom {
	fn get_val(&self) -> BigFloat {
		match self {
			Atom::Number(val) => *val,
			Atom::Const(c) => match c {
				Const::Pi => CONST_CACHE.lock().unwrap().pi(PRECISION, RoundingMode::None),
				Const::E => CONST_CACHE.lock().unwrap().e(PRECISION, RoundingMode::None),
				Const::Inf => astro_float::INF_POS,
				Const::Nan => astro_float::NAN,
			}
		}
	}
}*/

impl Const {
	fn get_val(&self) -> BigFloat {
		match self {
			Const::Pi => CONST_CACHE.lock().unwrap().pi(PRECISION, RoundingMode::None),
			Const::E => CONST_CACHE.lock().unwrap().e(PRECISION, RoundingMode::None),
			Const::Inf => astro_float::INF_POS,
			Const::Nan => astro_float::NAN,
		}
	}
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum Func {
	Sine,
	Cosine,
	Tangent,
	// sin-1, cos-1, tan-1
	ArcSine,
	ArcCosine,
	ArcTangent,
	Log2,
	Log10,
	LogN,
	Log,
	Square,
	SquareRoot,
	Abs,
}

impl Func {
	fn names() -> &'static [(Func, &'static [&'static str])] {
		&[
			(Func::Sine, &["sin", "sine"]),
			(Func::Cosine, &["cos", "cosine"]),
			(Func::Tangent, &["tan", "tangent"]),
			(Func::ArcSine, &["asin", "asine", "arcsin", "arcsine"]),
			(Func::ArcCosine, &["acos", "acosine", "arccos", "arccosine"]),
			(Func::ArcTangent, &["atan", "atangent", "arctan", "arctangent"]),
			(Func::Log2, &["log2"]),
			(Func::Log10, &["log10"]),
			(Func::LogN, &["ln", "logn"]),
			(Func::Log, &["log"]),
			(Func::Square, &["square", "squared"]),
			(Func::SquareRoot, &["sqrt", "squareroot", "√"]),
			(Func::Abs, &["abs", "absolute"]),
		]
	}
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum Const {
	Pi,
	E,
	Inf,
	Nan,
}


impl Const {
	fn names() -> &'static [(Const, &'static [&'static str])] {
		&[
			(Const::Pi, &["pi", "PI", "π"]),
			(Const::E, &["e", "euler"]),
			(Const::Inf, &["inf", "infinity", "∞"]),
			(Const::Nan, &["nan", "NaN"])
		]
	}
}


#[derive(Debug, Copy, Clone, PartialEq)]
enum ParseErr {
	Eof,
	Invalid,
}

// this can probably be swapped out with a lexer generator like Logos if needed

struct Lexer<'a> {
	data: &'a str,
	data_ptr: &'a str,
	// idx: usize,
	next_by: usize,
	next_tok: Result<Token, ParseErr>,
}

// TODO: refactor with iterator that returns Option(Token) where one token option is Eof (or a enum of Token(Token) and Eof, or just Option(Option(Token)))
impl Lexer<'_> {

	fn new(data: &str) -> Lexer<'_> {
		let mut n: Lexer = Lexer {data, data_ptr: data, next_by: 0, next_tok: Err(ParseErr::Eof)};
		n.next();
		debug!("New finished!");
		n
	}
	
	fn _next(&mut self) -> Result<Token, ParseErr> {
		self.data_ptr = &self.data_ptr[self.next_by..];
		match self.data_ptr.chars().nth(0) {
			Some(val) => {
				debug!("lexing char '{}' at idx {}", val, self.data.chars().count() - self.data_ptr.chars().count());
				self.next_by = 1;
				match val {
					'+' => Ok(Token::Op(Op::Add)),
					'-' => Ok(Token::Op(Op::Subtract)),
					'×' | '*' => Ok(Token::Op(Op::Multiply)),
					'÷' | '/' => Ok(Token::Op(Op::Divide)),
					'^' => Ok(Token::Op(Op::Exponent)),
					'(' => Ok(Token::Op(Op::LParen)),
					')' => Ok(Token::Op(Op::RParen)),
					_ if val.is_whitespace() => self._next(),
					_ if val.is_digit(10) => {
						let mut len: usize = 0;
						self.data_ptr.chars().take_while(|c| c.is_digit(10) || *c == '.').for_each(|_| len += 1);
						self.next_by = len;

						match self.data_ptr[..len].parse() {
							Ok(val) => Ok(Token::Atom(Atom::Number(val))),
							Err(e) => Err(ParseErr::Invalid),
						}
					},
					_ => {
						let len = self.data_ptr.chars().count();
						for (f, names) in Func::names() {
							for name in *names {
								let n_len = name.chars().count();
								if self.data_ptr.starts_with(name) && (len == n_len || !self.data_ptr.chars().nth(n_len).unwrap().is_alphabetic()) {
									self.next_by = name.chars().count();
									return Ok(Token::Op(Op::Func(*f)));
								}
							}
						}
						for (f, names) in Const::names() {
							for name in *names {
								let n_len = name.chars().count();
								if self.data_ptr.starts_with(name) && (len == n_len || !self.data_ptr.chars().nth(n_len).unwrap().is_alphabetic()) {
									self.next_by = name.chars().count();
									return Ok(Token::Atom(Atom::Const(*f)));
								}
							}
						}
						debug!("got invalid char '{}'", val);
						Err(ParseErr::Invalid)
					}
				}
			}
			None => {
				self.next_by = 0;
				Err(ParseErr::Eof)
			},
		}
	}

	fn next(&mut self) -> Result<Token, ParseErr> {
		let res = self._next();
		let val = mem::replace(&mut self.next_tok, res);
		// self.next_tok = self._next();
		val
	}

	fn peek(&mut self) -> &Result<Token, ParseErr> {
		&self.next_tok
	}

	// TODO: replace with iterator so I can do parser.parse(lexer.iter()) and parse does lex_iter.next() & such
	fn lex_all(&mut self) -> Option<Vec<Token>> {
		let mut tokens: Vec<Token> = vec![];
		loop {
			match self.next() {
				Err(ParseErr::Eof) => return Some(tokens),
				Err(ParseErr::Invalid) => return None,
				Ok(tok) => tokens.push(tok),
			}
			// debug!("tokens: {:?}", tokens);
		}
	}
}

fn matches(s: &str, check: &str) -> usize {
	// debug!("s: \"{}\", check: \"{}\"c_len: {}, s_len: {}, s[c_len]: {:?}, s[c_len + 1]: {:?}", s, check, check.chars().count(), s.chars().count(), s.chars().nth(check.chars().count()), s.chars().nth(check.chars().count() + 1));
	match (s.chars().count(), check.chars().count()) {
		(s_len, c_len) if s_len < c_len => 0,
		(s_len, c_len) if s_len == c_len && s == check => c_len - 1,
		(s_len, c_len) if s_len > c_len && s.starts_with(check) && s.chars().nth(c_len).unwrap().is_whitespace() => c_len,
		(_, _) => 0,
	}
}

struct Parser<'a> {
	lex: Lexer<'a>,
}

impl Parser<'_> {
	fn new(lex: Lexer) -> Parser { Parser {lex} }

	fn parse(&mut self) -> Option<Expr> {
		self.parse_expr(0.0).ok()
	}

	fn parse_expr(&mut self, min_bp: f64) -> Result<Expr, ParseErr> {
		/*while let Ok(val) = self.lex.next() {debug!("token: {:?}", val)}
		match self.lex.next().err() {
			
			_ => return Err(ParseErr::Invalid),
		}*/
		let mut lhs: Expr = match self.lex.next() {
			Ok(val) => match val {
				Token::Atom(val) => Ok(Expr::Atom(val)),
				Token::Op(op) => match op {
					Op::LParen => {
						let val = self.parse_expr(0.0);
						if self.lex.next() != Ok(Token::Op(Op::RParen)) {
							debug!("Unclosed parens");
							Err(ParseErr::Invalid)
						}
						else {
							val
						}
					},
					// Op::Func(f) => Ok(Expr::Node(Op::Func(f), vec![self.parse_expr(op.get_lbp())?])),
					_ => match op.bp_prefix() {
						Some(bp) => Ok(Expr::Node(op, vec![self.parse_expr(bp)?])),
						None => {debug!("Got unexpected {:?} as prefix", op); Err(ParseErr::Invalid)}
					}
				},
			},
			Err(err) => Err(err),
		}.map_err(|err| { debug!("Unexpected error at start of expr: {:?}", err); err })?;
		debug!("lhs of expression is {:?}, min_bp is {}", lhs, min_bp);
		loop {
			debug!("loop start");
			let op: Op = match self.lex.peek() {
				Err(ParseErr::Eof) => break,
				Err(e) => { debug!("In expr got err {:?}", e); Err(*e) },
				Ok(tok) => match tok {
					Token::Op(op) => match op {
						Op::RParen => {
							debug!("got RParen");
							break;
						},
						_ => Ok(*op),
					}
					v => { debug!("Got unexpected token {:?}", v); Err(ParseErr::Invalid) },
				}
			}.map_err(|err| { debug!("Unexpected error inside expr at {:?}", err); err })?;
			debug!("op is {:?}", op);
			if let Some((lbp, rbp)) = op.bp_infix() {
				if (lbp < min_bp) { break; }
				self.lex.next();
				let rhs: Expr = self.parse_expr(rbp)?;
				lhs = Expr::Node(op, vec![lhs, rhs]);
			} else {
				debug!("Got unexpected non-infix operator in expression: {:?}", op);
				return Err(ParseErr::Invalid);
			}
		}
		debug!("Returning expr {:?}", lhs);
		Ok(lhs)
	}
}

#[derive(Debug)]
enum Expr {
	Evaluated,
	Atom(Atom),
	Node(Op, Vec<Expr>),
}

impl Expr {
	fn eval(&mut self) -> BigFloat {
		let res = match self {
			Expr::Atom(_) => {
				let v = mem::replace(self, Expr::Evaluated);
				if let Expr::Atom(at) = v {
					match at {
						Atom::Number(n) => n,
						Atom::Const(c) => c.get_val(),
					}
				} else {
					unreachable!();
				}
				// at.get_val()
			}
			Expr::Node(op, exprs) => {
				*self = Expr::Atom(Atom::Number(op.apply_to(exprs)));
				self.eval()
			}
			Expr::Evaluated => unreachable!("Tried to evaluate an already evaluated node"),
		};
		// debug!("{:?} evaluated to {}", self, res);
		res
	}
}