import { Tensor, interpolate_4d, matmul, rfft, slice } from "../../src/transformers.js";
import { init } from "../init.js";

// Initialise the testing environment
init();

function expectToBeCloseToArray(actual, expected) {
  expect(actual.length).toEqual(expected.length);
  actual.forEach((x, i) => expect(x).toBeCloseTo(expected[i]), 6);
}

function range(start, stop = undefined, step = 1) {
  if (stop === undefined) {
    stop = start;
    start = 0;
  }

  const result = [];
  for (let i = start; i < stop; i += step) {
    result.push(i);
  }
  return result;
}

describe("Tensor operations", () => {
  describe("interpolate", () => {
    describe("downscale", () => {
      const input = new Tensor(
        "float32",
        new Float32Array(2 * 3 * 4 * 5).map((_, i) => i),
        [2, 3, 4, 5],
      );

      const size = [2, 3, 3, 2];
      it("bilinear", async () => {
        const resized = await interpolate_4d(input, { mode: "bilinear", size });
        const target = new Float32Array(
          [
            [
              [
                [1.5833335, 4.0833335],
                [8.25, 10.75],
                [14.916668, 17.416668],
              ],
              [
                [21.583332, 24.083334],
                [28.25, 30.75],
                [34.916668, 37.416668],
              ],
              [
                [41.583332, 44.083332],
                [48.25, 50.75],
                [54.916668, 57.416668],
              ],
            ],
            [
              [
                [61.583332, 64.083336],
                [68.25, 70.75],
                [74.916664, 77.41667],
              ],
              [
                [81.58333, 84.083336],
                [88.25, 90.75],
                [94.91667, 97.41667],
              ],
              [
                [101.583336, 104.08333],
                [108.25, 110.75],
                [114.916664, 117.416664],
              ],
            ],
          ].flat(Infinity),
        );

        expectToBeCloseToArray(target, resized.data);
      });

      it("bicubic", async () => {
        const resized = await interpolate_4d(input, { mode: "bicubic", size });

        const target = new Float32Array(
          [
            [
              [
                [1.2987545, 3.9628172],
                [8.167969, 10.832031],
                [15.037184, 17.701244],
              ],
              [
                [21.298756, 23.962818],
                [28.167969, 30.832031],
                [35.037186, 37.701252],
              ],
              [
                [41.298756, 43.96282],
                [48.16797, 50.83203],
                [55.037193, 57.701256],
              ],
            ],
            [
              [
                [61.29875, 63.96282],
                [68.16797, 70.83203],
                [75.03719, 77.701256],
              ],
              [
                [81.29875, 83.96282],
                [88.16797, 90.83203],
                [95.03721, 97.70126],
              ],
              [
                [101.29875, 103.962814],
                [108.16797, 110.83203],
                [115.03721, 117.70127],
              ],
            ],
          ].flat(Infinity),
        );

        expectToBeCloseToArray(target, resized.data);
      });
    });
    describe("upscale", () => {
      const input = new Tensor(
        "float32",
        new Float32Array(2 * 3 * 3 * 2).map((_, i) => i),
        [2, 3, 3, 2],
      );

      const size = [2, 3, 4, 5];
      it("bilinear", async () => {
        const resized = await interpolate_4d(input, { mode: "bilinear", size });
        const target = new Float32Array(
          [
            [
              [
                [0.0, 0.1, 0.5, 0.9, 1.0],
                [1.25, 1.35, 1.75, 2.15, 2.25],
                [2.75, 2.85, 3.25, 3.65, 3.75],
                [4.0, 4.1, 4.5, 4.9, 5.0],
              ],
              [
                [6.0, 6.1, 6.5, 6.9, 7.0],
                [7.25, 7.35, 7.75, 8.15, 8.25],
                [8.75, 8.85, 9.25, 9.65, 9.75],
                [10.0, 10.1, 10.5, 10.9, 11.0],
              ],
              [
                [12.0, 12.1, 12.5, 12.9, 13.0],
                [13.25, 13.35, 13.75, 14.15, 14.25],
                [14.75, 14.85, 15.25, 15.65, 15.75],
                [16.0, 16.1, 16.5, 16.9, 17.0],
              ],
            ],
            [
              [
                [18.0, 18.1, 18.5, 18.9, 19.0],
                [19.25, 19.35, 19.75, 20.15, 20.25],
                [20.75, 20.85, 21.25, 21.65, 21.75],
                [22.0, 22.1, 22.5, 22.9, 23.0],
              ],
              [
                [24.0, 24.1, 24.5, 24.9, 25.0],
                [25.25, 25.35, 25.75, 26.15, 26.25],
                [26.75, 26.85, 27.25, 27.65, 27.75],
                [28.0, 28.1, 28.5, 28.9, 29.0],
              ],
              [
                [30.0, 30.1, 30.5, 30.9, 31.0],
                [31.25, 31.35, 31.75, 32.15, 32.25],
                [32.75, 32.85, 33.25, 33.65, 33.75],
                [34.0, 34.1, 34.5, 34.9, 35.0],
              ],
            ],
          ].flat(Infinity),
        );

        expectToBeCloseToArray(target, resized.data);
      });

      it("bicubic", async () => {
        const resized = await interpolate_4d(input, { mode: "bicubic", size });

        const target = new Float32Array(
          [
            [
              [
                [-0.253804475069046, -0.06155451014637947, 0.3564453125, 0.7744455337524414, 0.9666945934295654],
                [0.9493208527565002, 1.1415706872940063, 1.5595703125, 1.977570652961731, 2.1698191165924072],
                [2.8301806449890137, 3.022430181503296, 3.4404296875, 3.8584301471710205, 4.050677299499512],
                [4.033306121826172, 4.225555419921875, 4.6435546875, 5.061554908752441, 5.253802299499512],
              ],
              [
                [5.746196269989014, 5.938446998596191, 6.3564453125, 6.774445533752441, 6.966691493988037],
                [6.94932222366333, 7.14157247543335, 7.5595703125, 7.977570056915283, 8.169816970825195],
                [8.830181121826172, 9.022432327270508, 9.4404296875, 9.858429908752441, 10.050675392150879],
                [10.033307075500488, 10.225557327270508, 10.6435546875, 11.061556816101074, 11.253799438476562],
              ],
              [
                [11.746198654174805, 11.938446998596191, 12.3564453125, 12.774446487426758, 12.966689109802246],
                [12.949322700500488, 13.141572952270508, 13.5595703125, 13.977571487426758, 14.16981315612793],
                [14.830183029174805, 15.022432327270508, 15.4404296875, 15.858430862426758, 16.05067253112793],
                [16.033309936523438, 16.225557327270508, 16.6435546875, 17.061555862426758, 17.25379753112793],
              ],
            ],
            [
              [
                [17.746200561523438, 17.938447952270508, 18.3564453125, 18.774446487426758, 18.966686248779297],
                [18.949325561523438, 19.14157485961914, 19.5595703125, 19.977571487426758, 20.169809341430664],
                [20.830184936523438, 21.02243423461914, 21.4404296875, 21.858430862426758, 22.050668716430664],
                [22.03331184387207, 22.225557327270508, 22.6435546875, 23.061555862426758, 23.25379180908203],
              ],
              [
                [23.746200561523438, 23.93844985961914, 24.3564453125, 24.77444839477539, 24.96668243408203],
                [24.949325561523438, 25.141576766967773, 25.5595703125, 25.977571487426758, 26.1698055267334],
                [26.830184936523438, 27.022436141967773, 27.4404296875, 27.858430862426758, 28.05066680908203],
                [28.033313751220703, 28.225557327270508, 28.6435546875, 29.061555862426758, 29.25379180908203],
              ],
              [
                [29.74620246887207, 29.93844985961914, 30.3564453125, 30.77444839477539, 30.96668243408203],
                [30.949325561523438, 31.141578674316406, 31.5595703125, 31.977571487426758, 32.16980743408203],
                [32.8301887512207, 33.022438049316406, 33.4404296875, 33.858428955078125, 34.050662994384766],
                [34.03330993652344, 34.22556686401367, 34.6435546875, 35.06155014038086, 35.253787994384766],
              ],
            ],
          ].flat(Infinity),
        );

        expectToBeCloseToArray(target, resized.data);
      });
    });
  });

  describe("matmul", () => {
    it("(2, 5) @ (5, 4) -> (2, 4)", async () => {
      const a = new Tensor("float32", range(10), [2, 5]);
      const b = new Tensor("float32", range(20), [5, 4]);

      const result = await matmul(a, b);

      const target = new Float32Array(
        [
          [120.0, 130.0, 140.0, 150.0],
          [320.0, 355.0, 390.0, 425.0],
        ].flat(),
      );

      expectToBeCloseToArray(target, result.data);
    });
  });

  describe("rfft", () => {
    it("non-power of 2", async () => {
      const rows = 2;
      const cols = 3;
      const input = new Tensor("float32", range(rows * cols), [rows, cols]);
      const dim = new Tensor("int64", [-1n], []);
      const result = await rfft(input, dim);

      const target = new Float32Array(
        [
          [
            [3, 0],
            [-1.5, 0.8660262823104858],
          ],
          [
            [12, 0],
            [-1.5, 0.866027295589447],
          ],
        ].flat(Infinity),
      );

      expectToBeCloseToArray(target, result.data);
    });

    it("power of 2", async () => {
      const rows = 2;
      const cols = 4;
      const input = new Tensor("float32", range(rows * cols), [rows, cols]);
      const dim = new Tensor("int64", [-1n], []);

      const result = await rfft(input, dim);
      const target = new Float32Array(
        [
          [
            [6, 0],
            [-2, 2],
            [-2, 0],
          ],
          [
            [22, 0],
            [-2, 2],
            [-2, 0],
          ],
        ].flat(Infinity),
      );

      expectToBeCloseToArray(target, result.data);
    });
  });

  describe("slice", () => {
    it("should slice", async () => {
      const input = new Tensor("float32", [1, 2, 3, 4, 5, 6, 7, 8, 9], [3, 3]);

      const target = new Float32Array(
        [
          [1, 2],
          [4, 5],
        ].flat(Infinity),
      );

      const starts = [0, 0];
      const ends = [2, 2];
      const axes = [0, 1];
      const steps = [1, 1];

      const result = await slice(input, starts, ends, axes, steps);
      expectToBeCloseToArray(target, result.data);
    });
  });
});
