import { test } from 'node:test';
import { strict as assert } from 'node:assert';
import {
  buildPairs,
  pickAntecedents,
  groupClusters,
} from '../src/pairs.js';

test('buildPairs: empty', () => {
  const [pi, pj] = buildPairs(0);
  assert.equal(pi.length, 0);
  assert.equal(pj.length, 0);
});

test('buildPairs: single mention has only DUMMY pair', () => {
  const [pi, pj] = buildPairs(1);
  assert.deepEqual(Array.from(pi), [1n]);
  assert.deepEqual(Array.from(pj), [0n]);
});

test('buildPairs: triangular shape matches Python', () => {
  // M=3 → mention 1 has [0]; mention 2 has [0,1]; mention 3 has [0,1,2].
  // Total = 1 + 2 + 3 = 6.
  const [pi, pj] = buildPairs(3);
  assert.deepEqual(Array.from(pi), [1n, 2n, 2n, 3n, 3n, 3n]);
  assert.deepEqual(Array.from(pj), [0n, 0n, 1n, 0n, 1n, 2n]);
});

test('pickAntecedents: argmax respects per-mention slicing', () => {
  const [pi, pj] = buildPairs(3);
  // Set up scores so mention 1 → DUMMY, mention 2 → mention 1,
  // mention 3 → mention 1.
  const scores = Float32Array.from([
    9.0,  // m=1, j=0 (DUMMY) — best
    -1.0, // m=2, j=0 (DUMMY)
    9.0,  // m=2, j=1           — best
    -1.0, // m=3, j=0 (DUMMY)
    9.0,  // m=3, j=1           — best
    -1.0, // m=3, j=2
  ]);
  const out = pickAntecedents(3, pi, pj, scores);
  assert.deepEqual(
    out.map((x) => x.antecedent),
    [0, 1, 1],
  );
});

test('groupClusters: chained pointers form one cluster', () => {
  // m0 → DUMMY, m1 → m0, m2 → m1 → all in one cluster.
  const decisions = [
    { antecedent: 0 },
    { antecedent: 1 },
    { antecedent: 2 },
  ];
  const { cluster, clusters } = groupClusters(decisions);
  assert.deepEqual(cluster, [0, 0, 0]);
  assert.deepEqual(clusters, [[0, 1, 2]]);
});

test('groupClusters: singletons have cluster=-1', () => {
  // Two independent singletons + one chain.
  const decisions = [
    { antecedent: 0 }, // m0: DUMMY → singleton
    { antecedent: 0 }, // m1: DUMMY → singleton
    { antecedent: 1 }, // m2 → m0 → cluster
    { antecedent: 3 }, // m3 → m2 → cluster
  ];
  const { cluster, clusters } = groupClusters(decisions);
  assert.equal(cluster[1], -1);
  assert.deepEqual(clusters, [[0, 2, 3]]);
});

test('groupClusters: empty input', () => {
  const { cluster, clusters } = groupClusters([]);
  assert.deepEqual(cluster, []);
  assert.deepEqual(clusters, []);
});