import type { Direction, Range } from "../vocab/index.ts";
import { combinations, compareIds, epsilonFloor, improvement } from "./basics.ts";
import { forwardBatchSize } from "./budget.ts";
import {
  forwardMergedRange,
  forwardMixtures,
  integrateUpperTail,
  mixtureBelow,
  toImprovement,
  type Mixture,
} from "./distribution.ts";
import { seededSample } from "./random.ts";

// 方法.md 6.4 入池、选批、分数尺读数；6.5 采纳线；4.3 ⑫ 采纳。

/** 选批规则的统一输入：同一份输入可以换选批规则重算。 */
export interface BatchInput {
  /** 本轮的全部候选（入池在规则里判）；ranges 是每位研究者第二次预测的三个数（分数原值），一人一份，带研究者编号。 */
  candidates: ReadonlyArray<{ id: string; ranges: ReadonlyArray<{ researcher: string; range: Range }> }>;
  /** 本批基线分数 b。 */
  baseline: number;
  direction: Direction;
  epsilon: number;
  /** 开放预留 λ。 */
  lambda: number;
  batchCap: number;
  /** 运行次数余额。 */
  balance: number;
  /** 本次决定的种子（decisionSeed）。批内最好期望规则不用它，照样记进记录。 */
  seed: string;
}

/** 选批规则的统一输出，含 6.4"选批记录写入"的全部内容；规则名和版本由登记表补上。 */
export interface BatchDecision {
  seed: string;
  params: { epsilon: number; lambda: number; batchCap: number; balance: number };
  /** 批量 k = min(每批上限, 余额)。 */
  k: number;
  /** 每个候选（按编号排）：大家合起来的范围（分数原值）、"比 b 低 ε 以上"的把握、是否入池。 */
  candidates: Array<{ id: string; merged: Range; lossProbability: number; inPool: boolean }>;
  /** 入池的候选编号（按编号排）。 */
  pool: string[];
  /** 列出的每个组合及其批内最好期望 V（均匀随机规则不列组合，是空数组）。 */
  combos: Array<{ ids: string[]; value: number }>;
  /** 选中的组合（按编号排）。 */
  chosen: string[];
  /** V(S*)。 */
  chosenValue: number;
  /** E[max(0, max_{c∈池} Y_c)] = V(候选池)。 */
  poolValue: number;
  /** 分数尺读数 = V(候选池) − V(S*)，只记录。 */
  scaleReading: number;
}

/**
 * 批内最好期望 V(S) = ∫_0^∞ [1 − ∏_{c∈S} F_c(t)] dt（6.4），F_c 是改进量上的合并分布，候选之间按互不相关，
 * 低于 b 的部分按 0 算。空组合是 0。
 */
export function batchValue(imps: readonly Mixture[]): number {
  return imps.length === 0 ? 0 : integrateUpperTail(imps, 0);
}

/** 入池（6.4）：合并分布里"比 b 低 ε 以上"的把握 P(Y < −ε) 严格小于一半才入池；正好低 ε 不算"低 ε 以上"。 */
export function lossProbability(imp: Mixture, epsilon: number): number {
  return mixtureBelow(imp, -epsilon);
}

/** 两条选批规则共用的部分：合并、入池、合并范围、批量、候选池的 V。 */
function assess(input: BatchInput) {
  const sorted = [...input.candidates].sort((a, b) => compareIds(a.id, b.id));
  const imps = forwardMixtures(sorted, input.lambda).map((m) => toImprovement(m, input.baseline, input.direction));
  const candidates = sorted.map((c, i) => {
    const loss = lossProbability(imps[i]!, input.epsilon);
    return {
      id: c.id,
      merged: forwardMergedRange(imps[i]!, input.baseline, input.direction),
      lossProbability: loss,
      inPool: loss < 0.5,
    };
  });
  const poolIdx = candidates.flatMap((c, i) => (c.inPool ? [i] : []));
  return {
    imps,
    candidates,
    poolIdx,
    k: forwardBatchSize(input.batchCap, input.balance),
    poolValue: batchValue(poolIdx.map((i) => imps[i]!)),
    params: { epsilon: input.epsilon, lambda: input.lambda, batchCap: input.batchCap, balance: input.balance },
  };
}

/**
 * 选批规则"批内最好期望"（6.4）：列出候选池里所有 min(k, 池里个数) 个候选的组合，取 V 最大的；
 * V 完全相等（不带容差）时取编号字典序小的组合（组合按字典序列出，后来者严格更大才替换）。
 */
export function selectBestExpected(input: BatchInput): BatchDecision {
  const a = assess(input);
  const size = Math.min(a.k, a.poolIdx.length);
  const picks = size === 0 ? [] : combinations(a.poolIdx.length, size);
  const combos = picks.map((pick) => {
    const idx = pick.map((p) => a.poolIdx[p]!);
    return { ids: idx.map((i) => a.candidates[i]!.id), value: batchValue(idx.map((i) => a.imps[i]!)) };
  });
  // 池空或批量为 0 时一个组合都没有，S* 是空组合，V = 0。
  let best = combos[0] ?? { ids: [], value: 0 };
  for (const c of combos) if (c.value > best.value) best = c;
  return {
    seed: input.seed,
    params: a.params,
    k: a.k,
    candidates: a.candidates,
    pool: a.poolIdx.map((i) => a.candidates[i]!.id),
    combos,
    chosen: best.ids,
    chosenValue: best.value,
    poolValue: a.poolValue,
    scaleReading: a.poolValue - best.value,
  };
}

/**
 * 选批规则"均匀随机"（第十一部分）：从全部候选里（不看入池门槛）不放回地均匀抽 min(k, 候选数) 个。
 * 入池、合并范围、分数尺读数照算照记；S* 是抽中的组合，可能含池外候选。
 */
export function selectUniformRandom(input: BatchInput): BatchDecision {
  const a = assess(input);
  const ids = a.candidates.map((c) => c.id);
  const chosen = seededSample(ids, Math.min(a.k, ids.length), input.seed);
  const chosenValue = batchValue(chosen.map((id) => a.imps[ids.indexOf(id)]!));
  return {
    seed: input.seed,
    params: a.params,
    k: a.k,
    candidates: a.candidates,
    pool: a.poolIdx.map((i) => a.candidates[i]!.id),
    combos: [],
    chosen,
    chosenValue,
    poolValue: a.poolValue,
    scaleReading: a.poolValue - chosenValue,
  };
}

/** 采纳线的输入（6.5）：B 是当前基线分数，K* 是已知最好（没有时 null）。 */
export interface AdoptionLineInput {
  baseline: number;
  knownBest: number | null;
  epsilon: number;
  direction: Direction;
}

/**
 * 采纳线（6.5）：没有已知最好时是 ε；有时是 min(ε, (K* − B) ÷ 4)（按方向折算），结果低于 ε/100 时取 ε/100，
 * B 已经超过 K*（(K* − B) 为负）时同样取 ε/100。quarter 是 (K* − B) ÷ 4 的原值，没有已知最好时为 null。
 */
export function adoptionLine(input: AdoptionLineInput): { line: number; quarter: number | null } {
  if (input.knownBest === null) return { line: input.epsilon, quarter: null };
  const quarter = improvement(input.baseline, input.knownBest, input.direction) / 4;
  return { line: Math.max(epsilonFloor(input.epsilon), Math.min(input.epsilon, quarter)), quarter };
}

/**
 * 采纳（6.5、4.3 ⑫）：只看实测分数。candidates 是有有效分数的正向候选；最好的那个（按方向，同分取编号小的）
 * 比基线的改进严格大于采纳线才采纳。
 */
export function adopt(
  input: AdoptionLineInput & { candidates: ReadonlyArray<{ id: string; score: number }> },
): {
  line: number;
  quarter: number | null;
  best: { id: string; score: number; improvement: number } | null;
  adopted: string | null;
} {
  const { line, quarter } = adoptionLine(input);
  let best: { id: string; score: number; improvement: number } | null = null;
  for (const c of [...input.candidates].sort((x, y) => compareIds(x.id, y.id))) {
    const gain = improvement(input.baseline, c.score, input.direction);
    if (best === null || gain > best.improvement) best = { id: c.id, score: c.score, improvement: gain };
  }
  return { line, quarter, best, adopted: best !== null && best.improvement > line ? best.id : null };
}