import assert from "node:assert/strict";
import { test } from "node:test";
import {
  ablationModel,
  cellGrid,
  chooseBet,
  evaluateCombos,
  gridSize,
  informationValue,
  informationValueExact,
  informationValueGrid,
  measuredGrid,
  pickByNet,
  remainingSpace,
  runCellRule,
  selectAll,
  selectByValue,
  selectNone,
  selectRandomSameCount,
  type AblationExplanation,
  type CellRuleInput,
} from "../src/rules/index.ts";

// 预期值来源写在每条测试旁边。手算用的模型（下称"模型甲"）：
//   基线 0、完整 1、越大越好，G = 1，撑几成 s = 1 − 分数；没有已知最好，h = 1；λ = 0（开放预留权重 0）。
//   一条解释：p1 交很窄的 (0.49, 0.5, 0.51)，s₁ 的折线是 (.48,0)(.49,.1)(.5,.5)(.51,.9)(.52,1)；
//   p2 交 (0.25, 0.5, 0.75)，s₂ 的折线是 (0,0)(.25,.1)(.5,.5)(.75,.9)(1,1)。E s₁ = E s₂ = 0.5，ĝ₀ = 0.5。
//   VOI({p2}) = E[max(0.5, s₂)] − 0.5 = ∫_{0.5}^{1} (1 − F) = 0.25 × 0.3 + 0.25 × 0.05 = 0.0875。
//   VOI({p1}) = E[max(s₁, 0.5)] − 0.5 = 0.4 × 0.005 + 0.1 × 0.015 = 0.0035（[.5,.51] 上把握 0.4、平均超出 0.005；
//   [.51,.52] 上把握 0.1、平均超出 0.015）。
//   VOI({p1, p2}) = E[max(s₁, s₂)] − 0.5 = 65669/750000 = 0.0875586666…（tests/fixtures/rules/rules_ref.py 的
//   AblationModel.voi_full_exact，有理数精确算）。

function close(actual: number, expected: number, tol = 1e-12): void {
  assert.ok(Math.abs(actual - expected) <= tol, `实际 ${actual}，预期 ${expected}`);
}

function explanation(id: string, author: string, cells: Record<string, [number, number, number]>): AblationExplanation {
  return {
    id,
    author,
    cells: Object.entries(cells).map(([factor, [q1, q5, q9]]) => ({ factor, range: { q1, q5, q9 } })),
  };
}

const modelA = {
  baseline: 0,
  full: 1,
  direction: "maximize" as const,
  knownBest: null,
  epsilon: 0.01,
  lambda: 0,
  factors: ["p2", "p1"],
  explanations: [explanation("e1", "A", { p1: [0.49, 0.5, 0.51], p2: [0.25, 0.5, 0.75] })],
};
const VOI_P1 = 0.0035;
const VOI_BOTH = 65669 / 750000;

function ruleInput(over: Partial<CellRuleInput>): CellRuleInput {
  return { ...modelA, balance: 21, batchCap: 2, u: 1, decay: 1, seed: "seed", ...over };
}

test("6.7 剩余空间 h：有已知最好 max(K* − f(完整), ε/100)，没有时 1", () => {
  close(remainingSpace(2.6, 2.635983, 0.01, "maximize"), 0.035983);
  assert.equal(remainingSpace(2.64, 2.635983, 0.01, "maximize"), 0.0001);
  assert.equal(remainingSpace(2.6, null, 0.01, "maximize"), 1);
  close(remainingSpace(1.2, 1, 0.01, "minimize"), 0.2);
});

test("6.7 起始权重：解释合起来 1 − λ，研究者等权、同一人平分；开放预留 λ", () => {
  // 手算：A 两条、B 一条，λ = 0.1：A 各 0.9 × 0.25 = 0.225，B 0.9 × 0.5 = 0.45，预留 0.1。
  const m = ablationModel({
    ...modelA,
    lambda: 0.1,
    explanations: [
      explanation("a2", "A", { p1: [0.4, 0.5, 0.6], p2: [0.4, 0.5, 0.6] }),
      explanation("b1", "B", { p1: [0.4, 0.5, 0.6], p2: [0.4, 0.5, 0.6] }),
      explanation("a1", "A", { p1: [0.4, 0.5, 0.6], p2: [0.4, 0.5, 0.6] }),
    ],
  });
  assert.deepEqual(m.components, ["a1", "a2", "b1", "openReserve"]);
  assert.deepEqual(m.weights, [0.9 * 0.25, 0.9 * 0.25, 0.9 * 0.5, 0.1]);
});

test("6.7 部分格网格的每格小段数 m：m^k ≤ 4,000,000 的最大整数", () => {
  // 手算：158³ = 3,944,312 ≤ 4e6 < 159³ = 4,019,679；44⁴ = 3,748,096 ≤ 4e6 < 45⁴ = 4,100,625。
  assert.deepEqual([1, 2, 3, 4].map(gridSize), [4_000_000, 2000, 158, 44]);
});

test("6.7 网格：每个成分在一格上的小段把握加起来是 1", () => {
  // 手算：一个成分 [0, 1] 均匀，一个成分 [0.5, 1] 均匀。折点 0、0.5、1 两段，k = 1 时 s = 20，共 40 个小段，宽 0.025。
  // 第一个成分每个小段 0.025；第二个成分前 20 段 0，后 20 段各 0.05。
  const g = cellGrid(
    [
      [
        [0, 0],
        [1, 1],
      ],
      [
        [0.5, 0],
        [1, 1],
      ],
    ],
    1,
  );
  assert.equal(g.mids.length, 40);
  close(g.mids[0]!, 0.0125);
  close(g.mass[0]!.reduce((s, x) => s + x, 0), 1);
  close(g.mass[1]!.reduce((s, x) => s + x, 0), 1);
  close(g.mass[0]![3]!, 0.025);
  assert.equal(g.mass[1]![19], 0);
  close(g.mass[1]![20]!, 0.05);
});

test("6.7 段数超过 m 时整格等宽切 m 份", () => {
  // 手算：k = 4 时 m = 44；45 个折点有 44 段 → s = 1；46 个折点有 45 段 > 44 → 整格等宽切 44 份。
  const line = (n: number): [number, number][] => Array.from({ length: n }, (_, i) => [i, i / (n - 1)]);
  assert.equal(cellGrid([line(45)], 4).mids.length, 44);
  const g = cellGrid([line(46)], 4);
  assert.equal(g.mids.length, 44);
  close(g.mids[0]!, 45 / 44 / 2);
});

test("6.7 信息价值：模型甲的手算值；全测的精确算法与网格互相对照", () => {
  const m = ablationModel(modelA);
  assert.deepEqual(m.factors, ["p1", "p2"]);
  close(m.prior[0]!, 0.5);
  close(m.prior[1]!, 0.5);
  close(informationValue(m, [1]), 0.0875);
  // 只测 p1：s₁ 每个小段内均匀，0.5 正好是折点，max(s, 0.5) 在每个小段上是一次函数，中点法在这里是精确的。
  close(informationValue(m, [0]), VOI_P1);
  close(informationValue(m, [0, 1]), VOI_BOTH);
  close(informationValueExact(m), VOI_BOTH);
  // 全测时网格是近似，按 ablation.json tolerance.gridVsExact 的 k ≤ 3 界（5e-4 相对）对照。
  close(informationValueGrid(m, [0, 1]), VOI_BOTH, 5e-4 * VOI_BOTH);
  assert.equal(informationValue(m, []), 0);
});

test("6.7 价值、成本、净值：u、m′、R̂、ĝ₀ 按原文", () => {
  // 手算：模型甲，余额 21、每批上限 K = 2、u = 1、r = 1。
  //   一格：R̂ = ⌈20/2⌉ = 10，m′ = 10；成本 = 2 × 0.5 × (3/4 − 2/3) = 1/12。
  //   两格：R̂ = ⌈19/2⌉ = 10，m′ = 10；成本 = 2 × 0.5 × (4/5 − 2/3) = 2/15。
  //   [p1]：价值 10 × 0.0035 = 0.035，净值 0.035 − 1/12；[p2]：价值 0.875，净值 0.875 − 1/12；
  //   [p1,p2]：价值 10 × VOI_BOTH，净值 10 × VOI_BOTH − 2/15。
  const { g0, combos } = evaluateCombos(ablationModel(modelA), { balance: 21, batchCap: 2, u: 1, decay: 1 });
  close(g0, 0.5);
  assert.deepEqual(
    combos.map((c) => [c.factors.join(","), c.rHat, c.mPrime]),
    [
      ["p1", 10, 10],
      ["p2", 10, 10],
      ["p1,p2", 10, 10],
    ],
  );
  close(combos[0]!.net, 10 * VOI_P1 - 1 / 12);
  close(combos[1]!.value, 0.875);
  close(combos[1]!.net, 0.875 - 1 / 12);
  close(combos[2]!.cost, 2 / 15);
  close(combos[2]!.net, 10 * VOI_BOTH - 2 / 15);
  // u = 0.5、r = 0.5、余额 5：一格 R̂ = ⌈4/2⌉ = 2，m′ = 1 + 0.5 = 1.5，价值 0.5 × 1.5 × 0.0875 = 0.065625。
  const half = evaluateCombos(ablationModel(modelA), { balance: 5, batchCap: 2, u: 0.5, decay: 0.5 });
  close(half.combos[1]!.mPrime, 1.5);
  close(half.combos[1]!.value, 0.065625);
});

test("6.7 选格（信息价值减成本）：选净值最大的；没选中的附原因和含它的最好组合", () => {
  const d = selectByValue(ruleInput({}));
  assert.deepEqual(d.chosen, ["p2"]);
  close(d.chosenNet, 0.875 - 1 / 12);
  assert.equal(d.cells[0]!.factor, "p1");
  assert.equal(d.cells[0]!.tested, false);
  assert.equal(d.cells[0]!.reason, "value");
  // 含 p1 的组合：[p1] 0.035 − 1/12 < 0，[p1,p2] 10 × VOI_BOTH − 2/15 ≈ 0.742，最好的是后者。
  assert.deepEqual(d.cells[0]!.bestCombo?.factors, ["p1", "p2"]);
  close(d.cells[0]!.bestCombo!.net, 10 * VOI_BOTH - 2 / 15);
  assert.deepEqual(d.cells[1], { factor: "p2", tested: true, reason: null, bestCombo: null });
});

test("6.7 选格：余额 0、1 时不做消融，原因是剩余次数不够、没有可选组合；余额 2 时只列一格的组合", () => {
  for (const balance of [0, 1]) {
    const d = selectByValue(ruleInput({ balance }));
    assert.deepEqual(d.combos, []);
    assert.deepEqual(d.chosen, []);
    assert.equal(d.chosenNet, 0);
    for (const c of d.cells) assert.deepEqual([c.tested, c.reason, c.bestCombo], [false, "budget", null]);
  }
  // 余额 2：只能测 1 格。u = 1、r = 1：R̂ = ⌈1/2⌉ = 1，m′ = 1，[p2] 价值 0.0875 − 成本 1/12 = 0.00416…> 0，选 p2；
  // p1 再加进来就是 2 格 > 余额 − 1 = 1，记"剩余次数不够"。
  const two = selectByValue(ruleInput({ balance: 2 }));
  assert.deepEqual(
    two.combos.map((c) => c.factors.join(",")),
    ["p1", "p2"],
  );
  assert.deepEqual(two.chosen, ["p2"]);
  close(two.chosenNet, 0.0875 - 1 / 12);
  assert.equal(two.cells[0]!.reason, "budget");
});

test("6.7 选格的净值容差 1e-12：不大于 1e-12 不选；相差不到 1e-12 时留先列出的（格数少、编号小）", () => {
  assert.deepEqual(pickByNet([{ factors: ["p1"], net: 1e-12 }]), []);
  assert.deepEqual(pickByNet([{ factors: ["p1"], net: 0 }]), []);
  assert.deepEqual(pickByNet([{ factors: ["p1"], net: 2e-12 }]), ["p1"]);
  assert.deepEqual(
    pickByNet([
      { factors: ["p1"], net: 0.5 },
      { factors: ["p2"], net: 0.5 + 5e-13 },
      { factors: ["p1", "p2"], net: 0.5 + 9e-13 },
    ]),
    ["p1"],
  );
  assert.deepEqual(
    pickByNet([
      { factors: ["p1"], net: 0.5 },
      { factors: ["p1", "p2"], net: 0.5 + 2e-12 },
    ]),
    ["p1", "p2"],
  );
});

test("12 局限第 7 条：大家都押同一个因子、没有不确定时信息价值为 0，一格都不测", () => {
  // 手算：p1 交 (0.24, 0.25, 0.26)，s₁ 在 [0.73, 0.77]；p2 交 (0.49, 0.5, 0.51)，s₂ 在 [0.48, 0.52]。
  // 两格的范围不重叠，不论测到什么都押 p1，VOI 为 0；成本 > 0，净值都 < 0。
  const d = selectByValue(
    ruleInput({ explanations: [explanation("e1", "A", { p1: [0.24, 0.25, 0.26], p2: [0.49, 0.5, 0.51] })] }),
  );
  assert.deepEqual(d.chosen, []);
  for (const c of d.combos) close(c.voi, 0);
  assert.deepEqual(
    d.cells.map((c) => c.reason),
    ["value", "value"],
  );
});

test("登记表：四条选格规则", () => {
  const three = ruleInput({
    factors: ["p1", "p2", "p3"],
    balance: 3,
    explanations: [explanation("e1", "A", { p1: [0.49, 0.5, 0.51], p2: [0.25, 0.5, 0.75], p3: [0.25, 0.5, 0.75] })],
  });
  // 全测：余额 3，最多 2 格，按编号测 p1、p2，p3 剩余次数不够。
  const all = runCellRule("all", three);
  assert.deepEqual(all.rule, { name: "all", version: 1 });
  assert.deepEqual(all.chosen, ["p1", "p2"]);
  assert.deepEqual(
    all.cells.map((c) => c.reason),
    [null, null, "budget"],
  );
  assert.deepEqual(selectAll(three).chosen, ["p1", "p2"]);
  // 都不测。
  const none = runCellRule("none", three);
  assert.deepEqual(none.chosen, []);
  assert.deepEqual(
    none.cells.map((c) => c.reason),
    ["ruleNone", "ruleNone", "ruleNone"],
  );
  assert.deepEqual(selectNone(three).chosen, []);
  // 随机同格数：格数等于信息价值减成本选的格数；同一个种子得到同一组。
  const n = runCellRule("valueMinusCost", three).chosen.length;
  for (const seed of ["a", "b", "c"]) {
    const r1 = selectRandomSameCount({ ...three, seed });
    assert.equal(r1.chosen.length, n);
    assert.deepEqual(runCellRule("randomSameCount", { ...three, seed }).chosen, r1.chosen);
    for (const c of r1.cells) assert.equal(c.reason, c.tested ? null : "notDrawn");
  }
});

test("同样的输入算两次，选格结果逐位相同", () => {
  const i = ruleInput({
    lambda: 0.1,
    knownBest: 1.2,
    explanations: [
      explanation("e1", "A", { p1: [0.3, 0.5, 0.6], p2: [0.2, 0.4, 0.75] }),
      explanation("e2", "B", { p1: [0.6, 0.7, 0.9], p2: [0.1, 0.15, 0.3] }),
    ],
  });
  assert.deepEqual(selectByValue(i), selectByValue(i));
});

test("6.8 押的因子：有效格用实测，其余格用更新后的权重估", () => {
  // 手算：基线 0、完整 1、λ = 0。e1(A)：p1 (0.19, 0.2, 0.21)（对称，E s = 0.8），p2 (0.25, 0.5, 0.75)；
  // e2(B)：p1 (0.79, 0.8, 0.81)（E s = 0.2），p2 (2, 2.25, 2.5)。p2 实测 0.5：e1 在 0.5 处密度取 [0.5, 0.75) 段 0.4/0.25 = 1.6，
  // e2 为 0 → 新权重 e1 = 1。
  // 收益：p1 = 0.8（估），p2 = 1 − 0.5 = 0.5（实测）。押 p1，预测收益 0.8。
  const input = {
    ...modelA,
    explanations: [
      explanation("e1", "A", { p1: [0.19, 0.2, 0.21], p2: [0.25, 0.5, 0.75] }),
      explanation("e2", "B", { p1: [0.79, 0.8, 0.81], p2: [2, 2.25, 2.5] }),
    ],
    cells: [
      { factor: "p1", outcome: { kind: "untested" as const } },
      { factor: "p2", outcome: { kind: "valid" as const, score: 0.5 } },
    ],
  };
  const b = chooseBet(input);
  assert.deepEqual(b.weights, [1, 0, 0]);
  assert.equal(b.weightsKept, false);
  close(b.gains[0]!.gain, 0.8);
  assert.equal(b.gains[0]!.measured, false);
  close(b.gains[1]!.gain, 0.5);
  assert.equal(b.gains[1]!.measured, true);
  assert.equal(b.factor, "p1");
  assert.equal(b.decision, "bet");
  // 没有有效格时保持起始权重：p1 = 0.5 × 0.8 + 0.5 × 0.2 = 0.5。
  const prior = chooseBet({ ...input, cells: [] });
  assert.deepEqual(prior.weights, [0.5, 0.5, 0]);
  close(prior.gains[0]!.gain, 0.5);
  // "一起关"的分数不参与更新：p2 一起关得 0.5，权重仍是起始权重。
  const joint = chooseBet({ ...input, cells: [{ factor: "p2", outcome: { kind: "joint", off: ["p1", "p2"], score: 0.5 } }] });
  assert.deepEqual(joint.weights, [0.5, 0.5, 0]);
  assert.equal(joint.gains[1]!.measured, false);
});

test("6.8 权重无从更新时保持起始权重；撑几成 < 0 时不押并记关掉反而更好多少", () => {
  // 手算：p2 实测 5，落在两条解释和开放预留之外（λ = 0 时预留权重 0，密度连乘全为 0）→ 保持起始权重。
  // 撑 = 1 − 5 = −4 < 0：不押，关掉 p2 比完整好 5 − 1 = 4。
  const b = chooseBet({
    ...modelA,
    explanations: [
      explanation("e1", "A", { p1: [0.19, 0.2, 0.21], p2: [0.25, 0.5, 0.75] }),
      explanation("e2", "B", { p1: [0.79, 0.8, 0.81], p2: [2, 2.25, 2.5] }),
    ],
    cells: [{ factor: "p2", outcome: { kind: "valid", score: 5 } }],
  });
  assert.equal(b.weightsKept, true);
  assert.deepEqual(b.weights, [0.5, 0.5, 0]);
  assert.equal(b.decision, "offBetter");
  assert.deepEqual(b.offBetter, [{ factor: "p2", better: 4 }]);
  close(b.predictedGain, 0.5);
  assert.equal(b.factor, "p1");
});

test("6.8 挑出的因子估计收益不为正时不押；平局取编号小的", () => {
  // 手算：p1 (1.24, 1.25, 1.26)（E s = −0.25），p2 (1.115, 1.125, 1.135)（E s = −0.125）：押 p2，收益 −0.125，不为正。
  const neg = chooseBet({
    ...modelA,
    explanations: [explanation("e1", "A", { p1: [1.24, 1.25, 1.26], p2: [1.115, 1.125, 1.135] })],
    cells: [],
  });
  assert.equal(neg.factor, "p2");
  close(neg.predictedGain, -0.125);
  assert.equal(neg.decision, "notPositive");
  // 两格都交 (0.75, 1, 1.25)，关于完整 1 对称，收益正好是 0（这几个数在二进制里精确）：取 p1，不押。
  const zero = chooseBet({
    ...modelA,
    explanations: [explanation("e1", "A", { p1: [0.75, 1, 1.25], p2: [0.75, 1, 1.25] })],
    cells: [],
  });
  assert.equal(zero.predictedGain, 0);
  assert.equal(zero.factor, "p1");
  assert.equal(zero.decision, "notPositive");
  // 两格收益相等且为正：取 p1。
  const tie = chooseBet({
    ...modelA,
    explanations: [explanation("e1", "A", { p1: [0.49, 0.5, 0.51], p2: [0.49, 0.5, 0.51] })],
    cells: [],
  });
  assert.equal(tie.gains[0]!.gain, tie.gains[1]!.gain);
  assert.equal(tie.factor, "p1");
  assert.equal(tie.decision, "bet");
});

test("6.7 部分格网格在分数原值上取折点：刻度换算后才合并的两个折点仍算两段", () => {
  // 反例来自随机压测：第二条的低端尾巴 0.6819999999999999 − (0.6869999999999999 − 0.6819999999999999) = 0.6769999999999999，
  // 与第三条的 q1 = 0.677 只差一个末位，原值上是两个折点；换成撑几成 (2.6276819772628 − x) ÷ G 后变成同一个数。
  // 方法 6.7"每格在自己的开放预留范围内按折点切段"，开放预留和折点都在分数原值上定，所以按原值算：17 个折点、16 段。
  const model = ablationModel({
    baseline: 0.9597642169962064,
    full: 2.6276819772628,
    direction: "maximize",
    knownBest: null,
    epsilon: 0.01,
    lambda: 0.1,
    factors: ["p1"],
    explanations: [
      explanation("e1", "A", { p1: [1.1840000000000002, 1.5510000000000002, 2.0010000000000003] }),
      explanation("e2", "B", { p1: [0.6819999999999999, 0.6869999999999999, 1.168] }),
      explanation("e3", "C", { p1: [0.677, 1.099, 1.3519999999999999] }),
    ],
  });
  const distinct = (lines: ReadonlyArray<ReadonlyArray<readonly [number, number]>>) =>
    new Set(lines.flatMap((k) => k.map(([x]) => x))).size;
  assert.equal(distinct(model.rawKnots[0]!), 17, "原值上 17 个折点");
  assert.equal(distinct(model.supportKnots[0]!), 16, "撑几成刻度上只剩 16 个（这是反例的前提）");
  // k = 1：m = 4,000,000，s = 20，16 段 × 20 = 320 个小段（按撑几成刻度切会是 15 × 20 = 300）。
  assert.equal(measuredGrid(model, 0, 1).mids.length, 320);
  // k = 3：m = 158，s = ⌊158 ÷ 16⌋ = 9，144 个小段（按撑几成刻度切会是 ⌊158 ÷ 15⌋ = 10，150 个）。
  assert.equal(measuredGrid(model, 0, 3).mids.length, 144);
  // 小段中点换成了撑几成：第一个小段在原值最低处，撑几成最大。
  const g = measuredGrid(model, 0, 1);
  const raw = cellGrid(model.rawKnots[0]!, 1);
  g.mids.forEach((m, i) => close(m, (2.6276819772628 - raw.mids[i]!) / model.gain));
});