import assert from "node:assert/strict";
import { test } from "node:test";
import {
  diagnosisModel,
  runDiagnosisRule,
  selectDiagnosis,
  settleDiagnosis,
  trackRecord,
  type DiagnosisExplanation,
  type DiagnosisRuleInput,
  type ExperimentOutcome,
} from "../src/rules/index.ts";
import { close, loadFixture, scenarioPicker, type Scenario } from "./rules-fixtures-helpers.ts";

// 预期值来自 tests/fixtures/rules/gen_diagnosis.py：只照方法.md 6.12、6.7、6.9 的文字、用有理数写成的独立参考，
// 没有读过 src/rules/diagnosis.ts 和 value.ts（重新生成：python3 tests/fixtures/rules/gen_diagnosis.py）。
// 场景里有手算的几个（opposite_lambda0 的 VOI = 0.1，same_lambda0 的 VOI = 0，说明写在场景的 note 里）。

const fixture = loadFixture("diagnosis.json");
const pick = scenarioPicker(fixture);

const EXACT = 1e-12;
const BITS = 1e-9;
/** 只做一个实验时的 VOI：实现用网格近似，参考是逐段闭式积分。 */
const closeK1 = (actual: number, expected: number, what: string): void =>
  close(actual, expected, 1e-9 + 2e-4 * Math.abs(expected), what);

type Ranges = Record<string, { q1: number; q5: number; q9: number }>;

function explanationsOf(s: Scenario): DiagnosisExplanation[] {
  return s.input.explanations.map((e: { id: string; author: string; ranges: Ranges }) => ({
    id: e.id,
    author: e.author,
    experiments: Object.entries(e.ranges).map(([experiment, range]) => ({ experiment, range })),
  }));
}

function ruleInput(s: Scenario): DiagnosisRuleInput {
  return {
    baseline: s.input.b,
    direction: s.input.direction,
    lambda: s.input.lambda,
    experiments: s.input.experiments,
    explanations: explanationsOf(s),
    balance: s.input.balance,
    batchCap: s.input.batchCap,
    maxExperiments: s.input.maxExperiments,
    u: s.input.u,
    decay: s.input.decay,
  };
}

/** 选实验场景：事前数字、组合表、选中的实验都和独立参考一致。 */
function checkSelection(name: string): ReturnType<typeof selectDiagnosis> {
  const s = pick.get(name);
  const got = selectDiagnosis(ruleInput(s));
  const e = s.expected;
  for (const [d, v] of Object.entries(e.priorGains as Record<string, number>)) {
    close(got.priorGains.find((x) => x.experiment === d)!.gain, v, EXACT, `${name} ${d} 事前期望收益`);
  }
  close(got.g0, e.g0, EXACT, `${name} ĝ₀`);
  e.weights.forEach((w: number, j: number) => close(got.weights[j]!, w, EXACT, `${name} 起始权重 ${j}`));
  assert.deepEqual(got.combos.map((c) => c.experiments), e.combos.map((c: any) => c.experiments), `${name} 组合列表`);
  e.combos.forEach((c: any, i: number) => {
    const g = got.combos[i]!;
    if (c.voi === null) return;
    const what = `${name} {${c.experiments}}`;
    if (c.experiments.length === s.input.experiments.length) close(g.voi, c.voi, EXACT, `${what} VOI`);
    else closeK1(g.voi, c.voi, `${what} VOI`);
    assert.equal(g.rHat, c.rHat, `${what} R̂`);
    close(g.mPrime, c.mPrime, EXACT, `${what} m′`);
    close(g.cost, c.cost, EXACT, `${what} 成本`);
    close(g.value, c.value, 1e-9 + 2e-4 * Math.abs(c.value), `${what} 价值`);
    close(g.net, c.net, 1e-9 + 2e-4 * Math.abs(c.value), `${what} 净值`);
  });
  if ("diagnose" in e) {
    assert.equal(got.diagnose, e.diagnose, `${name} 做不做`);
    assert.deepEqual(got.chosen, e.chosen, `${name} 选中的实验`);
  }
  return got;
}

test("6.12 信息价值：两条解释对实验预测相反时为正（手算 0.1），预测一样时为 0", () => {
  const opposite = checkSelection("opposite_lambda0");
  close(opposite.combos[0]!.voi, 0.1, EXACT, "手算：0.5 × 0.2 − 0");
  assert.equal(opposite.diagnose, true);

  const same = checkSelection("same_lambda0");
  close(same.combos[0]!.voi, 0, EXACT, "预测一样、范围都在 0 以上：做不做都不改变看法");
  assert.equal(same.diagnose, false, "净值 = −成本 < 0");
  assert.equal(same.experiments[0]!.reason, "value");
  assert.equal(same.chosenNet, 0);

  // λ = 0.1：预测一样时只剩开放预留带来的一小部分，比预测相反小得多。
  const oppositeReserve = checkSelection("opposite_lambda01");
  const sameReserve = checkSelection("same_lambda01");
  close(sameReserve.combos[0]!.voi, 0.0025, EXACT, "手算：λ × (E[max(0, U)] − E U) = 0.1 × 0.025");
  assert.ok(sameReserve.combos[0]!.voi > 0);
  assert.ok(oppositeReserve.combos[0]!.voi > 10 * sameReserve.combos[0]!.voi);
});

test("6.12 价值减成本：不管做几个实验成本相同；R̂ = ⌈(余额 − 1) ÷ 每批上限⌉", () => {
  const got = checkSelection("multi_all");
  // 手算：余额 9、K = 2：R̂ = ⌈8 ÷ 2⌉ = 4，m′ = 1 + 0.5 + 0.25 + 0.125 = 1.875；
  // 成本 = 2 × 0.0525 × (3/4 − 2/3) = 0.00875（ĝ₀ = 0.0525 来自参考脚本的事前期望）。
  assert.equal(got.combos.length, 1, "实验数 3 ≤ D = 4：只有全做一个组合");
  assert.equal(got.combos[0]!.rHat, 4);
  close(got.combos[0]!.mPrime, 1.875, EXACT, "m′");
  close(got.combos[0]!.cost, 0.00875, EXACT, "成本");
  assert.deepEqual(got.chosen, ["d1", "d2", "d3"]);
  assert.ok(got.experiments.every((x) => x.done && x.reason === null && x.bestCombo === null));
});

test("6.12 选实验：余额 < 2 时不诊断，没有可选组合", () => {
  const got = checkSelection("balance_one");
  assert.deepEqual(got.combos, []);
  assert.equal(got.diagnose, false);
  assert.deepEqual(got.chosen, []);
  assert.ok(got.experiments.every((x) => !x.done && x.bestCombo === null));
});

test("6.12 选实验：实验数超过 D 时只比较装满的组合", () => {
  const one = checkSelection("over_cap_d1");
  assert.equal(one.combos.length, 3, "C(3,1) 个单实验组合");
  assert.ok(one.combos.every((c) => c.experiments.length === 1), "没有空组合，也没有更大的组合");
  assert.deepEqual(one.chosen, ["d2"]);
  // 没选中的实验：放不下，并附上含它的组合里净值最大的那个（这里就是它自己的单实验组合）。
  const d3 = one.experiments.find((x) => x.experiment === "d3")!;
  assert.equal(d3.done, false);
  assert.equal(d3.reason, "budget");
  assert.deepEqual(d3.bestCombo!.experiments, ["d3"]);
  close(d3.bestCombo!.net, pick.get("over_cap_d1").expected.combos[2].net, 1e-9 + 2e-4 * 0.1, "d3 的净值");

  const two = checkSelection("over_cap_d2");
  assert.deepEqual(two.combos.map((c) => c.experiments), [["d1", "d2"], ["d1", "d3"], ["d2", "d3"]]);
  assert.equal(two.chosen.length, 2);
  // 多做一个实验信息价值只会不变或变大：每个两实验组合的 VOI 不小于它包含的任一单实验（参考脚本的精确值）。
  const singles = pick.get("over_cap_d1").expected.combos as Array<{ experiments: string[]; voi: number }>;
  for (const pair of two.combos) {
    for (const d of pair.experiments) {
      const single = singles.find((c) => c.experiments[0] === d)!.voi;
      assert.ok(pair.voi >= single - (1e-9 + 2e-4 * single), `{${pair.experiments}} 的 VOI ${pair.voi} 小于 {${d}} 的 ${single}`);
    }
  }
});

test("6.12 选实验：平票按实验编号（自然序，d2 排在 d10 前面）；越小越好按改进量算", () => {
  const tie = checkSelection("tie_by_id");
  assert.deepEqual(tie.chosen, ["d2"]);
  assert.deepEqual(tie.experiments.find((x) => x.experiment === "d10")!.bestCombo!.experiments, ["d10"]);
  checkSelection("minimize_all");
});

test("登记表：选诊断实验只有一条规则，结果带规则名和版本；同样的输入算两次逐位相同", () => {
  const s = pick.get("multi_all");
  const run = runDiagnosisRule("valueMinusCost", ruleInput(s));
  assert.deepEqual(run.rule, { name: "valueMinusCost", version: 1 });
  assert.deepEqual(JSON.parse(JSON.stringify(run)), JSON.parse(JSON.stringify({ rule: run.rule, ...selectDiagnosis(ruleInput(s)) })));
});

test("成分模型：每条解释一个成分加开放预留；解释合起来 1 − λ，同一人平分", () => {
  const s = pick.get("settle_mixed");
  const m = diagnosisModel({ ...ruleInput({ ...s, input: { ...s.input, balance: 9, batchCap: 2, maxExperiments: 4, u: 0.5, decay: 0.5 } }), });
  // 手算：researcher-1 一条 0.9 × 1/2，researcher-2 两条各 0.9 × 1/4，预留 0.1。
  assert.deepEqual(m.components, ["e1", "e2", "e3", "openReserve"]);
  [0.45, 0.225, 0.225, 0.1].forEach((w, j) => close(m.weights[j]!, w, EXACT, `权重 ${j}`));
});

type Settled = ReturnType<typeof settleDiagnosis>;

function settle(name: string): { got: Settled; e: any } {
  const s = pick.get(name);
  const outcomes = s.input.outcomes as Record<string, { kind: ExperimentOutcome["kind"]; score?: number }>;
  const got = settleDiagnosis({
    baseline: s.input.b,
    direction: s.input.direction,
    epsilon: s.input.epsilon,
    lambda: s.input.lambda,
    experiments: Object.entries(outcomes).map(([experiment, o]) => ({
      experiment,
      outcome: o.kind === "valid" ? { kind: "valid", score: o.score! } : ({ kind: o.kind } as ExperimentOutcome),
    })),
    explanations: explanationsOf(s),
  });
  const e = s.expected;
  assert.deepEqual(got.experiments.map((x) => x.outcome), Object.keys(e.experiments).map((d) => e.experiments[d].outcome), `${name} 结局`);
  for (const row of got.experiments) {
    const x = e.experiments[row.experiment];
    const what = `${name} ${row.experiment}`;
    close(row.merged.q1, x.merged.q1, BITS, `${what} 合并一成`);
    close(row.merged.q5, x.merged.q5, BITS, `${what} 合并一半`);
    close(row.merged.q9, x.merged.q9, BITS, `${what} 合并九成`);
    if (x.gain === undefined) {
      // 证据不足、没做：不判、不记账。
      assert.equal(row.gain, null, `${what} 收益`);
      assert.equal(row.worldBits, null, `${what} 世界比特`);
      assert.equal(row.surprise, null, `${what} 惊讶`);
      assert.equal(row.disagreement, null, `${what} 分歧量`);
    } else {
      close(row.gain!, x.gain, EXACT, `${what} 收益`);
      close(row.worldBits!, x.worldBits, BITS, `${what} 世界比特`);
      assert.equal(row.surprise!.cell, x.surprise.bin, `${what} 惊讶所在格`);
      close(row.surprise!.prior, x.surprise.p, BITS, `${what} 事前把握`);
      assert.equal(row.surprise!.bits.status, x.surprise.status === "ok" ? "measured" : "unmeasured");
      if (row.surprise!.bits.status === "measured") close(row.surprise!.bits.value, x.surprise.surprise, BITS, `${what} 惊讶`);
      close(row.disagreement!, x.disagreement, BITS, `${what} 分歧量`);
    }
  }
  // 确认与否定：只列有效实验。
  assert.deepEqual(
    Object.fromEntries(got.judgments.map((j) => [j.id, Object.fromEntries(j.verdicts.map((v) => [v.experiment, v.verdict]))])),
    e.judgments,
    `${name} 确认与否定`,
  );
  const dir = got.direction;
  assert.equal(dir.weightsKept, e.direction.weightsKept, `${name} 权重是否保持`);
  e.direction.weights.forEach((w: number, j: number) => close(dir.weights[j]!, w, EXACT, `${name} 新权重 ${j}`));
  for (const g of dir.gains) {
    const x = e.direction.gains[g.experiment];
    close(g.gain, x.gain, EXACT, `${name} ${g.experiment} 方向收益`);
    assert.equal(g.measured, x.measured, `${name} ${g.experiment} 是实测还是估计`);
  }
  assert.equal(dir.experiment, e.direction.experiment, `${name} 推出的方向`);
  close(dir.gain, e.direction.gain, EXACT, `${name} 方向收益`);
  assert.equal(dir.found, e.direction.found, `${name} 有没有指出方向`);
  return { got, e };
}

test("6.12 结算：证据不足的实验不判、不记账；没做的实验也不判", () => {
  const { got } = settle("settle_mixed");
  assert.deepEqual(got.experiments.map((x) => x.outcome), ["valid", "insufficient", "untested"]);
  // d2、d3 没有有效分数：每条解释只有 d1 一行判定，没有比特。
  assert.ok(got.judgments.every((j) => j.verdicts.length === 1 && j.verdicts[0]!.experiment === "d1"));
  assert.equal(got.experiments.filter((x) => x.worldBits !== null).length, 1);
  // 手算：d1 实测 1.045 = b + 0.045，收益 +0.045；落在 e1 的 [1.02, 1.08] 里确认，e2 的 [1.00, 1.06] 里确认，e3 的 [0.98, 1.10] 里确认。
  assert.deepEqual(got.judgments.map((j) => j.verdicts[0]!.verdict), ["confirmed", "confirmed", "confirmed"]);
});

test("6.12 结算：两个有效实验、确认与否定、权重更新", () => {
  const { got } = settle("settle_two_valid");
  // d2 实测 1.25，e2 的 [1.1, 1.3] 确认；e1 的 [0.95, 1.05]、e3 的 [0.8, 1.0] 否定。
  assert.deepEqual(got.judgments.find((j) => j.id === "e2")!.verdicts.find((v) => v.experiment === "d2")!.verdict, "confirmed");
  assert.deepEqual(got.judgments.find((j) => j.id === "e1")!.verdicts.find((v) => v.experiment === "d2")!.verdict, "denied");
  assert.deepEqual(got.judgments.find((j) => j.id === "e3")!.verdicts.find((v) => v.experiment === "d2")!.verdict, "denied");
});

test("6.12 结算：全部实验收益不为正时，记没有指出能涨分的方向", () => {
  const { got } = settle("settle_no_direction");
  assert.equal(got.direction.found, false);
  assert.ok(got.direction.gains.every((g) => g.gain <= 0));
  // 手算：收益分别是 0.97 − 1 = −0.03、0.8 − 1 = −0.2、0.9 − 1 = −0.1，最大的是 d1 的 −0.03。
  assert.equal(got.direction.experiment, "d1");
  close(got.direction.gain, -0.03, EXACT, "最大收益");
});

test("6.12 结算：没有有效实验、实测落在所有成分之外时保持起始权重", () => {
  const none = settle("settle_all_insufficient");
  assert.equal(none.got.direction.weightsKept, false, "没有有效实验时没有密度可乘，权重原样（不算'无从更新'）");
  assert.ok(none.got.judgments.every((j) => j.verdicts.length === 0));
  const outside = settle("settle_outside_all");
  assert.equal(outside.got.direction.weightsKept, true);
  assert.deepEqual(outside.got.direction.weights, diagnosisModel({
    baseline: 1, direction: "maximize", lambda: 0.1, experiments: ["d1", "d2", "d3"],
    explanations: explanationsOf(pick.get("settle_outside_all")),
  }).weights);
});

test("6.12 结算：越小越好按改进量 b − x", () => {
  const { got } = settle("settle_minimize");
  close(got.experiments[0]!.gain!, 0.12, EXACT, "手算：1.0 − 0.88");
  close(got.experiments[1]!.gain!, -0.05, EXACT, "手算：1.0 − 1.05");
  assert.equal(got.direction.experiment, "d1");
});

test("战绩：按人、按正向 / 诊断 / 消融分开数，次序是编号、正向、诊断、消融", () => {
  // 手算：researcher-2 的诊断样本 (q1, q5, q9) = (0.1, 0.2, 0.3)，实测 0.35 在 q9 之上：中位数偏低、被穿过上沿、没包住、落点 aboveQ9。
  const range = { q1: 0.1, q5: 0.2, q9: 0.3 };
  const rows = trackRecord([
    { researcher: "researcher-2", phase: "ablation", range, value: 0.2 },
    { researcher: "researcher-2", phase: "diagnosis", range, value: 0.35 },
    { researcher: "researcher-1", phase: "diagnosis", range, value: 0.05 },
    { researcher: "researcher-2", phase: "forward", range, value: 0.25 },
  ]);
  assert.deepEqual(rows.map((r) => `${r.researcher}/${r.phase}`), [
    "researcher-1/diagnosis", "researcher-2/forward", "researcher-2/diagnosis", "researcher-2/ablation",
  ]);
  const t = rows[2]!.tally;
  assert.deepEqual([t.total, t.medianLow, t.medianHigh, t.aboveTop, t.belowBottom, t.covered], [1, 1, 0, 1, 0, 0]);
  assert.equal(t.placements.aboveQ9, 1);
});

test("diagnosis.json 里的场景都被测到", () => {
  assert.deepEqual(pick.unused(), []);
});