/* -------------------------------------------------------------------------- *
 *                                   Lepton                                   *
 * -------------------------------------------------------------------------- *
 * This is part of the Lepton expression parser originating from              *
 * Simbios, the NIH National Center for Physics-Based Simulation of           *
 * Biological Structures at Stanford, funded under the NIH Roadmap for        *
 * Medical Research, grant U54 GM072970. See https://simtk.org.               *
 *                                                                            *
 * Portions copyright (c) 2009 Stanford University and the Authors.           *
 * Authors: Peter Eastman                                                     *
 * Contributors:                                                              *
 *                                                                            *
 * Permission is hereby granted, free of charge, to any person obtaining a    *
 * copy of this software and associated documentation files (the "Software"), *
 * to deal in the Software without restriction, including without limitation  *
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
 * and/or sell copies of the Software, and to permit persons to whom the      *
 * Software is furnished to do so, subject to the following conditions:       *
 *                                                                            *
 * The above copyright notice and this permission notice shall be included in *
 * all copies or substantial portions of the Software.                        *
 *                                                                            *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
 * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
 * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
 * -------------------------------------------------------------------------- */

#include "Parser.h"
#include "Exception.h"
#include "ExpressionTreeNode.h"
#include "Operation.h"
#include "ParsedExpression.h"
#include <iostream>

using namespace Lepton;
using namespace std;

static const string Digits = "0123456789";
static const string Operators = "+-*/^";
static const bool LeftAssociative[] = {true, true, true, true, false};
static const int Precedence[] = {0, 0, 1, 1, 3};
static const Operation::Id OperationId[] = {Operation::ADD, Operation::SUBTRACT, Operation::MULTIPLY, Operation::DIVIDE, Operation::POWER};

class Lepton::ParseToken {
public:
    enum Type {Number, Operator, Variable, Function, LeftParen, RightParen, Comma, Whitespace};

    ParseToken(string text, Type type) : text(text), type(type) {
    }
    const string& getText() const {
        return text;
    }
    Type getType() const {
        return type;
    }
private:
    string text;
    Type type;
};

ParseToken Parser::getNextToken(string expression, int start) {
    char c = expression[start];
    if (c == '(')
        return ParseToken("(", ParseToken::LeftParen);
    if (c == ')')
        return ParseToken(")", ParseToken::RightParen);
    if (c == ',')
        return ParseToken(",", ParseToken::Comma);
    if (Operators.find(c) != string::npos)
        return ParseToken(string(1, c), ParseToken::Operator);
    if (c == ' ') {
        // White space

        for (int pos = start+1; pos < expression.size(); pos++) {
            if (expression[pos] != ' ')
                return ParseToken(expression.substr(start, pos-start), ParseToken::Whitespace);
        }
        return ParseToken(expression.substr(start, string::npos), ParseToken::Whitespace);
    }
    if (c == '.' || Digits.find(c) != string::npos) {
        // A number

        bool foundDecimal = (c == '.');
        bool foundExp = false;
        int pos;
        for (pos = start+1; pos < expression.size(); pos++) {
            c = expression[pos];
            if (Digits.find(c) != string::npos)
                continue;
            if (c == '.' && !foundDecimal) {
                foundDecimal = true;
                continue;
            }
            if ((c == 'e' || c == 'E') && !foundExp) {
                foundExp = true;
                if (pos < expression.size()-1 && expression[pos+1] == '-')
                    pos++;
                continue;
            }
            break;
        }
        return ParseToken(expression.substr(start, pos-start), ParseToken::Number);
    }

    // A variable, function, or left parenthesis

    for (int pos = start; pos < expression.size(); pos++) {
        c = expression[pos];
        if (c == '(')
            return ParseToken(expression.substr(start, pos-start+1), ParseToken::Function);
        if (Operators.find(c) != string::npos || c == ',' || c == ')')
            return ParseToken(expression.substr(start, pos-start), ParseToken::Variable);
    }
    return ParseToken(expression.substr(start, string::npos), ParseToken::Variable);
}

vector<ParseToken> Parser::tokenize(string expression) {
    vector<ParseToken> tokens;
    int pos = 0;
    while (pos < expression.size()) {
        ParseToken token = getNextToken(expression, pos);
        if (token.getType() != ParseToken::Whitespace)
            tokens.push_back(token);
        pos += token.getText().size();
    }
    return tokens;
}

ParsedExpression Parser::parse(string expression) {
    vector<ParseToken> tokens = tokenize(expression);
    int pos = 0;
    ExpressionTreeNode result = parsePrecedence(tokens, pos, 0);
    if (pos != tokens.size())
        throw Exception("Parse error: unexpected text at end of expression");
    return ParsedExpression(result);
}

ExpressionTreeNode Parser::parsePrecedence(const vector<ParseToken>& tokens, int& pos, int precedence) {
    if (pos == tokens.size())
        throw Exception("Parse error: unexpected end of expression");

    // Parse the next value (number, variable, function, parenthesized expression)

    ParseToken token = tokens[pos];
    ExpressionTreeNode result;
    if (token.getType() == ParseToken::Number) {
        double value;
        stringstream(token.getText()) >> value;
        result = ExpressionTreeNode(new Operation::Constant(value));
        pos++;
    }
    else if (token.getType() == ParseToken::Variable) {
        Operation* op = new Operation::Variable(token.getText());
        result = ExpressionTreeNode(op);
        pos++;
    }
    else if (token.getType() == ParseToken::LeftParen) {
        pos++;
        result = parsePrecedence(tokens, pos, 0);
        if (pos == tokens.size() || tokens[pos].getType() != ParseToken::RightParen)
        throw Exception("Parse error: unbalanced parentheses");
        pos++;
    }
    else if (token.getType() == ParseToken::Function) {
        pos++;
        vector<ExpressionTreeNode> args;
        bool moreArgs;
        do {
            args.push_back(parsePrecedence(tokens, pos, 0));
            moreArgs = (pos < tokens.size() && tokens[pos].getType() == ParseToken::Comma);
            if (moreArgs)
                pos++;
        } while (moreArgs);
        if (pos == tokens.size() || tokens[pos].getType() != ParseToken::RightParen)
        throw Exception("Parse error: unbalanced parentheses");
        pos++;
        result = ExpressionTreeNode(getFunctionOperation(token.getText(), args.size()), args);
    }
    else if (token.getType() == ParseToken::Operator && token.getText() == "-") {
        pos++;
        ExpressionTreeNode toNegate = parsePrecedence(tokens, pos, 2);
        result = ExpressionTreeNode(new Operation::Negate(), toNegate);
    }
    else
        throw Exception("Parse error: unexpected token");

    // Now deal with the next binary operator.

    while (pos < tokens.size() && tokens[pos].getType() == ParseToken::Operator) {
        token = tokens[pos];
        int op = Operators.find(token.getText());
        int opPrecedence = Precedence[op];
        if (opPrecedence < precedence)
            return result;
        pos++;
        ExpressionTreeNode arg = parsePrecedence(tokens, pos, LeftAssociative[op] ? opPrecedence+1 : opPrecedence);
        result = ExpressionTreeNode(getOperatorOperation(token.getText()), result, arg);
    }
    return result;
}

Operation* Parser::getOperatorOperation(const std::string& name) {
    switch (OperationId[Operators.find(name)]) {
        case Operation::ADD:
            return new Operation::Add();
        case Operation::SUBTRACT:
            return new Operation::Subtract();
        case Operation::MULTIPLY:
            return new Operation::Multiply();
        case Operation::DIVIDE:
            return new Operation::Divide();
        case Operation::POWER:
            return new Operation::Power();
        default:
            throw Exception("Parse error: unknown operator");
    }
}

Operation* Parser::getFunctionOperation(const std::string& name, int arguments) {

    static map<string, Operation::Id> opMap;
    if (opMap.size() == 0) {
        opMap["sqrt"] = Operation::SQRT;
        opMap["exp"] = Operation::EXP;
        opMap["log"] = Operation::LOG;
        opMap["sin"] = Operation::SIN;
        opMap["cos"] = Operation::COS;
        opMap["sec"] = Operation::SEC;
        opMap["csc"] = Operation::CSC;
        opMap["tan"] = Operation::TAN;
        opMap["cot"] = Operation::COT;
        opMap["asin"] = Operation::ASIN;
        opMap["acos"] = Operation::ACOS;
        opMap["atan"] = Operation::ATAN;
        opMap["square"] = Operation::SQUARE;
        opMap["cube"] = Operation::CUBE;
        opMap["recip"] = Operation::RECIPROCAL;
        opMap["increment"] = Operation::INCREMENT;
        opMap["decrement"] = Operation::DECREMENT;
    }
    string trimmed = name.substr(0, name.size()-1);
    map<string, Operation::Id>::const_iterator iter = opMap.find(trimmed);
    if (iter == opMap.end())
        return new Operation::Custom(trimmed, arguments);
    switch (iter->second) {
        case Operation::SQRT:
            return new Operation::Sqrt();
        case Operation::EXP:
            return new Operation::Exp();
        case Operation::LOG:
            return new Operation::Log();
        case Operation::SIN:
            return new Operation::Sin();
        case Operation::COS:
            return new Operation::Cos();
        case Operation::SEC:
            return new Operation::Sec();
        case Operation::CSC:
            return new Operation::Csc();
        case Operation::TAN:
            return new Operation::Tan();
        case Operation::COT:
            return new Operation::Cot();
        case Operation::ASIN:
            return new Operation::Asin();
        case Operation::ACOS:
            return new Operation::Acos();
        case Operation::ATAN:
            return new Operation::Atan();
        case Operation::SQUARE:
            return new Operation::Square();
        case Operation::CUBE:
            return new Operation::Cube();
        case Operation::RECIPROCAL:
            return new Operation::Reciprocal();
        case Operation::INCREMENT:
            return new Operation::Increment();
        case Operation::DECREMENT:
            return new Operation::Decrement();
        default:
            throw Exception("Parse error: unknown function");
    }
}