import dagre from "@dagrejs/dagre";

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

const NODE_SIZE = 40;
// How far a node's value and derivative sit from its centre, above and below its shape
export const SLOT_OFFSET = 28;
// Room kept for a node with its value and derivative
const NODE_ROOM_WIDTH = 60;
const NODE_ROOM_HEIGHT = 76;
const CHAR_WIDTH = 7.2;
// Room kept around each edge for its name
const LABEL_WIDTH = 44;
const LABEL_HEIGHT = 20;
// Flatter than this, a label can stay centred on its edge
const LEVEL_SLOPE = 0.08;
// A hidden node before every input, pulling them all into the first column
const SOURCE = "\0source";
// The box at the end, f
export const OUTPUT = "\0output";

interface Point {
  x: number;
  y: number;
}

export interface PlacedShape {
  id: string;
  kind: GraphNode["kind"] | "output";
  label: string;
  // The centre
  x: number;
  y: number;
  width: number;
  height: number;
}

export interface PlacedEdge {
  id: string;
  from: string;
  to: string;
  points: Point[];
  // The point on the edge that its label straddles
  labelAt: Point;
  // The edge's direction there, going right: it rises (-1), falls (1) or is level (0)
  labelSlope: -1 | 0 | 1;
}

export interface GraphLayout {
  shapes: PlacedShape[];
  edges: PlacedEdge[];
  width: number;
  height: number;
}

const INFIX: Partial<Record<Operation, string>> = { add: "+", sub: "−", mul: "·", div: "/", pow: "^" };
const SYMBOLS: Partial<Record<Operation, string>> = { add: "+", sub: "−", mul: "*", div: "/", pow: "^", neg: "−", sqrt: "√" };

// What a node computes, e.g. "w1 + 2" or "sin(w1)"; an input's is its name.
// Pass show to write its arguments another way, such as their numbers.
export function expressionOf(node: GraphNode, show: (arg: string) => string = (arg) => arg): string {
  if (node.kind === "input") {
    return node.name;
  }
  const [a, b] = node.args.map(show);
  if (node.op in INFIX) {
    return `${a} ${INFIX[node.op]} ${b}`;
  }
  return node.op === "neg" ? `−${a}` : `${node.op}(${a})`;
}

// What a node's shape shows: its name, or its operator with any constant, e.g. "2 ^"
function labelOf(node: GraphNode, graph: Graph): string {
  if (node.kind === "input") {
    return node.name;
  }
  const symbol = SYMBOLS[node.op] ?? node.op;
  const constant = (arg: string | undefined) => arg !== undefined && arg in graph.constants;
  if (constant(node.args[0])) {
    return `${node.args[0]} ${symbol}`;
  }
  return constant(node.args[1]) ? `${symbol} ${node.args[1]}` : symbol;
}

// A smooth line through the points
export function edgePath({ points }: PlacedEdge): string {
  const [first, ...rest] = points;
  const last = rest.pop()!;
  let path = `M ${first.x} ${first.y}`;
  for (let i = 0; i < rest.length; i++) {
    const next = rest[i + 1] ?? last;
    path += ` Q ${rest[i].x} ${rest[i].y} ${(rest[i].x + next.x) / 2} ${(rest[i].y + next.y) / 2}`;
  }
  return `${path} L ${last.x} ${last.y}`;
}

// Where the line from a shape's centre towards a point leaves the shape
function onShape(shape: PlacedShape, toward: Point): Point {
  const dx = toward.x - shape.x;
  const dy = toward.y - shape.y;
  const across = Math.abs(dx / (shape.width / 2));
  const down = Math.abs(dy / (shape.height / 2));
  const scale = 1 / (shape.kind === "op" ? Math.hypot(across, down) : Math.max(across, down));
  return { x: shape.x + dx * scale, y: shape.y + dy * scale };
}

// Where the drawn curve passes a dagre point, and which way it is heading
function labelPlacement(points: Point[], label: Point): Pick<PlacedEdge, "labelAt" | "labelSlope"> {
  const distance = (point: Point) => Math.hypot(point.x - label.x, point.y - label.y);
  const i = points.reduce((best, point, j) => (distance(point) < distance(points[best]) ? j : best), 0);
  const before = points[Math.max(i - 1, 0)];
  const after = points[Math.min(i + 1, points.length - 1)];
  const slope = (after.y - before.y) / (after.x - before.x);
  return {
    // Edges are smoothed, so they cut the corner at each point
    labelAt: { x: (before.x + 6 * points[i].x + after.x) / 8, y: (before.y + 6 * points[i].y + after.y) / 8 },
    labelSlope: Math.abs(slope) < LEVEL_SLOPE ? 0 : slope < 0 ? -1 : 1,
  };
}

// Left to right in evaluation order, with the inputs in a column on the left
export function layoutGraph(graph: Graph): GraphLayout {
  const g = new dagre.graphlib.Graph({ multigraph: true });
  g.setGraph({ rankdir: "LR", nodesep: 12, ranksep: 16, edgesep: 28, marginx: 4, marginy: 4 });
  g.setDefaultEdgeLabel(() => ({}));

  const size = (label: string) => ({ width: Math.max(NODE_SIZE, label.length * CHAR_WIDTH + 24), height: NODE_SIZE });
  const unplaced: Omit<PlacedShape, "x" | "y">[] = graph.nodes.map((node) => {
    const label = labelOf(node, graph);
    return { id: node.id, kind: node.kind, label, ...size(label) };
  });
  unplaced.push({ id: OUTPUT, kind: "output", label: "f", ...size("f") });
  for (const { id, width } of unplaced) {
    g.setNode(id, { width: Math.max(width, NODE_ROOM_WIDTH), height: NODE_ROOM_HEIGHT });
  }

  const nodeIds = new Set(graph.nodes.map((node) => node.id));
  const edges: { from: string; to: string; id: string }[] = [{ from: graph.output, to: OUTPUT, id: `${graph.output}-output` }];
  for (const node of graph.nodes) {
    if (node.kind === "op") {
      node.args.forEach((from, i) => {
        if (nodeIds.has(from)) {
          edges.push({ from, to: node.id, id: `${from}-${node.id}-${i}` });
        }
      });
    }
  }
  for (const edge of edges) {
    g.setEdge(edge.from, edge.to, { width: LABEL_WIDTH, height: LABEL_HEIGHT, labelpos: "c" }, edge.id);
  }

  g.setNode(SOURCE, { width: 0, height: 0 });
  for (const name of graph.nodes.filter((node) => node.kind === "input")) {
    g.setEdge(SOURCE, name.id, { weight: 100 }, SOURCE);
  }

  dagre.layout(g);

  // dagre doesn't keep the inputs' order, so mirror the drawing if they came out reversed
  const inputY = graph.nodes.filter((node) => node.kind === "input").map((node) => g.node(node.id).y);
  const mirrored = inputY[0] > inputY[inputY.length - 1];

  const { height = 0 } = g.graph();
  // Measure without the hidden node's column
  const left = Math.min(...unplaced.map(({ id }) => g.node(id).x - g.node(id).width / 2)) - 4;
  const place = ({ x, y }: Point): Point => ({ x: x - left, y: mirrored ? height - y : y });

  const shapes: PlacedShape[] = unplaced.map((shape) => ({ ...shape, ...place(g.node(shape.id)) }));
  const shapeOf = new Map(shapes.map((shape) => [shape.id, shape]));
  const width = Math.max(...unplaced.map(({ id }) => g.node(id).x + g.node(id).width / 2)) - left + 4;

  return {
    shapes,
    width,
    height,
    edges: edges.map(({ from, to, id }) => {
      const edge = g.edge(from, to, id);
      const points: Point[] = edge.points.map(place);
      // dagre ends edges on the room kept for a node, not on its shape
      points[0] = onShape(shapeOf.get(from)!, points[1]);
      points[points.length - 1] = onShape(shapeOf.get(to)!, points[points.length - 2]);
      return { id, from, to, points, ...labelPlacement(points, place(edge)) };
    }),
  };
}