use std::collections::HashMap;

#[derive(Debug, PartialEq)]
enum Token {
    Number(f64),
    Operator(Operator),
    Function(String),
    Variable(String),
    RightParen,
    LeftParen,
    Comma
}

#[derive(Debug, PartialEq)]
enum Operator {
    Add, Sub, Mul, Div, Mod, Exp, Chs
}

#[derive(Debug, PartialEq)]
enum TokenizeError {
    TooManyDots(String),
    ConstantNotFound(String),
    BadCharacter(char),
    BadOperator(char)
}

#[derive(Debug, PartialEq)]
enum ParseError {
    MismatchedParens,
    MissingLeftParen,
    TrailingParen,
}

#[derive(Debug, PartialEq)]
enum CompileError {
    BadOperator(Operator),
    BadFunction(String),
    BadToken(Token)
}

#[derive(Debug, PartialEq)]
enum Word {
    Push(f64),
    PushVar(String),
    Add,
    Subtract,
    Multiply,
    Divide,
    Modulo,
    Power,
    Negate,
    Sin,
    Cos,
    Tan,
    Sqrt
}

#[derive(Debug, PartialEq)]
enum ExecuteError {
    DivideByZero,
    BadVariable(String),
    UnimplementedWord(Word),
    StackUnderflow
}

#[derive(Debug, PartialEq)]
enum ExpressionError {
    Tokenize(TokenizeError),
    Parse(ParseError),
    Compile(CompileError),
    Execute(ExecuteError)
}

impl Word {
    fn pops(&self) -> usize {
        match self {
            Word::Push(_)   => 0,
            Word::PushVar(_)=> 0,
            Word::Negate    => 1,
            Word::Sin       => 1,
            Word::Cos       => 1,
            Word::Tan       => 1,
            Word::Sqrt      => 1,
            Word::Add       => 2,
            Word::Subtract  => 2,
            Word::Multiply  => 2,
            Word::Divide    => 2,
            Word::Modulo    => 2,
            Word::Power     => 2,
        }
    }
}

impl Operator {
    fn precedence(&self) -> u8 {
	match self {
	    Operator::Add | Operator::Sub => 2,
	    Operator::Mul | Operator::Div | Operator::Mod => 3,
	    Operator::Exp => 4,
	    Operator::Chs => 5
	}
    }

    fn is_left_associative(&self) -> bool {
	match self {
	    Operator::Add | Operator::Sub | Operator::Mul | Operator::Div | Operator::Mod | Operator::Chs => true,
	    Operator::Exp => false
	}
    }

    fn is_right_associative(&self) -> bool {
	match self {
	    Operator::Add | Operator::Sub | Operator::Mul | Operator::Div | Operator::Mod | Operator::Chs => false,
	    Operator::Exp => true
	}
    }
}

impl TryFrom<char> for Operator {
    type Error = TokenizeError;
    
    fn try_from(c: char) -> Result<Self, Self::Error> {
	match c {
	    '+' => Ok(Operator::Add),
	    '-' => Ok(Operator::Sub),
	    '*' => Ok(Operator::Mul),
	    '/' => Ok(Operator::Div),
	    '^' => Ok(Operator::Exp),
	    '$' => Ok(Operator::Chs),
	    _ => Err(TokenizeError::BadOperator(c))
	}
    }
}

impl TryFrom<Token> for Word {
    type Error = CompileError;

    fn try_from(t: Token) -> Result<Self, Self::Error> {
	match t {
	    Token::Number(n) => Ok(Word::Push(n)),
	    Token::Variable(s) => Ok(Word::PushVar(s)),
	    Token::Operator(o) => match o {
		Operator::Add => Ok(Word::Add),
		Operator::Sub => Ok(Word::Subtract),
		Operator::Mul => Ok(Word::Multiply),
		Operator::Div => Ok(Word::Divide),
		Operator::Mod => Ok(Word::Modulo),
		Operator::Exp => Ok(Word::Power),
		Operator::Chs => Ok(Word::Negate)
		// _ => Err(CompileError::BadOperator(o))
	    },
	    Token::Function(s) => match s.as_str() {
		"sin" => Ok(Word::Sin),
		"cos" => Ok(Word::Cos),
		"tan" => Ok(Word::Tan),
		_ => Err(CompileError::BadFunction(s))
	    },
	    _ => Err(CompileError::BadToken(t))
	}
    }
}

impl From<TokenizeError> for ExpressionError {
    fn from(e: TokenizeError) -> Self {
        ExpressionError::Tokenize(e)
    }
}

impl From<ParseError> for ExpressionError {
    fn from(e: ParseError) -> Self {
        ExpressionError::Parse(e)
    }
}

impl From<CompileError> for ExpressionError {
    fn from(e: CompileError) -> Self {
        ExpressionError::Compile(e)
    }
}

impl From<ExecuteError> for ExpressionError {
    fn from(e: ExecuteError) -> Self {
        ExpressionError::Execute(e)
    }
}

struct Expression {
    words: Vec<Word>
}

impl Expression {
    pub fn try_new(expression: &str) -> Result<Expression, ExpressionError> {
	let tokens = tokenize(expression)?;
	let parsed = parse(tokens)?;
	let words = compile(parsed)?;
	Ok(Self { words })
    }

    pub fn eval(&self, vars: &[(&str, f64)]) -> Result<f64, ExpressionError> {
	let mut stack: Vec<f64> = Vec::new();
	
	for word in &self.words {
	    let mut a: f64 = 0.0;
	    let mut b: f64 = 0.0;
	    
	    match word.pops() {
		1 => {
		    a = stack.pop().ok_or(ExecuteError::StackUnderflow)?;
		}

		2 => {
		    b = stack.pop().ok_or(ExecuteError::StackUnderflow)?;
		    a = stack.pop().ok_or(ExecuteError::StackUnderflow)?;
		}

		_ => {}
	    }
	    
	    let y = match word {
		Word::Push(n) => *n,
		Word::PushVar(v) => vars.iter()
		    .find(|(k, _)| *k == v)
		    .map(|(_, a)| *a)
		    .ok_or(ExecuteError::BadVariable(v.to_string()))?,
		Word::Add => a + b,
		Word::Subtract => a - b,
		Word::Multiply => a * b,
		Word::Divide => a / b,
		Word::Modulo => a % b,
		Word::Power => a.powf(b),
		Word::Negate => -a,
		Word::Sin => a.sin(),
		Word::Cos => a.cos(),
		Word::Tan => a.tan(),
		Word::Sqrt => a.sqrt(),
		_ => 0.0
	    };

	    stack.push(y);
	}

	Ok(stack.pop().unwrap())
    }
}

fn str_to_const(s: &str) -> Option<f64> {
    match s {
	"PI" => Some(std::f64::consts::PI),
	"E" => Some(std::f64::consts::E),
	_ => None
    }
}

fn tokenize(input: &str) -> Result<Vec<Token>, TokenizeError> {
    let mut tokens = Vec::new();
    let mut chars = input.chars().peekable();

    while let Some(&ch) = chars.peek() {
	match ch {
	    '0'..='9' => {
		let mut num = String::new();
		let mut has_dot = false;
		while let Some(&c) = chars.peek() {
		    if c.is_ascii_digit() || (c == '.' && !has_dot) {
			num.push(c);
			chars.next();
			has_dot |= c == '.';
		    } else if c == '.' && has_dot {
			return Err(TokenizeError::TooManyDots(num));
		    } else {
			break;
		    }
		}
		tokens.push(Token::Number(num.parse().unwrap()));
	    }

	    '+' | '-' | '*' | '/' | '%' | '^' => {
		chars.next();
		if (ch == '+' || ch == '-') &&
		    (tokens.len() == 0 ||
		     (tokens.len() > 1 && matches!(tokens.last(), Some(Token::Operator(_))))) {
			if ch == '-' {
			    /* change sign */
			    tokens.push(Token::Operator(Operator::Chs));
			}
		    } else {
			tokens.push(Token::Operator(Operator::try_from(ch)?));
		    }
	    }

	    'a'..='z' | '_' => {
		let mut pt: Option<Token> = None;
		let mut s = String::new();
		while let Some(&c) = chars.peek() {
		    match c {
			'a'..='z' | 'A'..='Z' | '0'..='9' | '_' => {
			    s.push(c);
			    chars.next();
			}

			'(' => {
			    pt = Some(Token::Function(s.clone()));
			    break;
			}

			_ => {
			    break;
			}
		    }
		}
		tokens.push(match pt {
		    Some(token) => token,
		    None => Token::Variable(s.clone())
		});
	    }

	    'A'..='Z' => {
		let mut s = String::new();
		while let Some(&c @ ('A'..='Z' | '_')) = chars.peek() {
		    s.push(c);
		    chars.next();
		}
		let n = str_to_const(&s).ok_or(TokenizeError::ConstantNotFound(s))?;
		tokens.push(Token::Number(n));
	    }
	    
	    c if c.is_whitespace() => {
		while let Some(&c) = chars.peek() {
		    if c.is_whitespace() {
			chars.next();
		    } else {
			break;
		    }
		}
	    }

	    ',' => {
		chars.next();
		tokens.push(Token::Comma);
	    }

	    ')' => {
		chars.next();
		tokens.push(Token::RightParen);
	    }

	    '(' => {
		chars.next();
		tokens.push(Token::LeftParen);
	    }

	    _ => {
		return Err(TokenizeError::BadCharacter(ch));
	    }
	}
    }

    Ok(tokens)
}

fn parse(tokens: Vec<Token>) -> Result<Vec<Token>, ParseError> {
    let mut output: Vec<Token> = Vec::new();
    let mut operators: Vec<Token> = Vec::new();
    
    for token in tokens {
	match token {
	    Token::Number(_) | Token::Variable(_) => {
		output.push(token);
	    }

	    Token::Function(_) => {
		operators.push(token);
	    }
	    
	    Token::Operator(o1) => {
		while let Some(Token::Operator(o2)) = operators.last() &&
		    ((o2.precedence() > o1.precedence()) ||
		     ((o1.precedence() == o2.precedence()) && o1.is_left_associative())) {
			output.push(operators.pop().unwrap());
		}
		operators.push(Token::Operator(o1));
	    }
	    
	    Token::Comma => {
		while !matches!(operators.last(), Some(&Token::LeftParen)) {
		    output.push(operators.pop().unwrap());
		}
	    }

	    Token::LeftParen => {
		operators.push(token);
	    }

	    Token::RightParen => {
		while let Some(top) = operators.last() {
		    if matches!(top, Token::LeftParen) {
			break;
		    }
		    output.push(operators.pop().unwrap());
		}
		
		if !matches!(operators.pop(), Some(Token::LeftParen)) {
		    return Err(ParseError::MismatchedParens);
		}
	    }
	}
    }

    while let Some(t) = operators.pop() {
	if t == Token::LeftParen {
	    return Err(ParseError::TrailingParen);
	}
	output.push(t);
    }

    Ok(output)
}

fn compile(tokens: Vec<Token>) -> Result<Vec<Word>, CompileError> {
    tokens.into_iter().map(Word::try_from).collect()
}

fn main() {
    //let input = "-x + 4 * -(3 - sin((y+PI)) / 2.123)^2";
    /*println!("{}", input);
    let tokens = tokenize(input).unwrap();
    println!("{:#?}", tokens);
    let program = parse(tokens).unwrap();
    println!("{:#?}", program);
    let executable = compile(program).unwrap();
    println!("{:#?}", executable);*/

    let input = "x * sin(PI/2)";
    let expression = Expression::try_new(input).unwrap();
    let vars = [("x", 3.0)];
    println!("{:#?}", expression.words);
    println!("{}", expression.eval(&vars).unwrap());
}