Entropyk/apps/web/src/store/diagramStore.test.ts

import { describe, it, expect, beforeEach } from "vitest";
import {
  useDiagramStore,
  snapToGridValue,
  snapPosition,
  normaliseRotation,
  GRID_SIZE,
} from "./diagramStore";

function reset() {
  useDiagramStore.setState({
    nodes: [],
    edges: [],
    selectedNodeId: null,
    fluid: "R410A",
    fluidBackend: "CoolProp",
    controls: [],
    snapToGrid: true,
    result: null,
    simulating: false,
    simError: null,
  });
}

beforeEach(reset);

describe("controls import", () => {
  it("preserves imported co-solved controls for the next simulation run", () => {
    useDiagramStore.getState().loadFromConfig({
      fluid: "R134a",
      circuits: [
        {
          id: 0,
          components: [{ type: "IsentropicCompressor", name: "comp", liquid_injection: true }],
          edges: [],
        },
      ],
      controls: [
        {
          type: "SaturatedController",
          id: "dgt_limiter",
          measure: { component: "comp", output: "temperature" },
          actuator: { component: "comp", factor: "injection", initial: 0.15, min: 0, max: 0.3 },
          target: 330,
          gain: -0.5,
          band: 5,
          alpha: 0.002,
          objectives: [
            {
              component: "comp",
              output: "temperature",
              setpoint: 365,
              gain: -0.05,
              combine: "min",
            },
          ],
        },
      ],
      solver: { strategy: "newton", max_iterations: 300, tolerance: 1e-6 },
    });

    expect(useDiagramStore.getState().controls).toEqual([
      expect.objectContaining({
        id: "dgt_limiter",
        actuator: expect.objectContaining({ factor: "injection", initial: 0.15 }),
        target: 330,
      }),
    ]);
    expect(useDiagramStore.getState().nodes).toEqual(
      expect.arrayContaining([
        expect.objectContaining({
          data: expect.objectContaining({
            type: "SaturatedController",
            name: "dgt_limiter",
            params: expect.objectContaining({
              measure_component: "comp",
              actuator_factor: "injection",
              alpha: 0.002,
              objectives_json: JSON.stringify([
                {
                  component: "comp",
                  output: "temperature",
                  setpoint: 365,
                  gain: -0.05,
                  combine: "min",
                },
              ]),
            }),
          }),
        }),
      ]),
    );
  });
});

describe("grid helpers", () => {
  it("snaps a value to the nearest grid line", () => {
    expect(snapToGridValue(5, 12)).toBe(0);
    expect(snapToGridValue(7, 12)).toBe(12);
    expect(snapToGridValue(18, 12)).toBe(24);
  });

  it("snaps an {x, y} position", () => {
    expect(snapPosition({ x: 7, y: 5 }, 12)).toEqual({ x: 12, y: 0 });
  });

  it("normalises rotation to 0/90/180/270", () => {
    expect(normaliseRotation(0)).toBe(0);
    expect(normaliseRotation(90)).toBe(90);
    expect(normaliseRotation(360)).toBe(0);
    expect(normaliseRotation(-90)).toBe(270);
    expect(normaliseRotation(450)).toBe(90);
  });
});

describe("addComponent", () => {
  it("creates a node with defaults, rotation 0, and selects it", () => {
    useDiagramStore.getState().addComponent("Condenser", { x: 100, y: 100 });
    const { nodes, selectedNodeId } = useDiagramStore.getState();
    expect(nodes).toHaveLength(1);
    expect(nodes[0].data.type).toBe("Condenser");
    expect(nodes[0].data.rotation).toBe(0);
    expect(nodes[0].data.params.ua).toBe(5000);
    expect(selectedNodeId).toBe(nodes[0].id);
  });

  it("snaps the drop position when snap is enabled", () => {
    useDiagramStore.getState().addComponent("Condenser", { x: 103, y: 97 });
    const n = useDiagramStore.getState().nodes[0];
    expect(n.position.x % GRID_SIZE).toBe(0);
    expect(n.position.y % GRID_SIZE).toBe(0);
  });

  it("keeps the raw position when snap is disabled", () => {
    useDiagramStore.setState({ snapToGrid: false });
    useDiagramStore.getState().addComponent("Condenser", { x: 103, y: 97 });
    const n = useDiagramStore.getState().nodes[0];
    expect(n.position).toEqual({ x: 103, y: 97 });
  });

  it("generates unique names per type", () => {
    const s = useDiagramStore.getState();
    s.addComponent("Condenser", { x: 0, y: 0 });
    s.addComponent("Condenser", { x: 0, y: 0 });
    const names = useDiagramStore.getState().nodes.map((n) => n.data.name);
    expect(new Set(names).size).toBe(2);
  });
});

describe("node mutations", () => {
  it("updates params, name, and circuit", () => {
    useDiagramStore.getState().addComponent("Condenser", { x: 0, y: 0 });
    const id = useDiagramStore.getState().nodes[0].id;
    const s = useDiagramStore.getState();
    s.updateNodeParams(id, { ua: 9999 });
    s.updateNodeName(id, "my_cond");
    s.updateNodeCircuit(id, 2);
    const n = useDiagramStore.getState().nodes[0];
    expect(n.data.params.ua).toBe(9999);
    expect(n.data.name).toBe("my_cond");
    expect(n.data.circuit).toBe(2);
  });

  it("rotates clockwise and counter-clockwise with wrap-around", () => {
    useDiagramStore.getState().addComponent("Compressor", { x: 0, y: 0 });
    const id = useDiagramStore.getState().nodes[0].id;
    const { rotateNode } = useDiagramStore.getState();
    rotateNode(id, 1);
    expect(useDiagramStore.getState().nodes[0].data.rotation).toBe(90);
    rotateNode(id, 1);
    rotateNode(id, 1);
    rotateNode(id, 1);
    expect(useDiagramStore.getState().nodes[0].data.rotation).toBe(0);
    rotateNode(id, -1);
    expect(useDiagramStore.getState().nodes[0].data.rotation).toBe(270);
  });

  it("flips horizontal and vertical like Modelica", () => {
    useDiagramStore.getState().addComponent("RefrigerantPipe", { x: 0, y: 0 });
    const id = useDiagramStore.getState().nodes[0].id;
    expect(useDiagramStore.getState().nodes[0].data.flipH).toBe(false);
    useDiagramStore.getState().flipNodeH(id);
    expect(useDiagramStore.getState().nodes[0].data.flipH).toBe(true);
    useDiagramStore.getState().flipNodeV(id);
    expect(useDiagramStore.getState().nodes[0].data.flipV).toBe(true);
  });
});

describe("connections and removal", () => {
  it("adds an edge on connect", () => {
    const s = useDiagramStore.getState();
    s.addComponent("IsentropicCompressor", { x: 0, y: 0 });
    s.addComponent("Condenser", { x: 200, y: 0 });
    const [a, b] = useDiagramStore.getState().nodes;
    useDiagramStore.getState().onConnect({
      source: a.id,
      target: b.id,
      sourceHandle: "outlet",
      targetHandle: "inlet",
    });
    const edges = useDiagramStore.getState().edges;
    expect(edges).toHaveLength(1);
    expect(edges[0].source).toBe(a.id);
    expect(edges[0].target).toBe(b.id);
  });

  it("removes a node along with its connected edges and clears selection", () => {
    const s = useDiagramStore.getState();
    s.addComponent("IsentropicCompressor", { x: 0, y: 0 });
    s.addComponent("Condenser", { x: 200, y: 0 });
    const [a, b] = useDiagramStore.getState().nodes;
    useDiagramStore.getState().onConnect({
      source: a.id,
      target: b.id,
      sourceHandle: "outlet",
      targetHandle: "inlet",
    });
    useDiagramStore.getState().setSelected(a.id);
    useDiagramStore.getState().removeNode(a.id);
    const st = useDiagramStore.getState();
    expect(st.nodes).toHaveLength(1);
    expect(st.edges).toHaveLength(0);
    expect(st.selectedNodeId).toBeNull();
  });

  it("clears selection when the selected node is removed via onNodesChange", () => {
    useDiagramStore.getState().addComponent("Condenser", { x: 0, y: 0 });
    const id = useDiagramStore.getState().nodes[0].id;
    useDiagramStore.getState().setSelected(id);
    useDiagramStore.getState().onNodesChange([{ type: "remove", id }]);
    expect(useDiagramStore.getState().selectedNodeId).toBeNull();
    expect(useDiagramStore.getState().nodes).toHaveLength(0);
  });
});

describe("toggleSnapToGrid", () => {
  it("flips the snap flag", () => {
    expect(useDiagramStore.getState().snapToGrid).toBe(true);
    useDiagramStore.getState().toggleSnapToGrid();
    expect(useDiagramStore.getState().snapToGrid).toBe(false);
    useDiagramStore.getState().toggleSnapToGrid();
    expect(useDiagramStore.getState().snapToGrid).toBe(true);
  });
});

describe("loadFromConfig and clear", () => {
  const config = {
    fluid: "R290",
    fluid_backend: "Test",
    circuits: [
      {
        id: 0,
        components: [
          { type: "IsentropicCompressor", name: "comp", isentropic_efficiency: 0.7 },
          { type: "Condenser", name: "cond", ua: 5000 },
        ],
        edges: [{ from: "comp:outlet", to: "cond:inlet" }],
      },
    ],
  };

  it("rebuilds nodes, edges, rotation and fluid from a config", () => {
    useDiagramStore.getState().loadFromConfig(config);
    const st = useDiagramStore.getState();
    expect(st.nodes).toHaveLength(2);
    expect(st.edges).toHaveLength(1);
    expect(st.fluid).toBe("R290");
    expect(st.fluidBackend).toBe("Test");
    for (const n of st.nodes) expect(n.data.rotation).toBe(0);
    // the params should be carried over as primitives
    const comp = st.nodes.find((n) => n.data.name === "comp");
    expect(comp?.data.params.isentropic_efficiency).toBe(0.7);
  });

  it("ignores configs without circuits", () => {
    useDiagramStore.getState().addComponent("Condenser", { x: 0, y: 0 });
    useDiagramStore.getState().loadFromConfig({ foo: "bar" });
    expect(useDiagramStore.getState().nodes).toHaveLength(1);
  });

  it("flattens subsystem instances from system model JSON", () => {
    useDiagramStore.getState().loadFromConfig({
      fluid: "R134a",
      fluid_backend: "CoolProp",
      subsystems: {
        Circuit: {
          params: { ua_cond: 700 },
          components: [
            { type: "IsentropicCompressor", name: "comp", isentropic_efficiency: 0.7 },
            { type: "Condenser", name: "cond", ua: "$ua_cond" },
          ],
          edges: [{ from: "comp:outlet", to: "cond:inlet" }],
          ports: { discharge: "comp:outlet" },
        },
      },
      instances: [
        { of: "Circuit", name: "A", circuit: 0, params: { ua_cond: 800 } },
        { of: "Circuit", name: "B", circuit: 1 },
      ],
    });

    const st = useDiagramStore.getState();
    expect(st.nodes).toHaveLength(4);
    expect(st.edges).toHaveLength(2);
    expect(st.nodes.map((n) => n.data.name)).toContain("A.comp");
    expect(st.nodes.map((n) => n.data.name)).toContain("B.cond");
    expect(st.nodes.find((n) => n.data.name === "A.cond")?.data.params.ua).toBe(800);
    expect(st.nodes.find((n) => n.data.name === "B.cond")?.data.params.ua).toBe(700);
  });

  it("keeps explicit water circuits while flattening refrigerant subsystem instances", () => {
    useDiagramStore.getState().loadFromConfig({
      fluid: "R134a",
      circuits: [
        {
          id: 2,
          components: [
            { type: "BrineSource", name: "water_in", t_set_c: 30 },
            { type: "ThermalLoad", name: "water_load", mass_flow_kg_s: 1.0 },
            { type: "BrineSink", name: "water_out" },
          ],
          edges: [
            { from: "water_in:outlet", to: "water_load:inlet" },
            { from: "water_load:outlet", to: "water_out:inlet" },
          ],
        },
      ],
      subsystems: {
        Circuit: {
          components: [
            { type: "IsentropicCompressor", name: "comp", isentropic_efficiency: 0.7 },
            { type: "Condenser", name: "cond", ua: 700 },
          ],
          edges: [{ from: "comp:outlet", to: "cond:inlet" }],
        },
      },
      instances: [{ of: "Circuit", name: "A", circuit: 0 }],
    });

    const st = useDiagramStore.getState();
    expect(st.nodes.map((n) => n.data.name)).toEqual(
      expect.arrayContaining(["A.comp", "A.cond", "water_in", "water_load", "water_out"]),
    );
    expect(st.edges).toHaveLength(3);
    expect(st.nodes.find((n) => n.data.name === "water_load")?.data.circuit).toBe(2);
  });

  it("lays out imported chiller cycles as a horizontal HVAC schematic", () => {
    useDiagramStore.getState().loadFromConfig({
      fluid: "R134a",
      circuits: [
        {
          id: 0,
          components: [
            { type: "IsentropicCompressor", name: "comp", liquid_injection: true },
            { type: "Condenser", name: "cond", ua: 5000 },
            { type: "IsenthalpicExpansionValve", name: "exv", opening: 0.8 },
            { type: "Evaporator", name: "evap", ua: 6000 },
            { type: "BrineSource", name: "cond_water_in", t_set_c: 30 },
            { type: "BrineSink", name: "cond_water_out" },
            { type: "BrineSource", name: "evap_water_in", t_set_c: 12 },
            { type: "BrineSink", name: "evap_water_out" },
          ],
          edges: [
            { from: "comp:outlet", to: "cond:inlet" },
            { from: "cond:outlet", to: "exv:inlet" },
            { from: "exv:outlet", to: "evap:inlet" },
            { from: "evap:outlet", to: "comp:inlet" },
            { from: "cond_water_in:outlet", to: "cond:secondary_inlet" },
            { from: "cond:secondary_outlet", to: "cond_water_out:inlet" },
            { from: "evap_water_in:outlet", to: "evap:secondary_inlet" },
            { from: "evap:secondary_outlet", to: "evap_water_out:inlet" },
          ],
        },
      ],
      controls: [
        {
          id: "dgt_limiter",
          measure: { component: "comp", output: "temperature" },
          actuator: { component: "comp", factor: "injection", initial: 0.15, min: 0, max: 0.3 },
          target: 330,
        },
      ],
    });

    const nodes = useDiagramStore.getState().nodes;
    const byName = Object.fromEntries(nodes.map((node) => [node.data.name, node]));

    expect(byName.evap.position.x).toBeLessThan(byName.comp.position.x);
    expect(byName.comp.position.x).toBeLessThan(byName.cond.position.x);
    expect(byName.exv.position.x).toBeGreaterThan(byName.cond.position.x);
    expect(byName.exv.position.y).toBeGreaterThan(byName.cond.position.y);
    expect(byName.dgt_limiter.position.y).toBeLessThan(byName.comp.position.y);

    expect(byName.evap_water_out.position.y).toBeLessThan(byName.evap.position.y);
    expect(byName.evap_water_in.position.y).toBeGreaterThan(byName.evap.position.y);
    expect(byName.cond_water_out.position.y).toBeLessThan(byName.cond.position.y);
    expect(byName.cond_water_in.position.y).toBeGreaterThan(byName.cond.position.y);
    expect(useDiagramStore.getState().edges).toHaveLength(8);
  });

  it("uses deterministic fallback positions for partial imported configs", () => {
    useDiagramStore.getState().loadFromConfig({
      fluid: "R134a",
      circuits: [
        {
          id: 0,
          components: [
            { type: "Pipe", name: "pipe_b", length_m: 2 },
            { type: "Placeholder", name: "custom_a", n_equations: 1 },
          ],
          edges: [{ from: "custom_a:outlet", to: "pipe_b:inlet" }],
        },
      ],
    });

    const first = useDiagramStore.getState().nodes.map((node) => [node.data.name, node.position]);
    useDiagramStore.getState().loadFromConfig({
      fluid: "R134a",
      circuits: [
        {
          id: 0,
          components: [
            { type: "Pipe", name: "pipe_b", length_m: 2 },
            { type: "Placeholder", name: "custom_a", n_equations: 1 },
          ],
          edges: [{ from: "custom_a:outlet", to: "pipe_b:inlet" }],
        },
      ],
    });

    expect(useDiagramStore.getState().nodes.map((node) => [node.data.name, node.position])).toEqual(first);
    expect(useDiagramStore.getState().edges).toHaveLength(1);
  });

  it("places secondary boundaries near generic four-port heat exchangers", () => {
    useDiagramStore.getState().loadFromConfig({
      fluid: "R134a",
      circuits: [
        {
          id: 0,
          components: [
            { type: "HeatExchanger", name: "hx", ua: 3000 },
            { type: "BrineSource", name: "water_in", t_set_c: 20 },
            { type: "BrineSink", name: "water_out" },
            { type: "AirSource", name: "air_in", t_dry_c: 15 },
            { type: "AirSink", name: "air_out" },
          ],
          edges: [
            { from: "water_in:outlet", to: "hx:hot_inlet" },
            { from: "hx:hot_outlet", to: "water_out:inlet" },
            { from: "air_in:outlet", to: "hx:cold_inlet" },
            { from: "hx:cold_outlet", to: "air_out:inlet" },
          ],
        },
      ],
    });

    const byName = Object.fromEntries(useDiagramStore.getState().nodes.map((node) => [node.data.name, node]));
    expect(Math.abs(byName.water_in.position.x - byName.hx.position.x)).toBeLessThan(220);
    expect(Math.abs(byName.air_in.position.x - byName.hx.position.x)).toBeLessThan(220);
    expect(byName.water_out.position.y).toBeLessThan(byName.hx.position.y);
    expect(byName.air_out.position.y).toBeLessThan(byName.hx.position.y);
    expect(byName.water_in.position.y).toBeGreaterThan(byName.hx.position.y);
    expect(byName.air_in.position.y).toBeGreaterThan(byName.hx.position.y);
  });

  it("separates multiple imported refrigeration circuits horizontally", () => {
    useDiagramStore.getState().loadFromConfig({
      fluid: "R134a",
      circuits: [
        {
          id: 0,
          components: [
            { type: "IsentropicCompressor", name: "comp_a" },
            { type: "Condenser", name: "cond_a" },
          ],
          edges: [{ from: "comp_a:outlet", to: "cond_a:inlet" }],
        },
        {
          id: 1,
          components: [
            { type: "IsentropicCompressor", name: "comp_b" },
            { type: "Condenser", name: "cond_b" },
          ],
          edges: [{ from: "comp_b:outlet", to: "cond_b:inlet" }],
        },
      ],
    });

    const byName = Object.fromEntries(useDiagramStore.getState().nodes.map((node) => [node.data.name, node]));
    expect(byName.comp_b.position.x - byName.comp_a.position.x).toBeGreaterThan(900);
    expect(byName.cond_b.position.x - byName.cond_a.position.x).toBeGreaterThan(900);
  });

  it("clear() empties the sheet", () => {
    useDiagramStore.getState().loadFromConfig(config);
    useDiagramStore.getState().clear();
    const st = useDiagramStore.getState();
    expect(st.nodes).toHaveLength(0);
    expect(st.edges).toHaveLength(0);
    expect(st.selectedNodeId).toBeNull();
  });
});