import assert from "node:assert/strict";
import { test } from "node:test";
import {
  ablationFinding,
  ablationModel,
  chooseBet,
  evaluateCombos,
  gridSize,
  inseparableGroups,
  informationValue,
  informationValueExact,
  informationValueGrid,
  judgeExplanations,
  measuredGrid,
  OPEN_RESERVE,
  pickByNet,
  topCombo,
  selectAll,
  selectByValue,
  supportOf,
  synergyLabel,
  totalGain,
  type AblationCell,
  type AblationModel,
  type AblationModelInput,
  type CellOutcome,
  type CellRuleInput,
  type SynergyLabel,
  type UntestedReason,
} from "../src/rules/index.ts";
import { close, closeArray, closeGrid, loadFixture, range, scenarioPicker } from "./rules-fixtures-helpers.ts";

// ablation.json：方法.md 6.6–6.8、4.3 ⑰ 与约定 17–30，对照独立 Python 参考实现算出的预期值。
// 容差见 ablation.json tolerance：精确量 1e-12；网格值 1e-12 + 1e-9·|值|；离散结果完全相同；
// 网格对一维精确按 tolerance.gridVsExact。

const fixture = loadFixture("ablation.json");
const pick = scenarioPicker(fixture);
const TOL = 1e-12;
const SEED = "fixture-seed";

/** Python 的标签、结论、原因名 → rules 的代号。 */
const LABEL: Record<string, SynergyLabel> = {
  negative: "offBetter",
  incomplete: "incomplete",
  indispensable: "indispensable",
  complementary: "complementary",
  substitute: "substitute",
  independent: "independent",
};
const VERDICT: Record<string, string> = { confirmed: "confirmed", refuted: "denied" };
const REASON: Record<string, UntestedReason> = { insufficientBudget: "budget", notWorthCost: "value" };
const BET: Record<string, string> = { bet: "bet", noBetNegativeSupport: "offBetter", noBetNonPositiveGain: "notPositive" };
const componentName = (n: string): string => (n === "reserve" ? OPEN_RESERVE : n);

function modelInput(name: string): AblationModelInput & { params: Record<string, number> } {
  const m = fixture.models![name]!;
  const i = m.input;
  return {
    baseline: i.B,
    full: i.full,
    direction: i.direction,
    knownBest: i.knownBest,
    epsilon: m.params.eps!,
    lambda: m.params.lambda!,
    factors: i.cells,
    explanations: i.explanations.map((e: any) => ({
      id: e.id,
      author: e.author,
      cells: Object.entries(e.cells).map(([factor, r]) => ({ factor, range: range(r as any) })),
    })),
    params: m.params,
  };
}

/** Python 的一格 {factor, kind, value?} → rules 的 AblationCell。一起关的 off 只用于沿用（sharedSettings），这里用到的函数不读它。 */
function cellOf(factor: string, o: { kind: string; value?: number }): AblationCell {
  const outcome: CellOutcome =
    o.kind === "valid"
      ? { kind: "valid", score: o.value! }
      : o.kind === "joint"
        ? { kind: "joint", off: [factor], score: o.value ?? null }
        : { kind: o.kind as "insufficient" | "untested" | "cannotRun" };
  return { factor, outcome };
}

const MODEL_CACHE = new Map<string, AblationModel>();
function modelOf(name: string): AblationModel {
  if (!MODEL_CACHE.has(name)) MODEL_CACHE.set(name, ablationModel(modelInput(name)));
  return MODEL_CACHE.get(name)!;
}

// ---------------------------------------------------------------- 撑几成、标签（6.6）

for (const name of [
  "label_method_example_independent",
  "label_calc15_substitute",
  "label_r1_negative",
  "label_incomplete_with_other_outcomes",
  "label_negative_beats_incomplete",
  "label_indispensable",
  "label_complementary",
  "label_boundary_exact_0.9",
  "label_minimize_substitute",
  "label_G_not_positive",
  "label_G_negative_minimize",
]) {
  test(`ablation.json ${name}`, () => {
    const s = pick.get(name);
    const { B, full, direction, cells } = s.input;
    if (s.expected.error !== undefined) {
      assert.equal(s.expected.error, "totalGainNotPositive");
      assert.throws(() => totalGain(B, full, direction), /G = /);
      return;
    }
    close(totalGain(B, full, direction), s.expected.G, TOL, "G");
    const supports = cells.map((c: any) => (c.kind === "valid" ? supportOf(B, full, c.value, direction) : null));
    cells.forEach((c: any, i: number) => {
      const e = s.expected.supports[c.factor];
      if (e === null) assert.equal(supports[i], null, `${c.factor} 没有撑几成`);
      else close(supports[i]!, e, TOL, `${c.factor} 撑几成`);
    });
    const got = synergyLabel(supports);
    assert.equal(got.label, LABEL[s.expected.label], "标签");
    // Python 在"关掉反而更好""不全"两种标签下不给 i；rules 在"不全"时给 null，在"关掉反而更好"且每格都有效时照给数。
    if (s.expected.interaction !== null) close(got.interaction!, s.expected.interaction, TOL, "相对交互量 i");
    if (got.label === "incomplete") assert.equal(got.interaction, null);
  });
}

// ---------------------------------------------------------------- 确认与否定、分不开（6.6）

for (const name of ["confirm_r1_real", "confirm_calc15", "confirm_endpoint_inclusive"]) {
  test(`ablation.json ${name}`, () => {
    const s = pick.get(name);
    const input = modelInput(s.input.model);
    const cells = Object.entries(s.input.valid as Record<string, number>).map(([f, v]) => cellOf(f, { kind: "valid", value: v }));
    const judgments = judgeExplanations(input.explanations, cells);
    assert.deepEqual(
      Object.fromEntries(judgments.map((j) => [j.id, Object.fromEntries(j.verdicts.map((v) => [v.factor, v.verdict]))])),
      Object.fromEntries(
        Object.entries(s.expected.verdicts as Record<string, Record<string, string>>).map(([id, row]) => [
          id,
          Object.fromEntries(Object.entries(row).map(([f, v]) => [f, VERDICT[v]])),
        ]),
      ),
    );
    assert.deepEqual(inseparableGroups(judgments), s.expected.inseparableGroups);
  });
}

// ---------------------------------------------------------------- 消融中发现（4.3 ⑰）

for (const name of ["discovery_r1real", "discovery_calc15", "discovery_gainNotPositive"]) {
  test(`ablation.json ${name}`, () => {
    const s = pick.get(name);
    const input = modelInput(s.input.model);
    const got = ablationFinding({
      full: input.full,
      knownBest: input.knownBest,
      epsilon: s.params.eps,
      direction: input.direction,
      cells: s.input.cells.map((c: any) => cellOf(c.factor, c)),
    });
    close(got.line, s.expected.line, TOL, "采纳线");
    assert.deepEqual(got.passed.map((p) => p.factor), s.expected.overLine, "过线的格");
    assert.equal(got.chosen, s.expected.newBaseline, "新基线");
  });
}

// ---------------------------------------------------------------- 成分、先验、信息价值（6.7）

function gridVsExactTol(k: number): number {
  // tolerance.gridVsExact：k ≤ 3 时 5e-4·|精确值|，k = 4 时 3e-3·|精确值|。
  return k <= 3 ? 5e-4 : 3e-3;
}

for (const name of ["calc15", "r1real", "minimize", "fourCells", "negativeG0"]) {
  test(`ablation.json voi_${name}`, () => {
    const s = pick.get(`voi_${name}`);
    const e = s.expected;
    const model = modelOf(name);
    const idx = (f: string): number => model.factors.indexOf(f);
    close(model.gain, e.G, TOL, "G");
    close(model.space, e.h, TOL, "h");
    assert.deepEqual(model.components, e.componentNames.map(componentName), "成分名单");
    closeArray(model.weights, e.weights, TOL, "起始权重");
    for (const [f, ends] of Object.entries(e.reserve as Record<string, number[]>)) {
      const k = model.rawKnots[idx(f)]!.at(-1)!;
      closeArray([k[0]![0], k.at(-1)![0]], ends, TOL, `${f} 的开放预留`);
    }
    e.componentNames.forEach((n: string, j: number) => {
      for (const [f, v] of Object.entries(e.componentExpectedPayoff[n] as Record<string, number>)) {
        close(model.space * model.expected[idx(f)]![j]!, v, TOL, `成分 ${n} 在 ${f} 的期望收益`);
      }
    });
    for (const [f, v] of Object.entries(e.priorExpectedPayoff as Record<string, number>)) {
      close(model.space * model.prior[idx(f)]!, v, TOL, `${f} 的事前期望收益`);
    }
    close(evaluateCombos(model, { balance: 0, batchCap: 2, u: 0.5, decay: 0.5 }).g0, e.g0, TOL, "ĝ₀");

    const all = model.factors.map((_, c) => c);
    const exact = informationValueExact(model);
    close(exact, e.voiFullExact, TOL, "全测 VOI（一维精确）");
    close(informationValue(model, all), e.voiFullExact, TOL, "全测 VOI（informationValue 走精确）");

    // 网格：先核对每格折点数、小段数与 Python 一致（差一个 ulp 会多出一段），再比 VOI。
    const checkGrid = (S: string[], info: any, what: string): void => {
      assert.equal(gridSize(S.length), info.m, `${what}：m`);
      for (const f of S) {
        // 折点在分数原值上取（measuredGrid），两边都按浮点原式算、只合并完全相等的值。
        const brk = new Set(model.rawKnots[idx(f)]!.flatMap((k) => k.map(([x]) => x))).size;
        assert.equal(brk, info.cells[f].breakpoints, `${what} ${f}：折点数`);
        assert.equal(measuredGrid(model, idx(f), S.length).mids.length, info.cells[f].subsegments, `${what} ${f}：小段数`);
      }
    };
    for (const row of e.partial) {
      const S = row.cells as string[];
      checkGrid(S, row.grid, `测 ${S.join(",")}`);
      const v = informationValueGrid(model, S.map(idx));
      closeGrid(v, row.voiGrid, `测 ${S.join(",")} 的网格 VOI`);
      closeGrid(informationValue(model, S.map(idx)), row.voiGrid, `测 ${S.join(",")} 的 informationValue`);
      if (row.exactK1VOI !== undefined) {
        // 只测一格的精确参考值：rules 没有这个算法，只核对网格与它的差在 gridVsExact 的 k ≤ 3 界内。
        close(v, row.exactK1VOI, gridVsExactTol(1) * Math.abs(row.exactK1VOI), `测 ${S[0]} 的网格对精确`);
      }
    }
    checkGrid(model.factors, e.gridFullInfo, "全测走网格");
    const gridFull = informationValueGrid(model, all);
    closeGrid(gridFull, e.gridFullVOI, "全测走网格的 VOI");
    close(gridFull, exact, gridVsExactTol(all.length) * Math.abs(exact), "全测：网格对一维精确");
    // tolerance.gridVsExact：每格 s ≥ 1 时网格值不超过精确值（Jensen）；fourCells 有 s = 0 的格，不适用。
    if (name !== "fourCells") assert.ok(gridFull <= exact + 1e-12, `全测网格 ${gridFull} 不应超过精确 ${exact}`);
  });
}

// ---------------------------------------------------------------- 价值、成本、选格（6.7）

for (const name of ["calc15", "r1real", "minimize", "negativeG0"]) {
  test(`ablation.json select_${name}`, () => {
    const s = pick.get(`select_${name}`);
    const { params, ...input } = modelInput(name);
    const model = modelOf(name);
    for (const e of s.expected.byBalance) {
      const tag = `余额 ${e.balance}`;
      const ruleInput: CellRuleInput = {
        ...input,
        balance: e.balance,
        batchCap: params.batchCap!,
        u: params.u!,
        decay: params.r!,
        seed: SEED,
      };
      const { g0, combos } = evaluateCombos(model, ruleInput);
      close(g0, s.expected.g0, TOL, `${tag} ĝ₀`);
      if (name === "negativeG0") {
        // ĝ₀ < 0：原始值照记，成本按 max(0, ĝ₀) = 0 算（6.7"价值减成本"、约定 28）。
        assert.ok(g0 < 0, `${tag} ĝ₀ 应为负`);
        assert.ok(combos.length > 0 && combos.every((c) => c.cost === 0 && c.net === c.value), `${tag} 成本为 0、净值 = 价值`);
        assert.equal(selectByValue(ruleInput).g0, g0, `${tag} 选格记录里的 ĝ₀ 是原始负值`);
      }
      assert.deepEqual(combos.map((c) => c.factors), e.combos.map((c: any) => c.combo), `${tag} 列出的组合`);
      combos.forEach((c, i) => {
        const r = e.combos[i];
        const what = `${tag} {${c.factors.join(",")}}`;
        // 全测用一维精确（精确量容差），部分格用网格（网格容差）。
        const cmp = c.factors.length === model.factors.length ? (a: number, b: number, w: string) => close(a, b, TOL, w) : closeGrid;
        cmp(c.voi, r.voi, `${what} VOI`);
        assert.equal(c.rHat, r.Rhat, `${what} R̂`);
        close(c.mPrime, r.mPrime, TOL, `${what} m′`);
        cmp(c.value, r.value, `${what} 价值`);
        close(c.cost, r.cost, TOL, `${what} 成本`);
        closeGrid(c.net, r.net, `${what} 净值`);
      });
      const d = selectByValue(ruleInput);
      assert.deepEqual(d.chosen, e.selected, `${tag} 选中的格`);
      closeGrid(d.chosenNet, e.selectedNet, `${tag} 选中组合的净值`);
      for (const cell of d.cells) {
        const u = e.unselected[cell.factor];
        if (cell.tested) {
          assert.equal(u, undefined, `${tag} ${cell.factor} 应选中`);
          continue;
        }
        assert.equal(cell.reason, REASON[u.reason], `${tag} ${cell.factor} 没测的原因`);
        if (u.bestCombo === null) assert.equal(cell.bestCombo, null, `${tag} ${cell.factor} 没有可选组合`);
        else {
          assert.deepEqual(cell.bestCombo!.factors, u.bestCombo, `${tag} ${cell.factor} 含它的最好组合`);
          closeGrid(cell.bestCombo!.net, u.bestNet, `${tag} ${cell.factor} 含它的最好组合的净值`);
        }
      }
      const a = selectAll(ruleInput);
      assert.deepEqual(a.chosen, e.fullRule.selected, `${tag} 全测规则选中的格`);
      assert.deepEqual(
        a.cells.filter((c) => !c.tested).map((c) => [c.factor, c.reason]),
        e.fullRule.insufficientBudget.map((f: string) => [f, "budget"]),
        `${tag} 全测规则放不下的格`,
      );
    }
  });
}

// 纯选格：给定净值表，核对 1e-12 容差与平局（6.7"选格"第 2 条、约定 29）：pickByNet 选组合，topCombo 挑
// "含这一格的最好组合"。没选中的原因在 selectByValue 里由模型算出的净值表决定，这里覆盖不到。
for (const name of [
  "select_pure_net_at_tolerance",
  "select_pure_just_above_tolerance",
  "select_pure_tie_chain",
  "select_pure_budget_reason",
  "select_pure_tolerance_relay",
  "select_pure_filter_before_max",
]) {
  test(`ablation.json ${name}`, () => {
    const s = pick.get(name);
    const cells = s.input.cells as string[];
    const limit = Math.min(cells.length, s.input.balance - 1);
    // 组合按格数从少到多、同格数按因子编号字典序（约定 29），只列格数 ≤ 余额 − 1 的。
    const combos: Array<{ factors: string[]; net: number }> = [];
    for (let n = 1; n <= limit; n++) {
      const walk = (start: number, acc: string[]): void => {
        if (acc.length === n) {
          combos.push({ factors: acc, net: s.input.nets[acc.join(",")] });
          return;
        }
        for (let i = start; i < cells.length; i++) walk(i + 1, [...acc, cells[i]!]);
      };
      walk(0, []);
    }
    const chosen = pickByNet(combos);
    assert.deepEqual(chosen, s.expected.selected, "选中的格");
    const net = chosen.length === 0 ? 0 : s.input.nets[chosen.join(",")];
    assert.equal(net, s.expected.selectedNet, "选中组合的净值");
    // 没选中的格附"含这一格的净值最大组合"，按同一条规则挑（topCombo）。
    for (const [f, u] of Object.entries(s.expected.unselected as Record<string, any>)) {
      assert.ok(!chosen.includes(f));
      const best = topCombo(combos.filter((c) => c.factors.includes(f)));
      if (u.bestCombo === null) assert.equal(best, null, `${f} 没有可选组合`);
      else {
        assert.deepEqual(best!.factors, u.bestCombo, `${f} 含它的最好组合`);
        assert.equal(best!.net, u.bestNet, `${f} 含它的最好组合的净值`);
      }
    }
  });
}

// ---------------------------------------------------------------- 押的因子（6.8）

for (const name of [
  "bet_calc15_all_valid",
  "bet_calc15_p2_insufficient",
  "bet_calc15_joint_ignored",
  "bet_r1_negative_support",
  "bet_gain_not_positive",
  "bet_weights_not_updatable",
  "bet_minimize",
]) {
  test(`ablation.json ${name}`, () => {
    const s = pick.get(name);
    const { params: _, ...input } = modelInput(s.input.model);
    const e = s.expected;
    const got = chooseBet({
      ...input,
      cells: Object.entries(s.input.outcomes as Record<string, any>).map(([f, o]) => cellOf(f, o)),
    });
    assert.equal(!got.weightsKept, e.weightsUpdated, "权重是否更新");
    const names = Object.keys(e.weights);
    assert.deepEqual(modelOf(s.input.model).components, names.map(componentName), "成分名单");
    closeArray(got.weights, names.map((n) => e.weights[n]), TOL, "更新后的权重");
    assert.deepEqual(got.gains.map((g) => g.factor), Object.keys(e.payoffs), "因子");
    got.gains.forEach((g) => close(g.gain, e.payoffs[g.factor], TOL, `${g.factor} 方向收益`));
    assert.equal(got.factor, e.picked, "押的因子");
    close(got.predictedGain, e.predictedGain, TOL, "预测收益");
    assert.equal(got.decision, BET[e.status], "押不押");
    assert.deepEqual(got.offBetter.map((o) => o.factor), e.negativeFactors.map((n: any) => n.factor), "关掉反而更好的格");
    got.offBetter.forEach((o, i) => close(o.better, e.negativeFactors[i].offBetterBy, TOL, `${o.factor} 关掉后好多少`));
  });
}

test("ablation.json 每个场景都已覆盖（value_cost_pure 除外，见下）", () => {
  // value_cost_pure 直接给 ĝ₀ 和 VOI 算 R̂、m′、价值、成本；rules 只在 evaluateCombos 里从模型算这些，
  // 没有接受 (ĝ₀, VOI) 的对外函数，无法映射。R̂、m′、价值、成本已在 select_* 场景里逐组合核对，
  // ĝ₀ < 0 时成本为 0 由 select_negativeG0 核对。
  pick.get("value_cost_pure");
  assert.deepEqual(pick.unused(), []);
});