"use strict";

class ParseError extends Error {
  constructor(message) {
    super(message);
    this.name = "ParseError";
  }
}

const OPEN_TO_CLOSE = {
  "(": ")",
  "[": "]",
  "{": "}",
  "<": ">",
  "\u00ab": "\u00bb",
};
const OPEN_BRACKETS = new Set(Object.keys(OPEN_TO_CLOSE));
const CLOSE_BRACKETS = new Set(Object.values(OPEN_TO_CLOSE));
const BRACKET_RE = /\u00ab|\u00bb|[()[\]{}<>]|[^\s()[\]{}<>\u00ab\u00bb]+/g;

function isConst(e) {
  return e instanceof Const;
}
function isZero(e) {
  return isConst(e) && e._value === 0;
}
function isOne(e) {
  return isConst(e) && e._value === 1;
}

class Expression {
  diff(v) {
    throw new Error("diff not implemented");
  }
  simplify() {
    return this;
  }
  isConst() {
    return false;
  }
}

class Const extends Expression {
  constructor(value) {
    super();
    this._value = value;
  }
  evaluate(_x, _y, _z) {
    return this._value;
  }
  toString() {
    return String(this._value);
  }
  prefix() {
    return String(this._value);
  }
  postfix() {
    return String(this._value);
  }
  diff(_v) {
    return ZERO;
  }
  isConst() {
    return true;
  }
  contains(_v) {
    return false;
  }
}

const ZERO = new Const(0);
const ONE = new Const(1);
const TWO = new Const(2);

const VAR_INDEX = { x: 0, y: 1, z: 2 };

class Variable extends Expression {
  constructor(name) {
    super();
    this._name = name;
  }
  evaluate(x, y, z) {
    return [x, y, z][VAR_INDEX[this._name]];
  }
  toString() {
    return this._name;
  }
  prefix() {
    return this._name;
  }
  postfix() {
    return this._name;
  }
  diff(v) {
    return v === this._name ? ONE : ZERO;
  }
  contains(v) {
    return v === this._name;
  }
}

class Operation extends Expression {
  constructor(...args) {
    super();
    this._args = args;
  }
  evaluate(x, y, z) {
    return this._calc(...this._args.map((a) => a.evaluate(x, y, z)));
  }
  toString() {
    return (
      this._args.map((a) => a.toString()).join(" ") +
      " " +
      this.constructor.symbol
    );
  }
  postfix() {
    return (
      "(" +
      this._args.map((a) => a.postfix()).join(" ") +
      " " +
      this.constructor.symbol +
      ")"
    );
  }
  prefix() {
    return (
      "(" +
      this.constructor.symbol +
      " " +
      this._args.map((a) => a.prefix()).join(" ") +
      ")"
    );
  }
  isConst() {
    return this._args.every((a) => a.isConst());
  }
  simplify() {
    const args = this._args.map((a) => a.simplify());
    const ruled = this._simplifyArgs(args);
    if (!(ruled instanceof Operation)) return ruled;
    if (ruled._args.every((a) => a.isConst())) {
      const val = ruled._calc(...ruled._args.map((a) => a._value));
      return new Const(val);
    }
    return ruled;
  }
  contains(v) {
    return this._args.some((a) => a.contains(v));
  }
  diff(v) {
    if (this.isConst()) return ZERO;
    if (!this.contains(v)) return ZERO;
    return this._diff(v);
  }
  _diff(v) {
    throw new Error("_diff not implemented");
  }
  _simplifyArgs(args) {
    return new this.constructor(...args);
  }
}

function makeOperation(symbol, calc, diffFn, simplifyArgsFn) {
  class Op extends Operation {
    _calc(...vals) {
      return calc(...vals);
    }
    _diff(v) {
      return diffFn(v, ...this._args);
    }
    _simplifyArgs(args) {
      if (simplifyArgsFn) {
        const result = simplifyArgsFn(...args);
        if (result !== undefined) return result;
      }
      return new Op(...args);
    }
  }
  Op.symbol = symbol;
  return Op;
}

function exprEqual(a, b) {
  if (a instanceof Const && b instanceof Const) return a._value === b._value;
  if (a instanceof Variable && b instanceof Variable)
    return a._name === b._name;
  if (a.constructor !== b.constructor) return false;
  if (!(a instanceof Operation)) return false;
  if (a._args.length !== b._args.length) return false;
  return a._args.every((arg, i) => exprEqual(arg, b._args[i]));
}

function factors(e) {
  if (e.constructor === Multiply) {
    return [...factors(e._args[0]), ...factors(e._args[1])];
  }
  return [e];
}

function cancelDivide(f, g) {
  const numFactors = factors(f);
  const denFactors = factors(g);

  for (let i = 0; i < numFactors.length; i++) {
    for (let j = 0; j < denFactors.length; j++) {
      if (exprEqual(numFactors[i], denFactors[j])) {
        const newNum = rebuildProduct(numFactors.filter((_, k) => k !== i));
        const newDen = rebuildProduct(denFactors.filter((_, k) => k !== j));
        if (isOne(newDen)) return newNum;
        const further = cancelDivide(newNum, newDen);
        if (further !== undefined) return further;
        return new Divide(newNum, newDen);
      }
    }
  }
  return undefined;
}

function rebuildProduct(fs) {
  if (fs.length === 0) return ONE;
  return fs.reduce((acc, f) => new Multiply(acc, f));
}

const Add = makeOperation(
  "+",
  (a, b) => a + b,
  (v, f, g) => new Add(f.diff(v), g.diff(v)),
  (f, g) => (isZero(f) ? g : isZero(g) ? f : undefined),
);

const Subtract = makeOperation(
  "-",
  (a, b) => a - b,
  (v, f, g) => new Subtract(f.diff(v), g.diff(v)),
  (f, g) => (isZero(g) ? f : isZero(f) ? new Negate(g) : undefined),
);

const Multiply = makeOperation(
  "*",
  (a, b) => a * b,
  (v, f, g) => new Add(new Multiply(f.diff(v), g), new Multiply(f, g.diff(v))),
  (f, g) => {
    if (isZero(f) || isZero(g)) return ZERO;
    if (isOne(f)) return g;
    if (isOne(g)) return f;
    if (g.constructor === Divide && isOne(g._args[0]))
      return new Divide(f, g._args[1]);
    if (f.constructor === Divide && isOne(f._args[0]))
      return new Divide(g, f._args[1]);
    return undefined;
  },
);

const Divide = makeOperation(
  "/",
  (a, b) => a / b,
  (v, f, g) =>
    new Divide(
      new Subtract(new Multiply(f.diff(v), g), new Multiply(f, g.diff(v))),
      new Multiply(g, g),
    ),
  (f, g) => {
    if (isZero(f)) return ZERO;
    if (isOne(g)) return f;
    if (exprEqual(f, g)) return ONE;
    if (f.constructor === Divide) {
      return new Divide(f._args[0], new Multiply(f._args[1], g));
    }
    if (g.constructor === Divide) {
      return new Divide(new Multiply(f, g._args[1]), g._args[0]);
    }
    const cancelled = cancelDivide(f, g);
    if (cancelled) return cancelled;
    return undefined;
  },
);

const Negate = makeOperation(
  "negate",
  (a) => -a,
  (v, f) => new Negate(f.diff(v)),
  (f) => (isZero(f) ? ZERO : undefined),
);

const Power = makeOperation(
  "pow",
  (a, b) => Math.pow(a, b),
  (v, f, g) =>
    new Add(
      new Multiply(
        new Multiply(g, new Power(f, new Subtract(g, ONE))),
        f.diff(v),
      ),
      new Multiply(
        new Multiply(new Power(f, g), new Log(new Const(Math.E), f)),
        g.diff(v),
      ),
    ),
);
Power.symbol = "pow";

const Log = makeOperation(
  "log",
  (a, b) => Math.log(Math.abs(b)) / Math.log(Math.abs(a)),
  (v, base, x) =>
    new Divide(
      new Subtract(
        new Divide(x.diff(v), x),
        new Multiply(new Log(base, x), new Divide(base.diff(v), base)),
      ),
      new Log(new Const(Math.E), base),
    ),
);
Log.symbol = "log";

function makeSumN(n) {
  class SumNOp extends Operation {
    _calc(...vals) {
      return vals.reduce((s, v) => s + v, 0);
    }
    _diff(v) {
      return this._args
        .map((a) => a.diff(v))
        .reduce((acc, d) => new Add(acc, d));
    }
  }
  SumNOp.symbol = "sum" + n;
  return SumNOp;
}

const Sum1 = makeSumN(1);
const Sum2 = makeSumN(2);
const Sum3 = makeSumN(3);
const Sum4 = makeSumN(4);
const Sum5 = makeSumN(5);

function makeAvgN(n) {
  class AvgNOp extends Operation {
    _calc(...vals) {
      return vals.reduce((s, v) => s + v, 0) / n;
    }
    _diff(v) {
      const sumDiff = this._args
        .map((a) => a.diff(v))
        .reduce((acc, d) => new Add(acc, d));
      return new Divide(sumDiff, new Const(n));
    }
  }
  AvgNOp.symbol = "avg" + n;
  return AvgNOp;
}

const Avg1 = makeAvgN(1);
const Avg2 = makeAvgN(2);
const Avg3 = makeAvgN(3);
const Avg4 = makeAvgN(4);
const Avg5 = makeAvgN(5);

const Gauss = makeOperation(
  "gauss",
  (a, b, c, x) => a * Math.exp(-((x - b) ** 2) / (2 * c * c)),
  (v, a, b, c, x) => {
    if (a.isConst() && isZero(a)) return ZERO;
    const gauss = new Gauss(a, b, c, x);
    const xmb = new Subtract(x, b);
    const c2 = new Multiply(c, c);
    const da = new Divide(a.diff(v), a);
    const dxb = new Divide(
      new Multiply(xmb, new Subtract(x.diff(v), b.diff(v))),
      c2,
    );
    const dc = new Divide(
      new Multiply(new Multiply(xmb, xmb), c.diff(v)),
      new Multiply(c2, c),
    );
    return new Multiply(gauss, new Add(da, new Subtract(dc, dxb)));
  },
  (a, b, c, x) => (isZero(a) ? ZERO : undefined),
);
Gauss.symbol = "gauss";

function wrapCalc(x, mn, mx) {
  const range = mx - mn;
  if (range === 0) return mn;
  return mn + ((((x - mn) % range) + range) % range);
}

const Floor = makeOperation(
  "__floor__",
  (a) => Math.floor(a),
  () => ZERO,
);
Floor.symbol = "__floor__";

const Wrap = makeOperation("wrap", wrapCalc, (v, x, mn, mx) => {
  const q = new Floor(new Divide(new Subtract(x, mn), new Subtract(mx, mn)));
  return new Subtract(
    x.diff(v),
    new Multiply(q, new Subtract(mx.diff(v), mn.diff(v))),
  );
});
Wrap.symbol = "wrap";

const Sin = makeOperation(
  "__sin__",
  (a) => Math.sin(a),
  (v, f) => new Multiply(new Cos(f), f.diff(v)),
);
Sin.symbol = "__sin__";

const Cos = makeOperation(
  "__cos__",
  (a) => Math.cos(a),
  (v, f) => new Negate(new Multiply(new Sin(f), f.diff(v))),
);
Cos.symbol = "__cos__";

const Asin = makeOperation(
  "__asin__",
  (a) => Math.asin(a),
  (v, f) =>
    new Divide(
      f.diff(v),
      new Power(new Subtract(ONE, new Multiply(f, f)), new Const(0.5)),
    ),
);
Asin.symbol = "__asin__";

const Tanh = makeOperation(
  "__tanh__",
  (a) => Math.tanh(a),
  (v, f) => {
    const th = new Tanh(f);
    return new Multiply(new Subtract(ONE, new Multiply(th, th)), f.diff(v));
  },
);
Tanh.symbol = "__tanh__";

function softWrapCalc(x, mn, mx, lam) {
  const range = mx - mn;
  if (range === 0) return mn;
  const a = (Math.PI * (x - mn)) / range;
  return (
    (Math.asin(Math.cos(a) * Math.tanh(lam * -Math.sin(a))) * range) / Math.PI +
    (mx + mn) / 2
  );
}

class SoftWrap extends Operation {
  _calc(...vals) {
    return softWrapCalc(...vals);
  }

  _tree() {
    const [x, mn, mx, lam] = this._args;
    const PI = new Const(Math.PI);
    const range = new Subtract(mx, mn);
    const a = new Divide(new Multiply(PI, new Subtract(x, mn)), range);
    const th = new Tanh(new Multiply(lam, new Negate(new Sin(a))));
    const inner = new Asin(new Multiply(new Cos(a), th));
    return new Add(
      new Multiply(inner, new Divide(range, PI)),
      new Divide(new Add(mx, mn), TWO),
    );
  }

  _diff(v) {
    return this._tree().diff(v);
  }
}
SoftWrap.symbol = "softWrap";

class ArcTan12 extends Operation {
  _calc(...vals) {
    return vals.length === 1
      ? Math.atan(vals[0])
      : Math.atan2(vals[0], vals[1]);
  }
  _diff(v) {
    if (this._args.length === 1) {
      const [x] = this._args;
      return new Divide(x.diff(v), new Add(ONE, new Multiply(x, x)));
    }
    const [y, x] = this._args;
    return new Divide(
      new Subtract(new Multiply(x, y.diff(v)), new Multiply(y, x.diff(v))),
      new Add(new Multiply(x, x), new Multiply(y, y)),
    );
  }
}
ArcTan12.symbol = "atan12";

const Exp = makeOperation(
  "exp",
  (a) => Math.exp(a),
  (v, f) => new Multiply(new Exp(f), f.diff(v)),
);
Exp.symbol = "exp";

class MeanExp extends Operation {
  _calc(...vals) {
    return vals.reduce((s, v) => s + Math.exp(v), 0) / vals.length;
  }
  _diff(v) {
    const n = this._args.length;
    const sumDiff = this._args
      .map((a) => new Multiply(new Exp(a), a.diff(v)))
      .reduce((acc, d) => new Add(acc, d));
    return new Divide(sumDiff, new Const(n));
  }
}
MeanExp.symbol = "meanExp";

class SumExp extends Operation {
  _calc(...vals) {
    return vals.reduce((s, v) => s + Math.exp(v), 0);
  }
  _diff(v) {
    return this._args
      .map((a) => new Multiply(new Exp(a), a.diff(v)))
      .reduce((acc, d) => new Add(acc, d));
  }
}
SumExp.symbol = "sumExp";

class Lse extends Operation {
  _calc(...vals) {
    return Math.log(vals.reduce((s, v) => s + Math.exp(v), 0));
  }
  _diff(v) {
    const sumExpNode = new SumExp(...this._args);
    const sumDiff = this._args
      .map((a) => new Multiply(new Exp(a), a.diff(v)))
      .reduce((acc, d) => new Add(acc, d));
    return new Divide(sumDiff, sumExpNode);
  }
}
Lse.symbol = "lse";

class Lme extends Operation {
  _calc(...vals) {
    return Math.log(vals.reduce((s, v) => s + Math.exp(v), 0) / vals.length);
  }
  _diff(v) {
    const n = this._args.length;
    const sumExpNode = new SumExp(...this._args);
    const sumDiff = this._args
      .map((a) => new Multiply(new Exp(a), a.diff(v)))
      .reduce((acc, d) => new Add(acc, d));
    return new Divide(sumDiff, sumExpNode);
  }
}
Lme.symbol = "lme";

const OPERATIONS = {
  "+": [Add, 2],
  "-": [Subtract, 2],
  "*": [Multiply, 2],
  "/": [Divide, 2],
  negate: [Negate, 1],
  pow: [Power, 2],
  log: [Log, 2],
  sum1: [Sum1, 1],
  sum2: [Sum2, 2],
  sum3: [Sum3, 3],
  sum4: [Sum4, 4],
  sum5: [Sum5, 5],
  avg1: [Avg1, 1],
  avg2: [Avg2, 2],
  avg3: [Avg3, 3],
  avg4: [Avg4, 4],
  avg5: [Avg5, 5],
  gauss: [Gauss, 4],
  wrap: [Wrap, 3],
  softWrap: [SoftWrap, 4],
  __floor__: [Floor, 1],
  __sin__: [Sin, 1],
  __cos__: [Cos, 1],
  __asin__: [Asin, 1],
  __tanh__: [Tanh, 1],
  exp: [Exp, 1],
  atan12: [ArcTan12, null],
  meanExp: [MeanExp, null],
  sumExp: [SumExp, null],
  lse: [Lse, null],
  lme: [Lme, null],
};

function parse(expr) {
  const tokens = expr.trim().split(/\s+/);
  const stack = [];
  for (const tok of tokens) {
    if (tok in OPERATIONS) {
      const [Ctor, arity] = OPERATIONS[tok];
      const args = stack.splice(stack.length - arity, arity);
      stack.push(new Ctor(...args));
    } else if (tok in VAR_INDEX) {
      stack.push(new Variable(tok));
    } else {
      stack.push(new Const(Number(tok)));
    }
  }
  return stack[0];
}

function parsePostfix(expr) {
  const tokens = expr.trim().match(BRACKET_RE) || [];
  if (tokens.length === 0) {
    throw new ParseError("Empty expression");
  }
  let pos = 0;

  function peek() {
    return tokens[pos];
  }
  function consume() {
    return tokens[pos++];
  }

  function parseArg() {
    const val = parseExpr();
    if (typeof val === "string") {
      throw new ParseError(`Unexpected operator '${val}' in argument position`);
    }
    return val;
  }

  function parseExpr() {
    if (pos >= tokens.length) {
      throw new ParseError("Unexpected end of input");
    }
    const tok = peek();

    if (OPEN_BRACKETS.has(tok)) {
      const openBracket = consume();
      const closeBracket = OPEN_TO_CLOSE[openBracket];

      const args = [];
      while (peek() !== closeBracket) {
        if (peek() === undefined) {
          throw new ParseError(`Missing closing '${closeBracket}'`);
        }
        const item = parseExpr();
        args.push(item);
      }
      consume();

      if (args.length === 0) {
        throw new ParseError("Empty parenthesised expression");
      }

      const opSym = args[args.length - 1];
      if (typeof opSym !== "string") {
        throw new ParseError(
          "Expected operator at end of parenthesised expression, got expression",
        );
      }
      if (!(opSym in OPERATIONS)) {
        throw new ParseError(`Unknown operator: '${opSym}'`);
      }

      const exprArgs = args.slice(0, args.length - 1);
      for (let i = 0; i < exprArgs.length; i++) {
        if (typeof exprArgs[i] === "string") {
          throw new ParseError(
            `Unexpected operator '${exprArgs[i]}' in argument position`,
          );
        }
      }

      const [Ctor, arity] = OPERATIONS[opSym];
      if (arity !== null && exprArgs.length !== arity) {
        throw new ParseError(
          `Operator '${opSym}' expects ${arity} argument(s), got ${exprArgs.length}`,
        );
      }
      return new Ctor(...exprArgs);
    }

    if (CLOSE_BRACKETS.has(tok)) {
      throw new ParseError(`Unexpected closing bracket: '${tok}'`);
    }

    consume();
    if (tok in OPERATIONS) {
      return tok;
    }
    if (tok in VAR_INDEX) {
      return new Variable(tok);
    }
    const n = Number(tok);
    if (isNaN(n)) {
      throw new ParseError(`Unknown token: '${tok}'`);
    }
    return new Const(n);
  }

  const result = parseExpr();

  if (typeof result === "string") {
    throw new ParseError(`Unexpected operator token: '${result}'`);
  }
  if (pos < tokens.length) {
    throw new ParseError(`Unexpected token after expression: '${tokens[pos]}'`);
  }
  return result;
}

function parsePrefix(expr) {
  const tokens = expr.trim().match(BRACKET_RE) || [];
  let pos = 0;

  function peek() {
    return tokens[pos];
  }
  function consume() {
    return tokens[pos++];
  }

  function parseExpr() {
    if (pos >= tokens.length) {
      throw new ParseError("Unexpected end of input");
    }
    const tok = peek();

    if (OPEN_BRACKETS.has(tok)) {
      const openBracket = consume();
      const closeBracket = OPEN_TO_CLOSE[openBracket];

      const op = consume();
      if (op === undefined) {
        throw new ParseError(`Expected operator after '${openBracket}'`);
      }
      if (!(op in OPERATIONS)) {
        throw new ParseError(`Unknown operator: '${op}'`);
      }
      const [Ctor, arity] = OPERATIONS[op];

      const args = [];
      if (arity === null) {
        while (peek() !== closeBracket) {
          if (peek() === undefined) {
            throw new ParseError(
              `Expected '${closeBracket}' but reached end of input`,
            );
          }
          args.push(parseExpr());
        }
      } else {
        for (let i = 0; i < arity; i++) {
          args.push(parseExpr());
        }
        if (peek() !== closeBracket) {
          throw new ParseError(
            `Expected '${closeBracket}' after ${arity} argument(s) of '${op}', ` +
              `got ${peek() === undefined ? "end of input" : "'" + peek() + "'"}`,
          );
        }
      }
      consume();

      return new Ctor(...args);
    }

    consume();
    if (tok in VAR_INDEX) {
      return new Variable(tok);
    }
    if (CLOSE_BRACKETS.has(tok)) {
      throw new ParseError(`Unexpected closing bracket: '${tok}'`);
    }
    const n = Number(tok);
    if (!isNaN(n)) {
      return new Const(n);
    }
    throw new ParseError(`Unexpected token: '${tok}'`);
  }

  const result = parseExpr();

  if (pos < tokens.length) {
    throw new ParseError(`Unexpected token after expression: '${tokens[pos]}'`);
  }

  return result;
}

if (typeof module !== "undefined") {
  module.exports = {
    Const,
    Variable,
    Add,
    Subtract,
    Multiply,
    Divide,
    Negate,
    Power,
    Log,
    Sum1,
    Sum2,
    Sum3,
    Sum4,
    Sum5,
    Avg1,
    Avg2,
    Avg3,
    Avg4,
    Avg5,
    Gauss,
    Wrap,
    SoftWrap,
    ArcTan12,
    MeanExp,
    SumExp,
    Lse,
    Lme,
    parse,
    parsePrefix,
    parsePostfix,
  };
}