import assert from "node:assert/strict";
import { createHash } from "node:crypto";
import { test } from "node:test";
import {
  ablationCellLimit,
  ablationCellMixture,
  assignEntries,
  authorShares,
  authorVsMerged,
  bitWidth,
  cellIndex,
  coverage,
  decisionSeed,
  disagreement,
  forwardBatchSize,
  forwardMixtures,
  personKnots,
  placement,
  runBalance,
  seededRandom,
  seededSample,
  surprise,
  trackRecord,
  worldBits,
  type Mixture,
  type Placement,
  type RunState,
} from "../src/rules/index.ts";
import { canonical } from "../src/vocab/index.ts";
import { close, closeArray, loadFixture, range, scenarioPicker } from "./rules-fixtures-helpers.ts";

// info.json（方法.md 6.9、6.10 第 4 条，约定 31–34、37）与 misc.json（4.2、第七部分、约定 35–40），
// 对照独立 Python 参考实现算出的预期值。

const info = loadFixture("info.json");
const misc = loadFixture("misc.json");
const pickInfo = scenarioPicker(info);
const pickMisc = scenarioPicker(misc);
/** 世界比特、分歧量：info.json tolerance.bits（理由见 tolerance.note）。 */
const BITS_TOL: number = info.tolerance.bits;
/** 事前把握、覆盖率等普通数值：info.json tolerance.abs。惊讶另按 tolerance.surprise 的公式（见 checkSurprise）。 */
const TOL: number = info.tolerance.abs;

/** Python 用的落点、状态名 → rules 的代号。 */
/** 正向一个候选的各人范围：[{researcher, range}]。 */
const rangesOf = (rs: any[]) => rs.map((r) => ({ researcher: r.researcher as string, range: range(r.range) }));

const PLACEMENT: Record<string, Placement> = {
  below10: "belowQ1",
  from10to50: "q1ToQ5",
  from50to90: "q5ToQ9",
  above90: "aboveQ9",
};

function checkSurprise(m: Mixture, b: number, width: number, v: number, e: any, what: string): void {
  const got = surprise(m, b, width, v);
  assert.equal(got.cell, e.bin, `${what}：格号`);
  close(got.prior, e.p, TOL, `${what}：事前把握`);
  if (e.status === "ok") {
    assert.equal(got.bits.status, "measured", `${what}：状态`);
    // 惊讶 = −log₂p；p 是两个浮点 F 值之差，带几个 2⁻⁵³ 的绝对误差，除以 p 后放大（info.json tolerance.surprise）。
    close((got.bits as { value: number }).value, e.surprise, TOL + 2 ** -50 / (e.p * Math.LN2), `${what}：惊讶`);
  } else {
    assert.equal(e.status, "priorZero");
    assert.equal(got.bits.status, "unmeasured", `${what}：事前把握为 0 时是明确状态`);
  }
}

test("info.json bits_forward_r1_real", () => {
  const s = pickInfo.get("bits_forward_r1_real");
  const width = bitWidth(s.params.eps);
  close(width, s.params.delta, 0, "格宽 Δ");
  const ids = Object.keys(s.input.candidates);
  const raws = forwardMixtures(
    ids.map((id) => ({ ranges: rangesOf(s.input.candidates[id]) })),
    s.params.lambda,
  );
  ids.forEach((id, i) => close(worldBits(raws[i]!, s.input.b, width), s.expected.candidates[id].worldBits, BITS_TOL, `${id} 世界比特`));
  const c3 = ids.indexOf("r1-v1-c3");
  checkSurprise(raws[c3]!, s.input.b, width, s.input.actual["r1-v1-c3"], s.expected.candidates["r1-v1-c3"].surprise, "c3 惊讶");
  checkSurprise(raws[c3]!, s.input.b, width, s.input.farValue, s.expected.surpriseFarValue, "预留之外的实测");
  // 分歧量：正向一人一份范围，作者各不相同。
  const people = rangesOf(s.input.candidates["r1-v1-c3"]);
  const ranges = people.map((p) => p.range);
  const entries = people.map((p) => ({ id: p.researcher, author: p.researcher, range: p.range }));
  const e = s.expected.disagreementC3;
  close(disagreement(entries, s.input.b, width), e.disagreement, BITS_TOL, "c3 分歧量");
  // 分歧量的两项：等权平均后的比特、各人的比特（用 worldBits 在不掺预留的混合上算）。
  const avg: Mixture = ranges.map((r: any) => ({ weight: 1 / ranges.length, knots: personKnots(r) }));
  close(worldBits(avg, s.input.b, width), e.avgBits, BITS_TOL, "c3 平均后的比特");
  ranges.forEach((r: any, i: number) =>
    close(worldBits([{ weight: 1, knots: personKnots(r) }], s.input.b, width), e.eachBits[i], BITS_TOL, `c3 第 ${i} 人比特`),
  );
});

test("info.json bits_ablation_cell_r1_p1", () => {
  const s = pickInfo.get("bits_ablation_cell_r1_p1");
  const width = bitWidth(s.params.eps);
  const entries = s.input.explanations.map((e: any) => ({ id: e.id as string, author: e.author as string, range: range(e.range) }));
  const m = ablationCellMixture(entries, s.params.lambda);
  close(worldBits(m, s.input.b, width), s.expected.worldBits, BITS_TOL, "世界比特");
  checkSurprise(m, s.input.b, width, s.input.actual, s.expected.surprise, "惊讶");
  // 平均后的比特：研究者等权、同人几条平分，不掺预留（两边读法相同）。
  const shares = authorShares(entries.map((e: any) => e.author));
  const avg: Mixture = entries.map((e: any, i: number) => ({ weight: shares[i]!, knots: personKnots(e.range) }));
  close(worldBits(avg, s.input.b, width), s.expected.avgBits, BITS_TOL, "平均后的比特");
  // 各人的比特：每人几条解释等权平均成"这个人的折线"（6.9，只量人与人之间的差别）。
  const authors = [...new Set(entries.map((e: any) => e.author as string))];
  authors.forEach((a, i) => {
    const own = entries.filter((e: any) => e.author === a);
    const person: Mixture = own.map((e: any) => ({ weight: 1 / own.length, knots: personKnots(e.range) }));
    close(worldBits(person, s.input.b, width), s.expected.eachPersonBits[i], BITS_TOL, `${a} 的比特`);
  });
  close(disagreement(entries, s.input.b, width), s.expected.disagreement, BITS_TOL, "分歧量");
});

test("info.json bits_minimize_bin_edges", () => {
  const s = pickInfo.get("bits_minimize_bin_edges");
  const width = bitWidth(s.params.eps);
  const ids = Object.keys(s.input.candidates);
  const raws = forwardMixtures(
    ids.map((id) => ({ ranges: rangesOf(s.input.candidates[id]) })),
    s.params.lambda,
  );
  const m = raws[ids.indexOf(s.input.candidate)]!;
  close(worldBits(m, s.input.b, width), s.expected.worldBits, BITS_TOL, "世界比特");
  assert.equal(cellIndex(s.input.b, width, s.expected.edgeMinus1), -1, "正好在 b − Δ 边界上的值归第 −1 格");
  s.input.values.forEach((v: number, i: number) => checkSurprise(m, s.input.b, width, v, s.expected.surprise[i], `实测 ${v}`));
});

test("info.json landing_segments_endpoints", () => {
  const s = pickInfo.get("landing_segments_endpoints");
  const samples = s.input.cases.map((c: any) => ({ range: range(c.range), value: c.value as number }));
  assert.deepEqual(
    samples.map((c: any) => placement(c.range, c.value)),
    s.expected.segments.map((x: string) => PLACEMENT[x]),
  );
  const cov = coverage(samples);
  assert.equal(cov.rate.status, "measured");
  close((cov.rate as { value: number }).value, s.expected.coverage, TOL, "覆盖率");
});

test("info.json coverage_no_samples", () => {
  const s = pickInfo.get("coverage_no_samples");
  assert.equal(s.expected.coverage.status, "noSamples");
  assert.equal(coverage(s.input.cases).rate.status, "unmeasured", "没有样本时是明确状态");
});

test("info.json track_record_counts", () => {
  const s = pickInfo.get("track_record_counts");
  const samples = (["forward", "ablation"] as const).flatMap((phase) =>
    s.input[phase].map((x: any) => ({ researcher: "r", phase, range: range(x.range), value: x.value as number })),
  );
  const got = trackRecord(samples);
  assert.deepEqual(got.map((g) => g.phase), ["forward", "ablation"]);
  for (const g of got) {
    const e = s.expected[g.phase];
    assert.deepEqual(
      {
        medianHigh: g.tally.medianHigh,
        medianLow: g.tally.medianLow,
        pierceUpper: g.tally.aboveTop,
        pierceLower: g.tally.belowBottom,
        covered: g.tally.covered,
      },
      e,
      `${g.phase} 战绩`,
    );
  }
});

test("info.json author_vs_merged", () => {
  const s = pickInfo.get("author_vs_merged");
  // 只收跑过、有有效成绩的候选（由调用方筛）。
  const items = s.input.items
    .filter((i: any) => i.ran && i.value !== null)
    .map((i: any) => ({
      candidate: i.candidate,
      author: "author",
      authorRange: range(i.author),
      mergedRange: range(i.merged),
      value: i.value,
    }));
  const got = authorVsMerged(items, s.params.direction);
  assert.equal(got.length, s.expected.rows.length);
  got.forEach((g, i) => {
    const e = s.expected.rows[i];
    assert.equal(g.candidate, e.candidate);
    close(g.optimism, e.q5Diff, TOL, `${g.candidate} 作者比合并乐观多少`);
    assert.equal(g.authorPlacement, PLACEMENT[e.authorSegment]);
    assert.equal(g.mergedPlacement, PLACEMENT[e.mergedSegment]);
  });
});

for (const name of [
  "entries_rule1_unique",
  "entries_rule2_author_wins",
  "entries_rule2_smallest_index",
  "entries_identical_prefs",
  "entries_rule3_exhausted",
  "entries_N1",
  "entries_natural_rule2",
  "entries_natural_rule3",
]) {
  test(`misc.json ${name}`, () => {
    const s = pickMisc.get(name);
    const got = assignEntries({
      researchers: s.input.researchers.map((id: string) => ({ id, choices: s.input.preferences[id] })),
      entries: s.input.entries,
    });
    assert.deepEqual(
      Object.fromEntries(got.map((a) => [a.researcher, { entry: a.entry, choice: a.rank }])),
      s.expected.assignments,
    );
    // 结果按研究者编号（自然序）排，与 Python 的输出次序一致。
    assert.deepEqual(got.map((a) => a.researcher), Object.keys(s.expected.assignments));
  });
}

function balanceRow(total: number, runs: RunState[], batchCap: number) {
  const b = runBalance(total, runs);
  return {
    charged: b.charged,
    occupied: b.charged + b.reserved,
    balance: b.balance,
    forwardK: forwardBatchSize(batchCap, b.balance),
    ablationMax: ablationCellLimit(b.balance),
  };
}

test("misc.json balance_part7_table", () => {
  const s = pickMisc.get("balance_part7_table");
  let runs: RunState[] = [];
  s.input.steps.forEach((st: any, i: number) => {
    runs = [...runs, ...st.added];
    const { step, added, ...e } = s.expected.afterEachStep[i];
    assert.equal(step, st.step);
    assert.deepEqual(balanceRow(s.params.totalRuns, runs, s.params.batchCap), e, `第 ${i} 步 ${step}`);
  });
});

for (const name of ["balance_mixed_states", "balance_0", "balance_1", "balance_2"]) {
  test(`misc.json ${name}`, () => {
    const s = pickMisc.get(name);
    assert.deepEqual(balanceRow(s.params.totalRuns, s.input.runs, s.params.batchCap), s.expected);
  });
}

for (const name of ["seed_study-a_r1", "seed_15-real-1_r3"]) {
  test(`misc.json ${name}`, () => {
    const s = pickMisc.get(name);
    const { studyId, round, version } = s.input;
    assert.equal(canonical({ studyId, round, version }), s.expected.canonical, "规范 JSON");
    const seed = decisionSeed({ studyId, round, version });
    assert.equal(seed, s.expected.seed, "种子");
    // 第 counter 个数取 sha256(`${seed}:${counter}`) 的前 6 字节（约定 40），这里只核对两边对这个式子的理解相同。
    assert.deepEqual(
      [0, 1, 2].map((c) => createHash("sha256").update(`${seed}:${c}`).digest().subarray(0, 6).toString("hex")),
      s.expected.first6BytesHex,
    );
    const rng = seededRandom(seed);
    assert.deepEqual([0, 1, 2, 3, 4].map(() => rng.uniformInt(10)), s.expected.uniformInt10First5);
    assert.deepEqual(seededSample(s.input.candidates, s.input.batchSize, seed), s.expected.batchSample, "抽候选");
    assert.deepEqual(seededSample(s.input.cells, s.input.cellCount, seed), s.expected.cellSample, "抽格");
    // 十二个编号：候选和因子都按自然序（c2 < c10、p2 < p10）排好再抽。
    assert.deepEqual(seededSample(s.input.candidatesMany, s.input.manyCount, seed), s.expected.batchSampleMany, "抽候选（12 个）");
    assert.deepEqual(
      seededSample(s.input.cellsMany, s.input.manyCount, seed),
      s.expected.cellSampleMany,
      "抽格（12 个）",
    );
  });
}

test("info.json、misc.json 每个场景都已覆盖", () => {
  assert.deepEqual(pickInfo.unused(), []);
  assert.deepEqual(pickMisc.unused(), []);
});