import test from "node:test";
import assert from "node:assert/strict";
import { computeAmpMapForCircuit } from "../dist/ux/circuit-vis/state-viz/worker/stateCompute.js";
import { evaluateAngleExpression } from "../dist/ux/circuit-vis/angleExpression.js";

const approxEq = (a, b, eps = 1e-12) => Math.abs(a - b) <= eps;

test("π and pi keyword", () => {
  assert.ok(approxEq(evaluateAngleExpression("π"), Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("+π"), Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("-π"), -Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("pi"), Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("+pi"), Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("-pi"), -Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("Pi"), Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("+Pi"), Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("-Pi"), -Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("PI"), Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("+PI"), Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("-PI"), -Math.PI));
});

test("basic numbers", () => {
  assert.ok(approxEq(evaluateAngleExpression("5"), 5));
  assert.ok(approxEq(evaluateAngleExpression("+5"), 5));
  assert.ok(approxEq(evaluateAngleExpression("-5"), -5));
  assert.ok(approxEq(evaluateAngleExpression("3.5"), 3.5));
  assert.ok(approxEq(evaluateAngleExpression("+3.5"), 3.5));
  assert.ok(approxEq(evaluateAngleExpression("-3.5"), -3.5));
  assert.ok(approxEq(evaluateAngleExpression("5."), 5));
  assert.ok(approxEq(evaluateAngleExpression("+5."), 5));
  assert.ok(approxEq(evaluateAngleExpression("-5."), -5));
});

test("arithmetic operations", () => {
  assert.ok(approxEq(evaluateAngleExpression("π/2"), Math.PI / 2));
  assert.ok(approxEq(evaluateAngleExpression("-π/2"), -Math.PI / 2));
  assert.ok(approxEq(evaluateAngleExpression("2*pi"), 2 * Math.PI));
  assert.ok(approxEq(evaluateAngleExpression("π + 2 - 3"), Math.PI - 1));
  assert.ok(approxEq(evaluateAngleExpression("2 * (pi / 4)"), Math.PI / 2));
});

test("parentheses nesting", () => {
  assert.ok(approxEq(evaluateAngleExpression("((π))"), Math.PI));
});

test("invalid inputs return undefined", () => {
  assert.equal(evaluateAngleExpression("++π"), undefined);
  assert.equal(evaluateAngleExpression("--π"), undefined);
  assert.equal(evaluateAngleExpression("π // 2"), undefined);
  assert.equal(evaluateAngleExpression("1..2"), undefined);
  assert.equal(evaluateAngleExpression("(π"), undefined);
  assert.equal(evaluateAngleExpression("π / 0"), undefined); // Infinity -> undefined
  assert.equal(evaluateAngleExpression(""), undefined);
  assert.equal(evaluateAngleExpression(".5"), undefined);
  assert.equal(evaluateAngleExpression("+.5"), undefined);
  assert.equal(evaluateAngleExpression("-.5"), undefined);
});

const colUnitaryAt = (gate, target, args, opts) => ({
  components: [
    {
      kind: "unitary",
      gate,
      targets: [{ qubit: target }],
      ...(opts?.controls?.length
        ? { controls: opts.controls.map((qubit) => ({ qubit })) }
        : null),
      args,
      ...(opts?.isAdjoint ? { isAdjoint: true } : null),
    },
  ],
});
const colUnitary = (gate, args, opts) => colUnitaryAt(gate, 0, args, opts);
const colReset0 = () => ({
  components: [{ kind: "ket", gate: "0", targets: [{ qubit: 0 }] }],
});
const colMeasure0 = () => ({
  components: [
    {
      kind: "measurement",
      gate: "M",
      qubits: [{ qubit: 0 }],
      results: [{ qubit: 0, result: 0 }],
    },
  ],
});

const assertAmp = (amp, re, im) => {
  assert.ok(approxEq(amp.re, re, 1e-11), `re expected ${re} got ${amp.re}`);
  assert.ok(approxEq(amp.im, im, 1e-11), `im expected ${im} got ${amp.im}`);
};

test("Single adjoint: S† on |1⟩ yields -i|1⟩", () => {
  const qubits = [{ id: 0 }];
  const componentGrid = [
    // Prepare |1⟩ then apply S†: S†|1⟩ = -i|1⟩
    colUnitary("X"),
    colUnitary("S", undefined, { isAdjoint: true }),
  ];
  const ampMap = computeAmpMapForCircuit(qubits, componentGrid);
  assertAmp(ampMap["1"], 0, -1);
});

test("Single adjoint: T† on |1⟩ yields e^{-iπ/4}|1⟩", () => {
  const qubits = [{ id: 0 }];
  const componentGrid = [
    // Prepare |1⟩ then apply T†: T†|1⟩ = e^{-iπ/4}|1⟩
    colUnitary("X"),
    colUnitary("T", undefined, { isAdjoint: true }),
  ];
  const ampMap = computeAmpMapForCircuit(qubits, componentGrid);
  assertAmp(ampMap["1"], Math.SQRT1_2, -Math.SQRT1_2);
});

test("Single adjoint: SX† on |0⟩ matches expected amplitudes", () => {
  const qubits = [{ id: 0 }];
  const componentGrid = [colUnitary("SX", undefined, { isAdjoint: true })];
  const ampMap = computeAmpMapForCircuit(qubits, componentGrid);
  // SX†|0⟩ = (0.5-0.5i)|0⟩ + (0.5+0.5i)|1⟩
  assertAmp(ampMap["0"], 0.5, -0.5);
  assertAmp(ampMap["1"], 0.5, 0.5);
});

test("Gate then adjoint returns |0⟩ (Rx(π/3))", () => {
  const qubits = [{ id: 0 }];
  const componentGrid = [
    colUnitary("Rx", ["π/3"]),
    colUnitary("Rx", ["π/3"], { isAdjoint: true }),
  ];
  const ampMap = computeAmpMapForCircuit(qubits, componentGrid);
  assertAmp(ampMap["0"], 1, 0);
});

test("Single adjoint: Ry†(π/3) from |0⟩ flips |1⟩ sign", () => {
  const qubits = [{ id: 0 }];
  const componentGrid = [colUnitary("Ry", ["π/3"], { isAdjoint: true })];
  const ampMap = computeAmpMapForCircuit(qubits, componentGrid);
  const c = Math.cos(Math.PI / 6);
  const s = Math.sin(Math.PI / 6);
  assertAmp(ampMap["0"], c, 0);
  assertAmp(ampMap["1"], -s, 0);
});

test("Single adjoint: Rz†(π/3) on |1⟩ applies e^{-iπ/6}", () => {
  const qubits = [{ id: 0 }];
  const componentGrid = [
    colUnitary("X"),
    colUnitary("Rz", ["π/3"], { isAdjoint: true }),
  ];
  const ampMap = computeAmpMapForCircuit(qubits, componentGrid);
  assertAmp(ampMap["1"], Math.cos(-Math.PI / 6), Math.sin(-Math.PI / 6));
});

test("Single adjoint: S† after H gives |0⟩ - i|1⟩ over √2", () => {
  const qubits = [{ id: 0 }];
  const componentGrid = [
    colUnitary("H"),
    colUnitary("S", undefined, { isAdjoint: true }),
  ];
  const ampMap = computeAmpMapForCircuit(qubits, componentGrid);
  assertAmp(ampMap["0"], Math.SQRT1_2, 0);
  assertAmp(ampMap["1"], 0, -Math.SQRT1_2);
});

test("Single adjoint: controlled S† phases |11⟩ by -i", () => {
  const qubits = [{ id: 0 }, { id: 1 }];
  const componentGrid = [
    colUnitaryAt("X", 0),
    colUnitaryAt("X", 1),
    colUnitaryAt("S", 1, undefined, { controls: [0], isAdjoint: true }),
  ];
  const ampMap = computeAmpMapForCircuit(qubits, componentGrid);
  assertAmp(ampMap["11"], 0, -1);
});

test("Single adjoint: controlled Rz† no-ops when control is |0⟩", () => {
  const qubits = [{ id: 0 }, { id: 1 }];
  const componentGrid = [
    // Prepare |01⟩ (control qubit 0 is 0, target qubit 1 is 1)
    colUnitaryAt("X", 1),
    colUnitaryAt("Rz", 1, ["π/3"], { controls: [0], isAdjoint: true }),
  ];
  const ampMap = computeAmpMapForCircuit(qubits, componentGrid);
  assertAmp(ampMap["01"], 1, 0);
});

test("Reset is unsupported by stateCompute", () => {
  const qubits = [{ id: 0 }];
  const componentGrid = [colUnitary("H"), colReset0()];
  assert.throws(
    () => computeAmpMapForCircuit(qubits, componentGrid),
    (err) => {
      assert.equal(err?.name, "UnsupportedStateComputeError");
      assert.equal(
        err?.message,
        "State visualization does not currently support measurement or ResetZ / |0⟩ reset operations.",
      );
      return true;
    },
  );
});

test("Measurement is unsupported by stateCompute", () => {
  const qubits = [{ id: 0 }];
  const componentGrid = [colUnitary("H"), colMeasure0()];
  assert.throws(
    () => computeAmpMapForCircuit(qubits, componentGrid),
    (err) => {
      assert.equal(err?.name, "UnsupportedStateComputeError");
      assert.equal(
        err?.message,
        "State visualization does not currently support measurement or ResetZ / |0⟩ reset operations.",
      );
      return true;
    },
  );
});
