import type { Direction } from "../vocab/index.ts";
import type { AblationExplanation } from "./ablation.ts";
import { totalGain } from "./ablation.ts";
import {
  combinations,
  compareIds,
  epsilonFloor,
  GRID_COMBINATION_CAP,
  GRID_SUBDIVISIONS,
  improvement,
} from "./basics.ts";
import {
  ablationCellMixture,
  below,
  density,
  integrateUpperTail,
  knotsMean,
  mapKnots,
  type Knots,
} from "./distribution.ts";

// 方法.md 6.7：方向收益、成分与权重、信息价值、价值减成本。选格规则和押的因子在 cells.ts。
//
// 成分、权重更新、精确积分、部分格网格、价值减成本这套引擎消融（6.7）和诊断（6.12）共用，只差"实测 → 收益"的换算：
// - 消融：一律先在"撑几成"刻度 s(x) = [f(完整) − x] ÷ G（按方向折算）上算，最后乘剩余空间 h：
//   收益_p = h × s_p，h 对所有因子相同且为正，所以信息价值、期望收益、ĝ₀ 都等于 h × (撑几成刻度上的同一个量)；
//   不设"都不走"的选项。
// - 诊断（diagnosis.ts）：收益 g(x) = 改进量 x − b，h 取 1，另有收益为 0 的"都不走"，即 floor = 0。
// 例外是部分格网格的切段：在分数原值上取折点、切小段，再把中点换成收益刻度（见 gridOf）。

/** 引擎要的成分模型（消融和诊断各自建出来）。下标"项"是因子（消融）或实验（诊断）。 */
export interface ValueModel {
  /** 成分起始权重。 */
  weights: number[];
  /** rawKnots[项][成分]：分数原值上的折线。 */
  rawKnots: Knots[][];
  /** scaledKnots[项][成分]：收益刻度上的折线。 */
  scaledKnots: Knots[][];
  /** expected[项][成分]：这个成分下这一项在收益刻度上的期望。 */
  expected: number[][];
  /** prior[项]：按起始权重平均的期望。 */
  prior: number[];
  /** 收益 = space × 收益刻度上的值：消融是剩余空间 h，诊断是 1。 */
  space: number;
  /** "都不走"的收益：诊断是 0，消融没有这个选项，写 null。 */
  floor: number | null;
  /** 分数原值 → 收益刻度。 */
  scale: (x: number) => number;
}

/** 开放预留成分的名字（成分名单里排在最后）。 */
export const OPEN_RESERVE = "openReserve";

/** 建成分模型要的输入（6.7）。baseline 是本批基线 b，full 是被采纳的完整代码的分数。 */
export interface AblationModelInput {
  baseline: number;
  full: number;
  direction: Direction;
  knownBest: number | null;
  epsilon: number;
  lambda: number;
  /** 被采纳改法的因子编号。 */
  factors: readonly string[];
  /** 每条解释对每一格都给了三个数。 */
  explanations: readonly AblationExplanation[];
}

export interface AblationModel {
  /** 因子编号，按编号排；下文"格的下标"都按这个次序。 */
  factors: string[];
  /** 成分：各条解释的编号（按编号排），最后是开放预留。 */
  components: string[];
  /** 起始权重（6.7"成分与权重"）：解释合起来 1 − λ（研究者等权、同一人平分），开放预留 λ。 */
  weights: number[];
  /** 被采纳的完整代码的分数 f(完整) 和方向：分数原值换算成撑几成要用。 */
  full: number;
  direction: Direction;
  /** 总提升 G。 */
  gain: number;
  /** 剩余空间 h。 */
  space: number;
  /** rawKnots[格][成分]：分数原值上的折线（开放预留是这一格自己的均匀段）。 */
  rawKnots: Knots[][];
  /** supportKnots[格][成分]：撑几成刻度上的折线。 */
  supportKnots: Knots[][];
  /** expected[格][成分]：这个成分下这一格撑几成的期望。 */
  expected: number[][];
  /** prior[格]：按起始权重平均的期望撑几成；乘 h 就是 E(收益_p)。 */
  prior: number[];
}

/**
 * 剩余空间 h（6.7"方向收益"）：有已知最好时 max(K* − f(完整), ε/100)（按方向折算），没有时 1。
 */
export function remainingSpace(full: number, knownBest: number | null, epsilon: number, direction: Direction): number {
  return knownBest === null ? 1 : Math.max(improvement(full, knownBest, direction), epsilonFloor(epsilon));
}

/** 分数原值 x → 撑几成 [f(完整) − x] ÷ G（按方向折算）。 */
function supportScale(full: number, direction: Direction, gain: number): (x: number) => number {
  return (x) => improvement(x, full, direction) / gain;
}

/** 消融模型在引擎眼里的样子：收益刻度是撑几成，没有"都不走"。 */
export function ablationView(model: AblationModel): ValueModel {
  return {
    weights: model.weights,
    rawKnots: model.rawKnots,
    scaledKnots: model.supportKnots,
    expected: model.expected,
    prior: model.prior,
    space: model.space,
    floor: null,
    scale: supportScale(model.full, model.direction, model.gain),
  };
}

/**
 * 成分模型（6.7"成分与权重"）：每条解释一个成分，开放预留一个成分；每个成分对每一格有一个分布，
 * 解释用它自己的三个数补尾巴，开放预留用这一格的均匀分布（6.3"消融"）。同一成分内各格互相独立。
 */
export function ablationModel(input: AblationModelInput): AblationModel {
  const gain = totalGain(input.baseline, input.full, input.direction);
  const factors = [...input.factors].sort(compareIds);
  const explanations = [...input.explanations].sort((a, b) => compareIds(a.id, b.id));
  const mixtures = factors.map((f) =>
    ablationCellMixture(
      explanations.map((e) => ({ id: e.id, author: e.author, range: e.cells.find((c) => c.factor === f)!.range })),
      input.lambda,
    ),
  );
  const toSupport = supportScale(input.full, input.direction, gain);
  const rawKnots = mixtures.map((m) => m.map((c) => c.knots));
  const supportKnots = rawKnots.map((cell) => cell.map((k) => mapKnots(k, toSupport, input.direction === "maximize")));
  const weights = mixtures[0]!.map((c) => c.weight);
  const expected = supportKnots.map((cell) => cell.map(knotsMean));
  return {
    factors,
    components: [...explanations.map((e) => e.id), OPEN_RESERVE],
    weights,
    full: input.full,
    direction: input.direction,
    gain,
    space: remainingSpace(input.full, input.knownBest, input.epsilon, input.direction),
    rawKnots,
    supportKnots,
    expected,
    prior: expected.map((row) => row.reduce((s, e, j) => s + weights[j]! * e, 0)),
  };
}

/** 测 k 格时每格的小段数 m：满足 m^k ≤ 4,000,000 的最大整数（6.7"信息价值"），用整数运算核对。 */
export function gridSize(k: number): number {
  let m = Math.floor(GRID_COMBINATION_CAP ** (1 / k));
  while (m ** k > GRID_COMBINATION_CAP) m -= 1;
  while ((m + 1) ** k <= GRID_COMBINATION_CAP) m += 1;
  return m;
}

/**
 * 一格的网格（6.7"只测一部分格"）：取这一格全部成分的折点（含开放预留两端）去重排序，段数 = 折点数 − 1；
 * s = min(20, ⌊m ÷ 段数⌋)，s ≥ 1 时每段等分 s 份，段数超过 m 时整格等宽切 m 份。
 * 每个小段的代表点是中点；成分 j 在小段 [左, 右] 的把握 = F_j(右) − F_j(左)，各小段把握加起来是 1。
 */
export function cellGrid(knots: readonly Knots[], k: number): { mids: number[]; mass: number[][] } {
  const brk = [...new Set(knots.flatMap((kk) => kk.map(([x]) => x)))].sort((a, b) => a - b);
  const segs = brk.length - 1;
  const m = gridSize(k);
  const s = Math.min(GRID_SUBDIVISIONS, Math.floor(m / segs));
  const edges: number[] = [];
  if (s >= 1) {
    for (let i = 0; i < segs; i++) {
      const a = brk[i]!;
      const step = (brk[i + 1]! - a) / s;
      for (let j = 0; j < s; j++) edges.push(a + j * step);
    }
  } else {
    const step = (brk[segs]! - brk[0]!) / m;
    for (let j = 0; j < m; j++) edges.push(brk[0]! + j * step);
  }
  edges.push(brk[segs]!);
  const mids: number[] = [];
  for (let i = 0; i + 1 < edges.length; i++) mids.push((edges[i]! + edges[i + 1]!) / 2);
  const mass = knots.map((kk) => {
    const cdf = edges.map((x) => below(kk, x));
    return mids.map((_, i) => cdf[i + 1]! - cdf[i]!);
  });
  return { mids, mass };
}

/**
 * 模型里第 c 项的网格（6.7"只测一部分格"），测 k 项时用。在分数原值上取折点、切小段：方法说"每格在自己的开放预留范围内
 * 按折点切段"，开放预留和折点都是在分数原值上定的（尾巴按 q1 − (q5 − q1)、q9 + (q9 − q5)，预留按 L − W/2、U + W/2
 * 的浮点原式），去重只合并浮点上完全相等的值。先换到收益刻度再去重，相差一个末位的两个原值折点可能换算后变成同一个数，
 * 段数和 s 就会跟原值上的不一样。切好后只把小段中点换算成收益刻度；把握与刻度无关。
 */
export function gridOf(model: ValueModel, c: number, k: number): { mids: number[]; mass: number[][] } {
  const { mids, mass } = cellGrid(model.rawKnots[c]!, k);
  return { mids: mids.map(model.scale), mass };
}

/** 消融模型里第 c 格的网格，中点换算成撑几成（6.7"只测一部分格"）。 */
export function measuredGrid(model: AblationModel, c: number, k: number): { mids: number[]; mass: number[][] } {
  return gridOf(ablationView(model), c, k);
}

/**
 * 只测 S 里的项时的 E_{S 的结果}[max(floor, max_p E(收益_p | S 的结果))]（6.7"信息价值"、6.12），在收益刻度上，
 * 没有"都不走"（floor 为 null）时不和 floor 取大。
 * 对测到的项的全部小段组合求和：成分 j 的联合把握 = 起始权重 × 各测到项在所选小段的把握；它们的和 marg 是这个组合的把握，
 * 按它归一就是更新后的权重。网格见 gridOf。测到的项取小段中点（已换成收益刻度），没测的项用更新后的权重平均各成分的期望，取各项里最大的。
 */
function expectedBestGrid(model: ValueModel, S: readonly number[]): number {
  const k = S.length;
  const J = model.weights.length;
  const floor = model.floor ?? -Infinity;
  const grids = S.map((c) => gridOf(model, c, k));
  const rest = model.prior.map((_, c) => c).filter((c) => !S.includes(c));
  const partial = Array.from({ length: k + 1 }, () => new Float64Array(J));
  for (let j = 0; j < J; j++) partial[0]![j] = model.weights[j]!;
  let total = 0;
  const walk = (level: number, seen: number): void => {
    const cur = partial[level]!;
    if (level === k) {
      let marg = 0;
      for (let j = 0; j < J; j++) marg += cur[j]!;
      if (marg === 0) return;
      let best = Math.max(seen, floor);
      for (const c of rest) {
        const e = model.expected[c]!;
        let num = 0;
        for (let j = 0; j < J; j++) num += cur[j]! * e[j]!;
        if (num / marg > best) best = num / marg;
      }
      total += best * marg;
      return;
    }
    const g = grids[level]!;
    const next = partial[level + 1]!;
    for (let i = 0; i < g.mids.length; i++) {
      for (let j = 0; j < J; j++) next[j] = cur[j]! * g.mass[j]![i]!;
      walk(level + 1, Math.max(seen, g.mids[i]!));
    }
  };
  walk(0, -Infinity);
  return total;
}

/**
 * 全部项都测时的 E[max(floor, max_p 收益_p)]（6.7"信息价值"、6.12）：同一成分内"各项最大值低于 t"的把握 = 各项把握连乘，
 * 期望 = 下端 + ∫(1 − 连乘) dt，逐段多项式精确积分；再按起始权重相加。下端：有"都不走"时是 floor（t 低于它时
 * max(floor, ·) 一定不低于 t，被积函数是 1），没有时是各项折线的最低点。
 */
function expectedBestAll(model: ValueModel): number {
  let v = 0;
  for (let j = 0; j < model.weights.length; j++) {
    const lines = model.scaledKnots.map((cell) => cell[j]!);
    const lo = model.floor ?? Math.min(...lines.map((k) => k[0]![0]));
    v += model.weights[j]! * (lo + integrateUpperTail(lines.map((knots) => [{ weight: 1, knots }]), lo));
  }
  return v;
}

/** max(floor, max_p E(收益_p))：不测时能拿到的最好期望，在收益刻度上。 */
function bestPrior(model: ValueModel): number {
  return Math.max(model.floor ?? -Infinity, ...model.prior);
}

/** 用网格算的信息价值（不论测几项都走网格），S 是项的下标。 */
export function valueByGrid(model: ValueModel, S: readonly number[]): number {
  return model.space * (expectedBestGrid(model, [...S].sort((a, b) => a - b)) - bestPrior(model));
}

/** 全部项都测时的信息价值（一维精确积分）。 */
export function valueExact(model: ValueModel): number {
  return model.space * (expectedBestAll(model) - bestPrior(model));
}

/**
 * 信息价值 VOI(S) = E_{S 的结果}[max(floor, max_p E(收益_p | S 的结果))] − max(floor, max_p E(收益_p))（6.7、6.12），
 * S 是项的下标；全部项都测走一维精确积分，只测一部分走网格，空集是 0。不截断。
 */
export function valueOf(model: ValueModel, S: readonly number[]): number {
  if (S.length === 0) return 0;
  return S.length === model.prior.length ? valueExact(model) : valueByGrid(model, S);
}

/** 网格算法算的信息价值，S 是格的下标（不论测几格都走网格；用来和精确算法互相对照）。 */
export function informationValueGrid(model: AblationModel, S: readonly number[]): number {
  return valueByGrid(ablationView(model), S);
}

/** 全部格都测时的信息价值（一维精确积分）。 */
export function informationValueExact(model: AblationModel): number {
  return valueExact(ablationView(model));
}

/** 消融的信息价值 VOI(S)（6.7），S 是格的下标。 */
export function informationValue(model: AblationModel, S: readonly number[]): number {
  return valueOf(ablationView(model), S);
}

/**
 * 成分新权重（6.8、6.12）：起始权重 × 各有效项（scores 里不是 null 的）上这个成分在实测处的密度（分数原值上，连乘），
 * 再归一；实测落在所有成分之外、权重之和为 0 时保持起始权重。scores 与项的下标对齐，没有有效实测的项为 null。
 */
export function updatedWeights(
  model: ValueModel,
  scores: ReadonlyArray<number | null>,
): { weights: number[]; kept: boolean } {
  const raw = model.weights.map((w, j) =>
    scores.reduce<number>((v, x, c) => (x === null ? v : v * density(model.rawKnots[c]![j]!, x)), w),
  );
  const sum = raw.reduce((t, v) => t + v, 0);
  const kept = sum === 0;
  return { weights: kept ? [...model.weights] : raw.map((v) => v / sum), kept };
}

/** 第 c 项的方向收益（6.8、6.12）：有有效实测 x 就用实测换算，否则用 weights 对各成分的期望加权平均。 */
export function itemGain(model: ValueModel, c: number, weights: readonly number[], x: number | null): number {
  return x !== null
    ? model.space * model.scale(x)
    : model.space * model.expected[c]!.reduce((t, e, j) => t + weights[j]! * e, 0);
}

/** 一个测格组合的价值与成本（6.7"价值减成本"）。 */
export interface ComboValue {
  /** 组合里的因子，按编号排。 */
  factors: string[];
  voi: number;
  /** R̂ = ⌈(余额 − |S|) ÷ 每批上限⌉。 */
  rHat: number;
  /** m′ = 1 + r + … + r^(R̂−1)，R̂ = 0 时是 0。 */
  mPrime: number;
  /** 价值 = u × m′ × VOI。 */
  value: number;
  /** 成本 = 2 max(0, ĝ₀) × [(K + |S|)/(K + |S| + 1) − K/(K + 1)]。 */
  cost: number;
  /** 净值 = 价值 − 成本。 */
  net: number;
}

export interface ComboParams {
  balance: number;
  batchCap: number;
  /** 测出的知识会被用上的概率 u。 */
  u: number;
  /** 衰减 r。 */
  decay: number;
}

/** m′ = 1 + r + … + r^(R̂−1)，R̂ = 0 时是 0（6.7"价值减成本"）；消融和诊断共用。 */
export function usePotential(rHat: number, decay: number): number {
  let mPrime = 0;
  for (let j = 0; j < rHat; j++) mPrime += decay ** j;
  return mPrime;
}

/** 成本 = 2 max(0, ĝ₀) × [(K + n)/(K + n + 1) − K/(K + 1)]，n 是这次占用的运行次数（6.7、6.12）；消融和诊断共用。 */
export function runCost(g0: number, batchCap: number, runs: number): number {
  const K = batchCap;
  return 2 * Math.max(0, g0) * ((K + runs) / (K + runs + 1) - K / (K + 1));
}

/**
 * 列出格数 1..min(格数, 余额 − 1) 的全部非空组合并算价值、成本、净值（6.7"选格"第 1 条）。余额为 0 或 1 时一个都没有。
 * 组合按格数从少到多、同格数按因子编号字典序列出。ĝ₀ = h × max_p 期望撑几成，原值另给，成本里取 max(0, ĝ₀)。
 */
export function evaluateCombos(model: AblationModel, params: ComboParams): { g0: number; combos: ComboValue[] } {
  const g0 = model.space * Math.max(...model.prior);
  const K = params.batchCap;
  const combos: ComboValue[] = [];
  for (let size = 1; size <= Math.min(model.factors.length, params.balance - 1); size++) {
    for (const S of combinations(model.factors.length, size)) {
      const voi = informationValue(model, S);
      const rHat = Math.ceil((params.balance - size) / K);
      const mPrime = usePotential(rHat, params.decay);
      const value = params.u * mPrime * voi;
      const cost = runCost(g0, K, size);
      combos.push({ factors: S.map((c) => model.factors[c]!), voi, rHat, mPrime, value, cost, net: value - cost });
    }
  }
  return { g0, combos };
}