import { describe, it, expect } from "vitest";
import {
  orientationFromAffine,
  isCanonicalOrientation,
  applyCanonicalOrientation,
} from "./niftiLoader";

// Real affines measured from the two dev-stack files (via nibabel) — see the KiTS
// orientation investigation. ct.nii.gz (TotalSegmentator) is RAS; KIT23_00002_0000.nii.gz
// (RibFrac dataset) is ('I','P','L').
const RAS_AFFINE = [
  [1.5, 0, 0, 0],
  [0, 1.5, 0, 0],
  [0, 0, 1.5, 0],
  [0, 0, 0, 1],
];
const KIT23_IPL_AFFINE = [
  [0, 0, -0.9395, 0],
  [0, -0.9395, 0, 0],
  [-1.0, 0, 0, 0],
  [0, 0, 0, 1],
];

describe("orientationFromAffine", () => {
  it("treats a clean RAS (diagonal, positive) affine as canonical — no reorder/flip", () => {
    const o = orientationFromAffine(RAS_AFFINE)!;
    expect(o.perm).toEqual([0, 1, 2]);
    expect(o.flips).toEqual([false, false, false]);
    expect(isCanonicalOrientation(o)).toBe(true);
  });

  it("maps the KiTS ('I','P','L') affine to a full axis reversal + flip-all", () => {
    const o = orientationFromAffine(KIT23_IPL_AFFINE)!;
    expect(o.perm).toEqual([2, 1, 0]);
    expect(o.flips).toEqual([true, true, true]);
    expect(isCanonicalOrientation(o)).toBe(false);
  });

  it("handles a pure permutation (axes swapped, no flip)", () => {
    // array axis0→S, axis1→A, axis2→R, all positive
    const o = orientationFromAffine([
      [0, 0, 1, 0],
      [0, 1, 0, 0],
      [1, 0, 0, 0],
      [0, 0, 0, 1],
    ])!;
    expect(o.perm).toEqual([2, 1, 0]);
    expect(o.flips).toEqual([false, false, false]);
  });

  it("handles a pure flip (in-order axes, one reversed)", () => {
    const o = orientationFromAffine([
      [1, 0, 0, 0],
      [0, -1, 0, 0],
      [0, 0, 1, 0],
      [0, 0, 0, 1],
    ])!;
    expect(o.perm).toEqual([0, 1, 2]);
    expect(o.flips).toEqual([false, true, false]);
  });

  it("returns null for a missing / too-short / singular affine (safe fallback)", () => {
    expect(orientationFromAffine(null)).toBeNull();
    expect(orientationFromAffine([])).toBeNull();
    expect(orientationFromAffine([[1, 0], [0, 1]])).toBeNull(); // < 3 rows
    // singular: a zero column means an array axis maps nowhere
    expect(
      orientationFromAffine([
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 1],
      ]),
    ).toBeNull();
  });
});

describe("applyCanonicalOrientation", () => {
  // A 2x3x4 volume (axis0 fastest-varying) filled with its own flat index, so a remapped
  // voxel's value IS its original flat index — easy to assert the mapping.
  const dims: [number, number, number] = [2, 3, 4];
  const spacing: [number, number, number] = [1, 2, 3];
  const data = new Float32Array(2 * 3 * 4);
  for (let i = 0; i < data.length; i++) data[i] = i;

  it("reverses + flips a KiTS-style volume to canonical, swapping dims and spacing", () => {
    const o = { perm: [2, 1, 0] as [number, number, number], flips: [true, true, true] as [boolean, boolean, boolean] };
    const r = applyCanonicalOrientation(data, dims, spacing, o);
    expect(r.dims).toEqual([4, 3, 2]); // dims[perm]
    expect(r.spacing).toEqual([3, 2, 1]); // spacing[perm]
    expect(r.data.length).toBe(24);
    // Full reversal: canonical (0,0,0) ← source far corner; canonical last ← source (0,0,0).
    expect(r.data[0]).toBe(23);
    expect(r.data[23]).toBe(0);
    // Stepping the fastest output axis maps back along the (flipped) source axis 2.
    expect(r.data[1]).toBe(17);
  });

  it("is a value-preserving no-op for the canonical orientation", () => {
    const o = { perm: [0, 1, 2] as [number, number, number], flips: [false, false, false] as [boolean, boolean, boolean] };
    const r = applyCanonicalOrientation(data, dims, spacing, o);
    expect(r.dims).toEqual([2, 3, 4]);
    expect(r.spacing).toEqual([1, 2, 3]);
    expect(Array.from(r.data)).toEqual(Array.from(data));
  });
});
