import type { Graph, GraphNode, Mode, Operation } from "./types";

// Rounded for reading
export function formatNumber(value: number): string {
  return String(Number(value.toPrecision(5)));
}

// A multi-part argument, such as the constant "2 * 3", in brackets
function atom(arg: string): string {
  return arg.includes(" ") ? `(${arg})` : arg;
}

// Each operation's derivative with respect to each argument; backend/src/logic.py's PARTIALS as text
const PARTIALS: Record<Operation, ((a: string, b: string) => string)[]> = {
  add: [() => "1", () => "1"],
  sub: [() => "1", () => "−1"],
  mul: [(_, b) => b, (a) => a],
  div: [(_, b) => `1 / ${b}`, (a, b) => `−${a} / ${b}^2`],
  pow: [(a, b) => `${b} · ${a}^(${b} − 1)`, (a, b) => `${a}^${b} · log(${a})`],
  neg: [() => "−1"],
  sin: [(a) => `cos(${a})`],
  cos: [(a) => `−sin(${a})`],
  tan: [(a) => `1 / cos(${a})^2`],
  exp: [(a) => `exp(${a})`],
  log: [(a) => `1 / ${a}`],
  sqrt: [(a) => `1 / (2 · sqrt(${a}))`],
};

export function partialOf(op: Operation, index: number, args: string[]): string {
  return PARTIALS[op][index](atom(args[0]), atom(args[1] ?? ""));
}

// An id with a bar over its first letter: the adjoint of that node
export function adjoint(id: string): string {
  return `${id[0]}̄${id.slice(1)}`;
}

// An id with a dot over its first letter: the tangent of that node
function tangent(id: string): string {
  return `${id[0]}\u0307${id.slice(1)}`;
}

export function derivativeSymbol(mode: Mode, id: string): string {
  return mode === "forward" ? tangent(id) : adjoint(id);
}

// A derivative times a local partial, leaving out a factor of 1
function term(symbol: string, partial: string): string {
  if (partial === "1") {
    return symbol;
  }
  if (partial === "−1") {
    return `−${symbol}`;
  }
  return `${symbol} · ${partial.includes(" ") ? `(${partial})` : partial}`;
}

function sumOf(terms: string[]): string {
  return terms.join(" + ").replaceAll("+ −", "− ");
}

// How reverse mode gets a node's adjoint: one term for each use of the node
export function adjointFormula(graph: Graph, node: GraphNode): string {
  if (node.id === graph.output) {
    return `${adjoint(node.id)} = 1`;
  }
  const terms: string[] = [];
  for (const consumer of graph.nodes) {
    if (consumer.kind === "op") {
      consumer.args.forEach((arg, i) => {
        if (arg === node.id) {
          terms.push(term(adjoint(consumer.id), partialOf(consumer.op, i, consumer.args)));
        }
      });
    }
  }
  return `${adjoint(node.id)} = ${sumOf(terms)}`;
}

// A node's arguments that have derivatives, in the order of a step's partials; constants don't
export function nodeArgs(graph: Graph, node: GraphNode & { kind: "op" }): { arg: string; index: number }[] {
  return node.args.map((arg, index) => ({ arg, index })).filter(({ arg }) => !(arg in graph.constants));
}

// How forward mode gets a node's tangent: one term for each of its arguments
export function tangentFormula(graph: Graph, node: GraphNode, seed: string): string {
  if (node.kind === "input") {
    return `${tangent(node.id)} = ${node.name === seed ? 1 : 0}`;
  }
  const terms = nodeArgs(graph, node).map(({ arg, index }) => term(tangent(arg), partialOf(node.op, index, node.args)));
  return `${tangent(node.id)} = ${sumOf(terms)}`;
}