InfoRSI/src/rules/ablation.ts

import type { Direction, Measured, Range } from "../vocab/index.ts";
import { compareIds, improvement, INDISPENSABLE_SUPPORT, INTERACTION_CUT } from "./basics.ts";
import { bitWidth, covers, disagreement, surprise, worldBits } from "./bits.ts";
import { ablationCellMixture, mixtureRange } from "./distribution.ts";
import { adoptionLine } from "./forward.ts";

// 方法.md 6.6 消融结算，4.3 ⑬ 列格、⑰ 结算与"消融中发现"。

/**
 * 一格的结局（6.6"一格的几种结局"）：
 * - valid：有效分数（含合法化后的分数），只有它进标签的求和、判确认与否定、记比特和惊讶；
 * - insufficient：证据不足：启动了（扣了次数），但跑崩、超时、没有产物、修不了或判为违规；
 *   或者消融核验不通过、重做次数用完，这一格不启动、不扣次数；
 * - untested：没测；
 * - cannotRun：跑不起来（不启动、不扣次）；
 * - joint：一起关，off 是实际关掉的全部因子；score 是那次运行的有效分数，没有时为 null。
 */
export type CellOutcome =
  | { kind: "valid"; score: number }
  | { kind: "insufficient" }
  | { kind: "untested" }
  | { kind: "cannotRun" }
  | { kind: "joint"; off: readonly string[]; score: number | null };

export interface AblationCell {
  factor: string;
  outcome: CellOutcome;
}

/** 一条解释：作者，以及对每一格给的三个数（关掉这个因子后的分数，原值）。 */
export interface AblationExplanation {
  id: string;
  author: string;
  cells: ReadonlyArray<{ factor: string; range: Range }>;
}

/** 列格（4.3 ⑬）：被采纳改法的每个因子各一格"只关掉它"，按因子编号排。 */
export function listCells(factors: readonly string[]): string[] {
  return [...factors].sort(compareIds);
}

/**
 * 总提升 G = f(完整) − f(基线)（6.6"撑几成"，按方向折算）。被消融的一定是被采纳的改法，
 * 所以 G 不小于采纳线；G ≤ 0 说明上游选错了改法，直接报错。
 */
export function totalGain(baseline: number, full: number, direction: Direction): number {
  const gain = improvement(baseline, full, direction);
  if (!(gain > 0)) {
    throw new Error(
      `消融结算（方法.md 6.6"撑几成"）：总提升 G = ${gain} ≤ 0（基线 ${baseline}，完整 ${full}，方向 ${direction}）；` +
        `被消融的必须是被采纳的改法`,
    );
  }
  return gain;
}

/** 撑几成（6.6）= [f(完整) − f(关掉 p)] ÷ G，按方向折算，不截断。 */
export function supportOf(baseline: number, full: number, off: number, direction: Direction): number {
  return improvement(off, full, direction) / totalGain(baseline, full, direction);
}

/**
 * 配合关系标签代号（6.6 表，从上往下判，命中即停）：
 * offBetter 有因子关掉反而更好；incomplete 不全；indispensable 缺一不可；complementary 互补；
 * substitute 替代；independent 各管各。中文和"给下一轮的话"由 materials 写。
 */
export type SynergyLabel = "offBetter" | "incomplete" | "indispensable" | "complementary" | "substitute" | "independent";

/**
 * 配合关系标签（6.6）。supports 是每个因子"只关它、有效分数"的撑几成，没有这样的格时为 null。
 * 每个因子都有有效格时给出 Σ撑 和相对交互量 i = Σ撑 − 1，否则这两个数为 null。
 */
export function synergyLabel(supports: ReadonlyArray<number | null>): {
  label: SynergyLabel;
  supportSum: number | null;
  interaction: number | null;
} {
  const complete = supports.every((s) => s !== null);
  const supportSum = complete ? (supports as number[]).reduce((t, s) => t + s, 0) : null;
  const interaction = supportSum === null ? null : supportSum - 1;
  const label: SynergyLabel = supports.some((s) => s !== null && s < 0)
    ? "offBetter"
    : !complete
      ? "incomplete"
      : (supports as number[]).every((s) => s > INDISPENSABLE_SUPPORT)
        ? "indispensable"
        : interaction! > INTERACTION_CUT
          ? "complementary"
          : interaction! < -INTERACTION_CUT
            ? "substitute"
            : "independent";
  return { label, supportSum, interaction };
}

/**
 * "一起关"是对称的（6.6）：设置由实际关掉的因子集合决定，关 p 连带关 q 和关 q 连带关 p 是同一种设置。
 * 给出每一格的设置键，以及同一次消融里因子编号更小、设置相同的第一格（有它就沿用它的结果，不重跑、不重复扣次）。
 */
export function sharedSettings(
  cells: ReadonlyArray<{ factor: string; off: readonly string[] }>,
): Array<{ factor: string; key: string; sameAs: string | null }> {
  const first = new Map<string, string>();
  return [...cells]
    .sort((a, b) => compareIds(a.factor, b.factor))
    .map((c) => {
      const key = [...new Set(c.off)].sort(compareIds).join("+");
      const sameAs = first.get(key) ?? null;
      if (sameAs === null) first.set(key, c.factor);
      return { factor: c.factor, key, sameAs };
    });
}

export type ExplanationVerdict = "confirmed" | "denied";

/**
 * 确认与否定（6.6）：每条解释在每个有效格上，实测落在它自己的 [q1, q9] 里（含两端）记确认，否则记否定。
 * 解释按编号排，格按因子编号排，只列有效格。
 */
export function judgeExplanations(
  explanations: readonly AblationExplanation[],
  cells: readonly AblationCell[],
): Array<{ id: string; verdicts: Array<{ factor: string; verdict: ExplanationVerdict }> }> {
  const valid = [...cells].sort((a, b) => compareIds(a.factor, b.factor)).filter((c) => c.outcome.kind === "valid");
  return [...explanations]
    .sort((a, b) => compareIds(a.id, b.id))
    .map((e) => ({
      id: e.id,
      verdicts: valid.map((c) => {
        const range = e.cells.find((x) => x.factor === c.factor)!.range;
        const score = (c.outcome as { score: number }).score;
        return { factor: c.factor, verdict: covers(range, score) ? ("confirmed" as const) : ("denied" as const) };
      }),
    }));
}

/**
 * 现有测法分不开（6.6）：在所有有效格上确认与否定序列完全相同的解释归为一组，只列两条以上的组。
 * 没有有效格时不标。组内按解释编号排，组之间按第一条的编号排。
 */
export function inseparableGroups(judgments: ReturnType<typeof judgeExplanations>): string[][] {
  if (judgments.length === 0 || judgments[0]!.verdicts.length === 0) return [];
  const groups = new Map<string, string[]>();
  for (const j of judgments) {
    const key = j.verdicts.map((v) => v.verdict).join(",");
    groups.set(key, [...(groups.get(key) ?? []), j.id]);
  }
  return [...groups.values()]
    .filter((g) => g.length > 1)
    .map((g) => g.sort(compareIds))
    .sort((a, b) => compareIds(a[0]!, b[0]!));
}

/**
 * 消融中发现（4.3 ⑰）：有效格和有分数的"一起关"格里，比当前基线（刚被采纳的完整代码，B = f(完整)）
 * 的改进严格大于采纳线的格过线；几格过线取分数最好的，同分取因子编号小的。
 */
export function ablationFinding(input: {
  full: number;
  knownBest: number | null;
  epsilon: number;
  direction: Direction;
  cells: readonly AblationCell[];
}): {
  line: number;
  quarter: number | null;
  passed: Array<{ factor: string; score: number; improvement: number }>;
  chosen: string | null;
} {
  const { line, quarter } = adoptionLine({ ...input, baseline: input.full });
  const passed = [...input.cells]
    .sort((a, b) => compareIds(a.factor, b.factor))
    .flatMap((c) => {
      const o = c.outcome;
      const score = o.kind === "valid" || (o.kind === "joint" && o.score !== null) ? o.score : null;
      if (score === null) return [];
      const gain = improvement(input.full, score, input.direction);
      return gain > line ? [{ factor: c.factor, score, improvement: gain }] : [];
    });
  let chosen: (typeof passed)[number] | null = null;
  for (const p of passed) if (chosen === null || p.improvement > chosen.improvement) chosen = p;
  return { line, quarter, passed, chosen: chosen?.factor ?? null };
}

/** 一次消融结算的输入（4.3 ⑰）。baseline 是本批基线 b，full 是被采纳的完整代码的分数。 */
export interface SettleInput {
  baseline: number;
  full: number;
  direction: Direction;
  epsilon: number;
  knownBest: number | null;
  lambda: number;
  /** 被采纳改法的因子编号：按它列格、贴标签，每个因子都要有一格。 */
  factors: readonly string[];
  cells: readonly AblationCell[];
  explanations: readonly AblationExplanation[];
}

/**
 * 一次消融的结算（4.3 ⑰、6.6、6.9）：每格的撑几成、合并范围，有效格的世界比特、惊讶、分歧量（格宽 ε/100，
 * 从 b 起切）；配合关系标签；每条解释的确认与否定；分不开的组；消融中发现。押的因子另由 chooseBet 算。
 * 格按因子清单列（4.3 ⑬ 每个因子一格）：清单里的因子没有格、或者有清单外的格，都是上游的错，直接报错，
 * 不把漏掉的因子当成"不全"。
 */
export function settleAblation(input: SettleInput): {
  gain: number;
  cells: Array<{
    factor: string;
    outcome: CellOutcome["kind"];
    support: number | null;
    merged: Range;
    worldBits: number | null;
    surprise: { cell: number; prior: number; bits: Measured } | null;
    disagreement: number | null;
  }>;
  label: ReturnType<typeof synergyLabel>;
  judgments: ReturnType<typeof judgeExplanations>;
  inseparable: string[][];
  finding: ReturnType<typeof ablationFinding>;
} {
  const gain = totalGain(input.baseline, input.full, input.direction);
  const width = bitWidth(input.epsilon);
  const factors = listCells(input.factors);
  for (const c of input.cells) {
    if (!factors.includes(c.factor)) {
      throw new Error(`消融结算（方法.md 4.3 ⑬ 列格）：格 ${c.factor} 不在被采纳改法的因子清单 ${factors.join("、")} 里`);
    }
  }
  const cells = factors
    .map((f) => {
      const c = input.cells.find((x) => x.factor === f);
      if (c === undefined) throw new Error(`消融结算（方法.md 4.3 ⑬ 列格）：因子 ${f} 没有对应的格`);
      return c;
    })
    .map((c) => {
      const entries = input.explanations.map((e) => ({
        id: e.id,
        author: e.author,
        range: e.cells.find((x) => x.factor === c.factor)!.range,
      }));
      const mixture = ablationCellMixture(entries, input.lambda);
      const o = c.outcome;
      const valid = o.kind === "valid";
      return {
        factor: c.factor,
        outcome: o.kind,
        support: valid ? improvement(o.score, input.full, input.direction) / gain : null,
        merged: mixtureRange(mixture),
        worldBits: valid ? worldBits(mixture, input.baseline, width) : null,
        surprise: valid ? surprise(mixture, input.baseline, width, o.score) : null,
        disagreement: valid ? disagreement(entries, input.baseline, width) : null,
      };
    });
  const judgments = judgeExplanations(input.explanations, input.cells);
  return {
    gain,
    cells,
    label: synergyLabel(cells.map((c) => c.support)),
    judgments,
    inseparable: inseparableGroups(judgments),
    finding: ablationFinding(input),
  };
}