expression.js

// Expression Evaluator - JavaScript port of expression.py

const TokenType = Object.freeze({
    VAR: 'VAR',
    ADD: 'ADD',
    SUB: 'SUB',
    MUL: 'MUL',
    DIV: 'DIV',
    NUMBER: 'NUMBER',
    PLUSMINUS: 'PLUSMINUS',
    PLUSMINUSPERCENT: 'PLUSMINUSPERCENT',
    LEFT_PAR: 'LEFT_PAR',
    RIGHT_PAR: 'RIGHT_PAR',
    ASSIGN: 'ASSIGN',
    DEFINE: 'DEFINE',
    END: 'END',
    LVAR: 'LVAR',
    PERCENT: 'PERCENT',
    NULL: 'NULL',
    ARG_SEP: 'ARG_SEP',
    FUNCTION: 'FUNCTION',
    NUMBERS: 'NUMBERS',
    LFUNCTION: 'LFUNCTION',
    ARG: 'ARG',
});

class Token {
    constructor(type, text = "", value = null) {
        this.type = type;
        this.text = text;
        this.value = value;
    }
    toString() {
        return `Token(${this.type}, "${this.text}", ${this.value})`;
    }
}

// Box-Muller transform for normal distribution (replaces numpy)
function randomNormal(mean, std, N) {
    const values = new Float64Array(N);
    for (let i = 0; i < N; i += 2) {
        const u1 = Math.random();
        const u2 = Math.random();
        const z0 = Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
        const z1 = Math.sqrt(-2 * Math.log(u1)) * Math.sin(2 * Math.PI * u2);
        values[i] = mean + std * z0;
        if (i + 1 < N) values[i + 1] = mean + std * z1;
    }
    return values;
}

// Array math helpers (replaces numpy array operations)
function resample(arr, size) {
    const result = new Float64Array(size);
    for (let i = 0; i < size; i++) result[i] = arr[Math.floor(Math.random() * arr.length)];
    return result;
}

function arrayOp(a, b, op) {
    const aIsArray = a instanceof Float64Array || Array.isArray(a);
    const bIsArray = b instanceof Float64Array || Array.isArray(b);

    if (aIsArray && bIsArray) {
        if (a.length !== b.length) {
            if (a.length < b.length) a = resample(a, b.length);
            else b = resample(b, a.length);
        }
        const len = a.length;
        const result = new Float64Array(len);
        for (let i = 0; i < len; i++) result[i] = op(a[i], b[i]);
        return result;
    } else if (aIsArray) {
        const result = new Float64Array(a.length);
        for (let i = 0; i < a.length; i++) result[i] = op(a[i], b);
        return result;
    } else if (bIsArray) {
        const result = new Float64Array(b.length);
        for (let i = 0; i < b.length; i++) result[i] = op(a, b[i]);
        return result;
    } else {
        return op(a, b);
    }
}

function arrayFunc(a, fn) {
    if (a instanceof Float64Array || Array.isArray(a)) {
        const result = new Float64Array(a.length);
        for (let i = 0; i < a.length; i++) result[i] = fn(a[i]);
        return result;
    }
    return fn(a);
}

function arrayMean(a) {
    if (a instanceof Float64Array || Array.isArray(a)) {
        let sum = 0;
        for (let i = 0; i < a.length; i++) sum += a[i];
        return sum / a.length;
    }
    return a;
}

function arrayStd(a) {
    if (a instanceof Float64Array || Array.isArray(a)) {
        if (a.length < 2) return 0;
        const mean = arrayMean(a);
        let sumSq = 0;
        for (let i = 0; i < a.length; i++) sumSq += (a[i] - mean) ** 2;
        const denom = a.length < 30 ? a.length - 1 : a.length;
        return Math.sqrt(sumSq / denom);
    }
    return 0;
}

function isArray(v) {
    return v instanceof Float64Array || Array.isArray(v);
}

class Expression {
    // Token regexes - ORDER MATTERS (same as Python)
    // VAR must not be followed by word char + '(' (that would be a FUNCTION)
    static tokenPatterns = [
        [TokenType.VAR,                /^([πa-zA-Z][\w]*)(?![\w(].*)/],
        [TokenType.FUNCTION,           /^([\w][\w\d]*)(\(.*)/],
        [TokenType.ADD,                /^(\+)/],
        [TokenType.SUB,                /^(-)/],
        [TokenType.MUL,                /^(\*)/],
        [TokenType.DIV,                /^([/÷])/],
        [TokenType.PLUSMINUS,          /^(±)(?!%)/],
        [TokenType.PLUSMINUSPERCENT,   /^(±%)/],
        [TokenType.PERCENT,            /^(%)/],
        [TokenType.NUMBERS,            /^(\[[^\]]+\])/],
        [TokenType.NUMBER,             /^(-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)/],
        [TokenType.LEFT_PAR,           /^(\()/],
        [TokenType.RIGHT_PAR,          /^(\))/],
        [TokenType.DEFINE,             /^(:=)/],
        [TokenType.ASSIGN,             /^(=)/],
        [TokenType.ARG_SEP,            /^(,)/],
    ];

    static precedence = {
        [TokenType.VAR]: 0,
        [TokenType.LVAR]: 0,
        [TokenType.NUMBER]: 0,
        [TokenType.NUMBERS]: 0,
        [TokenType.ADD]: 5,
        [TokenType.SUB]: 5,
        [TokenType.MUL]: 6,
        [TokenType.DIV]: 6,
        [TokenType.FUNCTION]: 7,
        [TokenType.LFUNCTION]: 7,
        [TokenType.PLUSMINUS]: 5,
        [TokenType.PLUSMINUSPERCENT]: 5,
        [TokenType.PERCENT]: 4,
        [TokenType.ASSIGN]: 3,
        [TokenType.DEFINE]: 3,
        [TokenType.LEFT_PAR]: -1,
        [TokenType.RIGHT_PAR]: -1,
        [TokenType.ARG_SEP]: -2,
        [TokenType.END]: -2,
        [TokenType.NULL]: -2,
    };

    static operators = new Set([
        TokenType.ADD, TokenType.SUB, TokenType.MUL, TokenType.DIV,
        TokenType.ASSIGN, TokenType.DEFINE, TokenType.PLUSMINUS, TokenType.PLUSMINUSPERCENT,
        TokenType.PERCENT, TokenType.LEFT_PAR, TokenType.RIGHT_PAR,
        TokenType.ARG_SEP, TokenType.END, TokenType.FUNCTION, TokenType.LFUNCTION,
    ]);

    static functions = {
        cos: Math.cos,
        sin: Math.sin,
        tan: Math.tan,
        acos: Math.acos,
        asin: Math.asin,
        atan: Math.atan,
        sqrt: Math.sqrt,
        exp: Math.exp,
        ln: Math.log,
        log: Math.log,
        log10: Math.log10,
        log2: Math.log2,
        abs: Math.abs,
        round: Math.round,
    };

    static defaultGlobalVariables = { pi: Math.PI, "π": Math.PI, e: Math.E, format: 2 };
    static globalVariables = { ...Expression.defaultGlobalVariables };
    static evalStack = [];
    static defaultN = 500_000;

    constructor(expression) {
        this.expression = expression;
        this._showHist = false;
        this.expressionTokens = Expression.tokenize(expression);
        this.rpnList = Expression.buildRpnList(this.expressionTokens);

        const lvalue = this.lvalue;
        if (lvalue !== null) {
            if (lvalue.type === TokenType.LFUNCTION) {
                Expression.functions[lvalue.text] = this;
            } else if (lvalue.type === TokenType.LVAR) {
                if (this.isDefinition) {
                    // := stores Expression (lazy, re-evaluated each time, independent)
                    Expression.globalVariables[lvalue.text] = this;
                } else {
                    // = evaluates immediately, stores result (correlated on reuse)
                    const result = this.evaluate(Expression.defaultN);
                    Expression.globalVariables[lvalue.text] = result;
                }
            }
        }
    }

    static resetGlobalVariables() {
        Expression.globalVariables = { ...Expression.defaultGlobalVariables };
        Expression.evalStack = [];
    }

    static tokenize(workString) {
        const tokens = [];
        while (workString.length > 0) {
            const [token, remaining] = Expression.nextToken(workString);
            workString = remaining;

            if ((token.type === TokenType.ASSIGN || token.type === TokenType.DEFINE) && tokens.length > 0) {
                const tokenTypes = tokens.map(t => t.type);
                if (tokenTypes.includes(TokenType.FUNCTION)) {
                    for (const t of tokens) {
                        if (t.type === TokenType.FUNCTION) t.type = TokenType.LFUNCTION;
                        else if (t.type === TokenType.VAR) t.type = TokenType.ARG;
                    }
                } else if (tokenTypes.includes(TokenType.VAR)) {
                    for (const t of tokens) {
                        if (t.type === TokenType.VAR) t.type = TokenType.LVAR;
                    }
                }
            }

            tokens.push(token);
        }
        tokens.push(new Token(TokenType.END, ''));
        return tokens;
    }

    static nextToken(workString) {
        // Strip leading whitespace
        workString = workString.replace(/^\s+/, '');

        for (const [tokenType, regex] of Expression.tokenPatterns) {
            const match = workString.match(regex);
            if (match !== null) {
                const tokenText = match[1];
                let remaining;

                if (tokenType === TokenType.FUNCTION) {
                    // FUNCTION regex captures the rest including '(' in group 2
                    remaining = match[2];
                } else {
                    remaining = workString.slice(match[0].length);
                }

                if (tokenType === TokenType.NUMBER) {
                    return [new Token(tokenType, tokenText, parseFloat(tokenText)), remaining];
                } else if (tokenType === TokenType.NUMBERS) {
                    const content = tokenText.slice(1, -1); // strip [ and ]
                    const values = new Float64Array(content.split(',').map(s => parseFloat(s.trim())));
                    return [new Token(tokenType, tokenText, values), remaining];
                } else {
                    return [new Token(tokenType, tokenText), remaining];
                }
            }
        }
        throw new Error(`Unrecognized token: ${workString}`);
    }

    get isAssignment() {
        if (this.rpnList.length === 0) return false;
        const lastType = this.rpnList[this.rpnList.length - 1].type;
        return lastType === TokenType.ASSIGN || lastType === TokenType.DEFINE;
    }

    get isDefinition() {
        return this.rpnList.length > 0 && this.rpnList[this.rpnList.length - 1].type === TokenType.DEFINE;
    }

    get lvalue() {
        if (this.isAssignment) return this.expressionTokens[0];
        return null;
    }

    get hasDistributions() {
        const [mean, std] = this.evaluateUncertainty(10);
        return std > 0;
    }

    static buildRpnList(tokens) {
        const rpnList = [];
        const operatorStack = [];

        for (let token of tokens) {
            if (Expression.operators.has(token.type)) {
                while (operatorStack.length > 0) {
                    if (token.type === TokenType.LEFT_PAR) break;

                    const topOperatorToken = operatorStack.pop();

                    if (token.type === TokenType.RIGHT_PAR && topOperatorToken.type === TokenType.LEFT_PAR) {
                        break;
                    }

                    if (token.type === TokenType.ARG_SEP && topOperatorToken.type === TokenType.LEFT_PAR) {
                        operatorStack.push(topOperatorToken);
                        break;
                    }

                    if (topOperatorToken.type === TokenType.PLUSMINUS && token.type === TokenType.PERCENT) {
                        operatorStack.push(new Token(TokenType.PLUSMINUSPERCENT, "±%"));
                        token = new Token(TokenType.NULL, "");
                        break;
                    }

                    if (Expression.precedence[token.type] <= Expression.precedence[topOperatorToken.type]) {
                        rpnList.push(topOperatorToken);
                    } else {
                        operatorStack.push(topOperatorToken);
                        break;
                    }
                }

                if (token.type !== TokenType.RIGHT_PAR && token.type !== TokenType.NULL && token.type !== TokenType.ARG_SEP) {
                    operatorStack.push(token);
                }
            } else if ([TokenType.VAR, TokenType.NUMBER, TokenType.NUMBERS, TokenType.ARG, TokenType.LVAR].includes(token.type)) {
                rpnList.push(token);
            } else {
                throw new Error(`Unable to treat token: ${token}`);
            }
        }

        return rpnList;
    }

    evaluateUncertainty(N = null, variables = null) {
        if (N === null) N = Expression.defaultN;
        const mean = this.evaluateScalar(variables);
        const values = this.evaluate(N, variables);
        const std = arrayStd(values);
        return [mean, std];
    }

    evaluateScalar(variables = null) {
        return this.evaluate(1, variables, false, true);
    }

    evaluate(N = 1_000_000, variables = null, _nested = false, _scalar = false) {
        const localVariables = { ...Expression.globalVariables };
        if (variables !== null) Object.assign(localVariables, variables);

        // Clear stack at top-level calls to prevent corruption from prior errors
        if (!_nested) Expression.evalStack.length = 0;
        const evalStack = Expression.evalStack;

        for (const token of this.rpnList) {
            if (Expression.operators.has(token.type)) {
                switch (token.type) {
                    case TokenType.ADD: {
                        const second = evalStack.pop();
                        const first = evalStack.pop();
                        evalStack.push(arrayOp(first, second, (a, b) => a + b));
                        break;
                    }
                    case TokenType.SUB: {
                        const second = evalStack.pop();
                        const first = evalStack.pop();
                        evalStack.push(arrayOp(first, second, (a, b) => a - b));
                        break;
                    }
                    case TokenType.DIV: {
                        const second = evalStack.pop();
                        const first = evalStack.pop();
                        evalStack.push(arrayOp(first, second, (a, b) => a / b));
                        break;
                    }
                    case TokenType.MUL: {
                        const second = evalStack.pop();
                        const first = evalStack.pop();
                        evalStack.push(arrayOp(first, second, (a, b) => a * b));
                        break;
                    }
                    case TokenType.FUNCTION: {
                        if (token.text === 'hist') {
                            const first = evalStack.pop();
                            this._showHist = true;
                            evalStack.push(first);
                        } else if (token.text === 'random') {
                            const second = evalStack.pop();
                            const first = evalStack.pop();
                            if (_scalar) {
                                evalStack.push(isArray(first) ? arrayMean(first) : first);
                            } else if (isArray(first) || isArray(second)) {
                                const len = isArray(first) ? first.length : (isArray(second) ? second.length : N);
                                const result = new Float64Array(len);
                                for (let i = 0; i < len; i++) {
                                    const m = isArray(first) ? first[i] : first;
                                    const s = isArray(second) ? second[i] : second;
                                    const u1 = Math.random(), u2 = Math.random();
                                    result[i] = m + s * Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
                                }
                                evalStack.push(result);
                            } else {
                                evalStack.push(randomNormal(first, second, N));
                            }
                        } else {
                            const fctOrExpr = Expression.functions[token.text];
                            if (fctOrExpr instanceof Expression) {
                                evalStack.push(fctOrExpr.evaluate(N, null, true, _scalar));
                            } else {
                                const first = evalStack.pop();
                                evalStack.push(arrayFunc(first, fctOrExpr));
                            }
                        }
                        break;
                    }
                    case TokenType.LFUNCTION: {
                        evalStack.push(token);
                        break;
                    }
                    case TokenType.PLUSMINUS: {
                        const second = evalStack.pop();
                        const first = evalStack.pop();
                        if (_scalar) {
                            // Scalar mode: just return the mean value
                            evalStack.push(isArray(first) ? arrayMean(first) : first);
                        } else if (isArray(first) || isArray(second)) {
                            const len = isArray(first) ? first.length : (isArray(second) ? second.length : N);
                            const result = new Float64Array(len);
                            for (let i = 0; i < len; i++) {
                                const m = isArray(first) ? first[i] : first;
                                const s = isArray(second) ? second[i] : second;
                                const u1 = Math.random(), u2 = Math.random();
                                result[i] = m + s * Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
                            }
                            evalStack.push(result);
                        } else {
                            evalStack.push(randomNormal(first, second, N));
                        }
                        break;
                    }
                    case TokenType.PLUSMINUSPERCENT: {
                        const second = evalStack.pop();
                        const first = evalStack.pop();
                        if (_scalar) {
                            evalStack.push(isArray(first) ? arrayMean(first) : first);
                        } else if (isArray(first) || isArray(second)) {
                            const len = isArray(first) ? first.length : (isArray(second) ? second.length : N);
                            const result = new Float64Array(len);
                            for (let i = 0; i < len; i++) {
                                const m = isArray(first) ? first[i] : first;
                                const s = isArray(second) ? second[i] : second;
                                const u1 = Math.random(), u2 = Math.random();
                                result[i] = m + (m * s / 100) * Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
                            }
                            evalStack.push(result);
                        } else {
                            evalStack.push(randomNormal(first, first * (second / 100), N));
                        }
                        break;
                    }
                    case TokenType.ASSIGN:
                    case TokenType.DEFINE: {
                        const second = evalStack.pop();
                        const first = evalStack.pop();
                        if (!(first instanceof Token)) {
                            throw new Error(`Invalid lvalue for assignment: ${first}`);
                        }
                        if (first.type === TokenType.LVAR || first.type === TokenType.LFUNCTION) {
                            evalStack.push(second);
                        } else {
                            throw new Error(`Invalid lvalue for assignment: ${first}`);
                        }
                        break;
                    }
                }
            } else {
                switch (token.type) {
                    case TokenType.VAR: {
                        if (token.text in localVariables) {
                            const exprOrValue = localVariables[token.text];
                            if (exprOrValue instanceof Expression) {
                                evalStack.push(exprOrValue.evaluate(N, localVariables, true, _scalar));
                            } else if (_scalar && isArray(exprOrValue)) {
                                evalStack.push(arrayMean(exprOrValue));
                            } else {
                                evalStack.push(exprOrValue);
                            }
                        } else {
                            throw new Error(`Variable '${token.text}' is undefined`);
                        }
                        break;
                    }
                    case TokenType.LVAR: {
                        evalStack.push(token);
                        break;
                    }
                    case TokenType.NUMBER: {
                        evalStack.push(token.value);
                        break;
                    }
                    case TokenType.NUMBERS: {
                        if (_scalar) {
                            evalStack.push(arrayMean(token.value));
                        } else {
                            evalStack.push(token.value);
                        }
                        break;
                    }
                    case TokenType.ARG: {
                        const arg = evalStack.pop();
                        localVariables[token.text] = arg;
                        break;
                    }
                    default:
                        throw new Error(`Unknown token when evaluating: ${token}`);
                }
            }
        }

        return evalStack.pop();
    }
}

// Export for Node.js testing, but also work in browser
if (typeof module !== 'undefined' && module.exports) {
    module.exports = { Expression, Token, TokenType, arrayMean, arrayStd, isArray };
}