import assert from "node:assert/strict";
import { readFileSync } from "node:fs";
import { test } from "node:test";
import {
  ablationModel,
  bitWidth,
  cellGrid,
  chooseBet,
  evaluateCombos,
  forwardMixtures,
  gridSize,
  informationValue,
  informationValueExact,
  informationValueGrid,
  selectBestExpected,
  supportOf,
  surprise,
  synergyLabel,
  totalGain,
  worldBits,
  type AblationExplanation,
  type AblationModelInput,
  type BatchInput,
} from "../src/rules/index.ts";
import type { Range } from "../src/vocab/index.ts";

// 用参考项目 RivalBits 的独立计算（numpy 脚本）当预期值，对照 src/rules。只收口径与 docs/方法.md 第六部分一致的数。
//
// 一、消融算例：tests/fixtures/calc15/expected.json，原样复制自 RivalBits tests/fixtures/calc15/expected.json，
//    由 RivalBits research-wiki/设计-15-计算/calc15.py 在 14-real-1 账本副本上算出（dump.py 写成全精度 JSON）。
//    口径核对（逐项读 calc15.py）：
//    - 格 0 = onlyB = 关掉网格起点 = p1；格 1 = onlyA = 关掉求解流程 = p2（calc15.py 第 128–129 行）。
//    - 成分权重 (1 − λ)/人数/同人条数，预留 λ（第 131 行）= 方法 6.7"成分与权重"。
//    - 每格自己的预留 [L − W/2, U + W/2]（第 122–125、142–145 行）= 方法 6.3"消融"。
//    - h = max(K* − f(完整), 0.0001)，0.0001 = ε/100（第 132 行）= 方法 6.7"方向收益"。
//    - gain(x) = h·(完整 − x)/G（第 133 行）：voi、g0、opp、hand.* 都已乘 h（收益刻度）；
//      priorSupport* 是 prior/h（dump.py），在撑几成刻度上。
//    - 网格：折点含预留两端，s = min(20, ⌊m_k/段数⌋)，s = 0 时整格切 m_k 份，小段取中点（第 147–157 行）= 方法 6.7"信息价值"。
//    - 全测：每成分各格把握连乘、逐段高斯-勒让德（k 点，对 k 次多项式精确）积分，再按权重相加（第 159–180 行）= 方法 6.7。
//    - 成本 opp(s, K) = 2ĝ₀[(K+s)/(K+s+1) − K/(K+1)]（第 218–220 行）= 方法 6.7"价值减成本"。
//    不用的部分：choose（旧规则：tight 时才计成本、R − r 固定剩余轮数）、simulate、fwd_after_testing（旧的份额规则）。
//    3 格例的第三格是构造的：每条解释取它对前两格三个数的逐项平均（calc15.py 第 249 行；RivalBits tests/controller-ablation.test.ts 第 96 行附近）。
//
// 二、正向算例：12-real-2 第 1、2 轮三位研究者的范围（RivalBits tests/info.test.ts ROUND1 第 64 行起、第 2 轮第 280 行附近，
//    与 research-wiki/设计-14-计算/calc14b.py 第 50–52、185–187 行相同），b = 0.9597642169962064。
//    - 世界比特、惊讶：RivalBits tests/fixtures/calc15/forward_bits.py。每个候选的合并分布 = 0.9 × 三人折线平均 + 0.1 × 预留，
//      预留由全部候选全部人的折线定（= 方法 6.3"正向"），从 b 起按 Δ = 0.0001 切格（在改进量上从 0 切，等价），F 是连续折线。
//    - 选批 V：calc14b.py 的 V = ∫_b^4.0（第 2 轮 ∫_b^2.8）[1 − ∏F] dt，40 万段梯形积分；积分上限高于预留上端（3.45、2.669），
//      低于 b 的部分不算 = 方法 6.4。
//    - 分数尺读数：calc14b 用全部三个候选取 max（第 64–66、194 行），方法用候选池。两轮三个候选都入池（下面断言），两者相同，可以用。
//    不用的部分：calc14b 的 bits@eps（格宽 ε，不是 ε/100）、"== split" 整段（旧的解释分格做法）。

interface Calc15 {
  input: {
    baseline: number;
    full: number;
    knownBest: number;
    offGrid: number;
    offSolverLegalized: number;
    explanations: Array<{ id: string; author: string; cells: [number, number, number][] }>;
  };
  G: number;
  h: number;
  supportOffGrid: number;
  supportOffSolver: number;
  priorSupport2: number[];
  priorSupport3: number[];
  voi2: Record<string, number>;
  voi3: Record<string, number>;
  g0: number;
  opp: Record<string, number>;
  hand: {
    reservePrior: number;
    reservePosterior: number;
    explanationsPosterior: number;
    gainOffGrid: number;
    solverGainPrior: number;
    solverGainPosterior: number;
  };
  gridCap: Record<string, number>;
  cellMiniSegments: { k1: number[]; k2: number[]; k3: number[] };
  labels: Record<string, string>;
  label2: string;
}

const X = JSON.parse(readFileSync(new URL("./fixtures/calc15/expected.json", import.meta.url), "utf8")) as Calc15;
const IN = X.input;

function close(actual: number, expected: number, tol: number, what: string): void {
  assert.ok(Math.abs(actual - expected) <= tol, `${what}：实际 ${actual}，预期 ${expected}，容差 ${tol}`);
}

const range = ([q1, q5, q9]: readonly number[]): Range => ({ q1: q1!, q5: q5!, q9: q9! });

/** calc15 的九条解释；n = 3 时补上构造的第三格（前两格三个数逐项平均，calc15.py 第 249 行）。 */
function explanations(n: 2 | 3): AblationExplanation[] {
  return IN.explanations.map((e) => {
    const cells = n === 2 ? e.cells : [...e.cells, e.cells[0]!.map((a, i) => (a + e.cells[1]![i]!) / 2)];
    return { id: e.id, author: e.author, cells: cells.map((c, i) => ({ factor: `p${i + 1}`, range: range(c) })) };
  });
}

function modelInput(n: 2 | 3): AblationModelInput {
  return {
    baseline: IN.baseline,
    full: IN.full,
    direction: "maximize",
    knownBest: IN.knownBest, // 2.635：calc15.py 第 132 行写死的已知最好
    epsilon: 0.01,
    lambda: 0.1,
    factors: n === 2 ? ["p1", "p2"] : ["p1", "p2", "p3"],
    explanations: explanations(n),
  };
}

const M2 = ablationModel(modelInput(2));
const M3 = ablationModel(modelInput(3));

// 容差：下面的量 calc15 和 rules 都是同一组浮点运算的不同次序（求和次序、高斯-勒让德 vs 多项式系数积分），
// 实测差在 1e-15 量级；取 1e-12 只容浮点舍入，不容算法差异（信息价值本身在 1e-4 量级）。
const FP = 1e-12;

test("calc15：总提升 G、剩余空间 h、两格撑几成", () => {
  close(totalGain(IN.baseline, IN.full, "maximize"), X.G, FP, "G（expected.json G）");
  close(M2.gain, X.G, FP, "模型里的 G（expected.json G）");
  close(M2.space, X.h, FP, "h（expected.json h，= 2.635 − 完整）");
  close(supportOf(IN.baseline, IN.full, IN.offGrid, "maximize"), X.supportOffGrid, FP, "撑(p1)（supportOffGrid）");
  close(supportOf(IN.baseline, IN.full, IN.offSolverLegalized, "maximize"), X.supportOffSolver, FP, "撑(p2)（supportOffSolver）");
});

test("calc15：起始权重与各因子事前期望撑几成", () => {
  // 预留起始权重 λ = 0.1（expected.json hand.reservePrior）；三人各 0.3，每人三条各 0.1。
  close(M2.weights[M2.weights.length - 1]!, X.hand.reservePrior, FP, "预留起始权重（hand.reservePrior）");
  for (let j = 0; j < 9; j++) close(M2.weights[j]!, 0.1, FP, `解释 e${j + 1} 起始权重（手算 0.9/3/3）`);
  // priorSupport2/3 = prior/h（dump.py），撑几成刻度。
  X.priorSupport2.forEach((p, c) => close(M2.prior[c]!, p, FP, `两格例 p${c + 1} 事前期望撑几成（priorSupport2[${c}]）`));
  X.priorSupport3.forEach((p, c) => close(M3.prior[c]!, p, FP, `三格例 p${c + 1} 事前期望撑几成（priorSupport3[${c}]）`));
  // ĝ₀ = h × max 期望撑几成，收益刻度（expected.json g0）；p2 的事前期望收益就是 ĝ₀（hand.solverGainPrior）。
  close(M2.space * M2.prior[1]!, X.hand.solverGainPrior, FP, "p2 事前期望收益（hand.solverGainPrior）");
  close(evaluateCombos(M2, { balance: 5, batchCap: 2, u: 0.5, decay: 0.5 }).g0, X.g0, FP, "ĝ₀（g0）");
});

test("calc15：网格规模 m_k 与每格小段数", () => {
  for (const k of [1, 2, 3, 4]) assert.equal(gridSize(k), X.gridCap[String(k)], `m_${k}（gridCap）`);
  // 两格例：p1 31 段、p2 30 段，m_1、m_2 都够每段 20 份 → 620、600（cellMiniSegments.k1/k2）。
  assert.deepEqual([0, 1].map((c) => cellGrid(M2.supportKnots[c]!, 1).mids.length), X.cellMiniSegments.k1);
  assert.deepEqual([0, 1].map((c) => cellGrid(M2.supportKnots[c]!, 2).mids.length), X.cellMiniSegments.k2);
  // 三格例 m_3 = 158：s = ⌊158/段数⌋ → 155、150、126（cellMiniSegments.k3）。
  assert.deepEqual([0, 1, 2].map((c) => cellGrid(M3.supportKnots[c]!, 3).mids.length), X.cellMiniSegments.k3);
});

test("calc15：信息价值（收益刻度，已乘 h）——只测 p1 走网格，全测走一维精确积分", () => {
  // voi2["1"]：两格例只测 p1，calc15 voi_grid，k = 1。
  close(informationValueGrid(M2, [0]), X.voi2["1"]!, FP, "两格例只测 p1（voi2['1']）");
  close(informationValue(M2, [0]), X.voi2["1"]!, FP, "两格例只测 p1，经 informationValue（voi2['1']）");
  // voi2["1+2"]：两格都测，calc15 emax_all（高斯-勒让德逐段精确）。
  close(informationValueExact(M2), X.voi2["1+2"]!, FP, "两格例全测（voi2['1+2']）");
  close(informationValue(M2, [0, 1]), X.voi2["1+2"]!, FP, "两格例全测，经 informationValue（voi2['1+2']）");
  // voi3["1+2+3"]：三格例三格都测。
  close(informationValue(M3, [0, 1, 2]), X.voi3["1+2+3"]!, FP, "三格例全测（voi3['1+2+3']）");
});

test("calc15：机会成本（K = 2）", () => {
  // opp(s) = 2ĝ₀[(2+s)/(3+s) − 2/3]（calc15.py 第 218–220 行），和余额无关；余额取 5 让两格组合可选。
  const { combos } = evaluateCombos(M2, { balance: 5, batchCap: 2, u: 0.5, decay: 0.5 });
  const cost = (key: string): number => combos.find((c) => c.factors.join("+") === key)!.cost;
  close(cost("p1"), X.opp["1"]!, FP, "测 1 格的成本（opp['1']）");
  close(cost("p2"), X.opp["1"]!, FP, "测 1 格的成本（opp['1']）");
  close(cost("p1+p2"), X.opp["2"]!, FP, "测 2 格的成本（opp['2']）");
});

test("calc15：测到 p1 = 2.599 后权重更新与押的因子（押 p2，预测收益约 0.0028939）", () => {
  // calc15 手算例（第 233–240 行）用 k = 1 网格里那一小段的把握更新权重；小段落在每个成分同一折线段内时，
  // 小段把握 = 密度 × 小段宽，与方法 6.8"新权重 ∝ 旧权重 × 密度"只差一个公共因子，归一后相同。
  const bet = chooseBet({
    ...modelInput(2),
    cells: [
      { factor: "p1", outcome: { kind: "valid", score: IN.offGrid } },
      { factor: "p2", outcome: { kind: "untested" } },
    ],
  });
  assert.equal(bet.weightsKept, false);
  close(bet.weights[9]!, X.hand.reservePosterior, FP, "预留更新后权重（hand.reservePosterior）");
  close(
    bet.weights.slice(0, 9).reduce((s, w) => s + w, 0),
    X.hand.explanationsPosterior,
    FP,
    "九条解释更新后权重和（hand.explanationsPosterior）",
  );
  close(bet.gains[0]!.gain, X.hand.gainOffGrid, FP, "p1 实测收益（hand.gainOffGrid）");
  assert.equal(bet.gains[0]!.measured, true);
  close(bet.gains[1]!.gain, X.hand.solverGainPosterior, FP, "p2 更新后期望收益（hand.solverGainPosterior）");
  assert.equal(bet.factor, "p2");
  assert.equal(bet.decision, "bet");
  close(bet.predictedGain, X.hand.solverGainPosterior, FP, "预测收益（hand.solverGainPosterior）");
});

test("calc15：配合关系标签（示例与 14-real-1 两格实测）", () => {
  // calc15 label 只返回中文，这里按方法 6.6 表对成代号；calc15 不判"有因子关掉反而更好"，这些示例里没有负数，不受影响。
  const code: Record<string, string> = {
    缺一不可: "indispensable",
    互补: "complementary",
    替代: "substitute",
    各管各: "independent",
    "不全，不贴标签": "incomplete",
  };
  for (const [supports, name] of Object.entries(X.labels)) {
    assert.equal(synergyLabel(JSON.parse(supports) as Array<number | null>).label, code[name], `labels ${supports}`);
  }
  assert.equal(synergyLabel([X.supportOffGrid, X.supportOffSolver]).label, code[X.label2], "label2");
});

// ---------------------------------------------------------------------------
// 正向：12-real-2

const R = (q1: number, q5: number, q9: number): Range => ({ q1, q5, q9 });
const B0 = 0.9597642169962064;
const ROUND1 = [
  { id: "c1", ranges: [{ researcher: "r1", range: R(2.44, 2.52, 2.62) }, { researcher: "r2", range: R(2.44, 2.5, 2.6) }, { researcher: "r3", range: R(2.44, 2.52, 2.61) }] },
  { id: "c2", ranges: [{ researcher: "r1", range: R(2.0, 2.3, 2.55) }, { researcher: "r2", range: R(2.0, 2.3, 2.5) }, { researcher: "r3", range: R(2.1, 2.35, 2.6) }] },
  { id: "c3", ranges: [{ researcher: "r1", range: R(2.0, 2.25, 2.5) }, { researcher: "r2", range: R(2.2, 2.35, 2.6) }, { researcher: "r3", range: R(1.9, 2.15, 2.45) }] },
];
const B2 = 2.61549458171637;
const ROUND2 = [
  { id: "c1", ranges: [{ researcher: "r1", range: R(2.6, 2.615, 2.632) }, { researcher: "r2", range: R(2.608, 2.6155, 2.625) }, { researcher: "r3", range: R(2.605, 2.618, 2.628) }] },
  { id: "c2", ranges: [{ researcher: "r1", range: R(2.618, 2.624, 2.632) }, { researcher: "r2", range: R(2.617, 2.625, 2.633) }, { researcher: "r3", range: R(2.617, 2.622, 2.63) }] },
  { id: "c3", ranges: [{ researcher: "r1", range: R(2.617, 2.623, 2.631) }, { researcher: "r2", range: R(2.617, 2.625, 2.632) }, { researcher: "r3", range: R(2.616, 2.621, 2.628) }] },
];

function batch(candidates: BatchInput["candidates"], baseline: number): BatchInput {
  return { candidates, baseline, direction: "maximize", epsilon: 0.01, lambda: 0.1, batchCap: 2, balance: 10, seed: "calc14b" };
}

test("12-real-2 第 1 轮：世界比特与惊讶按 Δ = 0.0001 从 b 起切", () => {
  // 预期值：forward_bits.py 输出 c1 bits 12.1453096896 surprise 11.7641504235；c2 bits 13.4884609551 surprise 14.4997227824
  // （RivalBits tests/info.test.ts 第 255 行附近）。打印到小数第 10 位，容差 1e-9 只容打印截断。
  const width = bitWidth(0.01);
  assert.equal(width, 0.0001);
  const [c1, c2] = forwardMixtures(ROUND1, 0.1);
  close(worldBits(c1!, B0, width), 12.1453096896, 1e-9, "c1 世界比特");
  close(worldBits(c2!, B0, width), 13.4884609551, 1e-9, "c2 世界比特");
  const s1 = surprise(c1!, B0, width, 2.6086313936944947).bits;
  const s2 = surprise(c2!, B0, width, 2.61549458171637).bits;
  assert.ok(s1.status === "measured" && s2.status === "measured");
  close(s1.value, 11.7641504235, 1e-9, "c1 惊讶");
  close(s2.value, 14.4997227824, 1e-9, "c2 惊讶");
});

// 选批 V 的容差 1e-6：calc14b 用 40 万段梯形积分（第 1 轮步长约 7.6e-6，第 2 轮约 4.6e-7），rules 是逐段多项式精确积分；
// 两者差来自梯形误差，实测在 1e-9 量级。预期值写到小数第 7 位（RivalBits tests/info.test.ts 第 228、286–291 行）。
const QUAD = 1e-6;

test("12-real-2 第 1 轮选批：c1 + c2 1.6009991 最好，分数尺读数 0.03471002", () => {
  const d = selectBestExpected(batch(ROUND1, B0));
  // 三个候选都远高于 b − ε，全部入池，所以 calc14b 对全部候选取 max 与方法 6.4 对候选池取 max 相同。
  assert.deepEqual(d.pool, ["c1", "c2", "c3"]);
  const v = (ids: string): number => d.combos.find((c) => c.ids.join("+") === ids)!.value;
  close(v("c1+c2"), 1.6009991, QUAD, "calc14b ('c1','c2')");
  close(v("c1+c3"), 1.5974972, QUAD, "calc14b ('c1','c3')");
  close(v("c2+c3"), 1.4845937, QUAD, "calc14b ('c2','c3')");
  assert.deepEqual(d.chosen, ["c1", "c2"]);
  close(d.scaleReading, 0.03471002, QUAD, "calc14b EVPI batch");
});

test("12-real-2 第 2 轮选批：c2 + c3 0.0136813 最好，分数尺读数 0.0020944", () => {
  const d = selectBestExpected(batch(ROUND2, B2));
  // 入池把握 P(Y < −ε)：本 Agent 用 calc14b.py 的 Combined.F 在 b − ε 处复算，c1 0.16951174225193、c2 c3 0.04101139196591，
  // 都低于一半，全部入池（读数口径与方法一致的前提）。numpy 插值与 rules 同式，容差 1e-12。
  assert.deepEqual(d.pool, ["c1", "c2", "c3"]);
  close(d.candidates[0]!.lossProbability, 0.1695117422519303, FP, "c1 入池把握（calc14b F(b − ε)）");
  close(d.candidates[1]!.lossProbability, 0.04101139196591438, FP, "c2 入池把握（calc14b F(b − ε)）");
  close(d.candidates[2]!.lossProbability, 0.04101139196591438, FP, "c3 入池把握（calc14b F(b − ε)）");
  const v = (ids: string): number => d.combos.find((c) => c.ids.join("+") === ids)!.value;
  close(v("c2+c3"), 0.0136813, QUAD, "calc14b ('c2','c3')");
  close(v("c1+c2"), 0.0126357, QUAD, "calc14b ('c1','c2')");
  close(v("c1+c3"), 0.0120484, QUAD, "calc14b ('c1','c3')");
  assert.deepEqual(d.chosen, ["c2", "c3"]);
  close(d.scaleReading, 0.0020944, QUAD, "calc14b EVPI");
});