import { judgeSameChange, type DedupFactor } from "../../jev/index.ts";
import { Board, candidateId, freezeProtocol, type DedupOutput } from "../../materials/index.ts";
import { compareIds, decisionSeed, runBalance, runBatchRule, type BatchInput, type BatchRuleName } from "../../rules/index.ts";
import { protocolBody } from "../../submissions/index.ts";
import type { JevJudgment, Json } from "../../vocab/index.ts";
import { freshOf, poolOf, stageIdOf } from "../next-step.ts";
import type { Study } from "../study.ts";

// 第 r 轮里程序自己做的三步（方法.md 4.3 ③⑥⑧）：判重合并、冻结协议、选批。每步写一条记录（冻结协议一个事务），
// 写下之前被打断就从账本重新算一遍；写下之后恢复时照原样用，不重算（方法.md 第九部分）。

/**
 * 同一轮里几个要随机的决定各用一个"版本"算种子（方法.md 第九部分）：选批和选格都用随机规则时，
 * 两次抽取不会用同一串随机数。选格由 flow 第 4 部分接上。
 */
const SEED_VERSIONS = { batch: 1, cells: 2 } as const;
export { SEED_VERSIONS };

/** 一对提案（或提案与留池候选）的 Jev 判断，原样存进判重合并决定的输入。 */
export type SamePair = { a: string; b: string; judgment: JevJudgment };

/**
 * ③ 判重合并：本轮新提案两两问一次 Jev，每条新提案再和每个留池候选问一次（留池候选看它冻结协议里的因子），
 * 按顺序发。全部问完才写决定：中途被打断就整轮重问，判重不扣运行次数。
 */
export async function dedupe(s: Study, round: number): Promise<void> {
  const board = new Board(s.ledger.records());
  const proposals = board.published("proposal").filter((p) => p.round === round).sort((a, b) => compareIds(a.item!, b.item!));
  const carried = board.carried(round);
  const model = s.frozen.config.models.jev;
  const factorsOf = new Map<string, readonly DedupFactor[]>([
    ...proposals.map((p) => [p.item!, p.content.factors] as const),
    ...carried.map((c) => [c, board.protocol(c).factors] as const),
  ]);
  const asks = [
    ...proposals.flatMap((p, i) => proposals.slice(i + 1).map((q) => [p.item!, q.item!] as const)),
    ...proposals.flatMap((p) => carried.map((c) => [p.item!, c] as const)),
  ];
  const pairs: SamePair[] = [];
  for (const [a, b] of asks) pairs.push({ a, b, judgment: await judgeSameChange(model, factorsOf.get(a)!, factorsOf.get(b)!) });
  const output = groupSameChange(round, proposals.map((p) => p.item!), carried, pairs);
  // SAFETY: JevJudgment 和 DedupOutput 都是纯 Json 值；账本在写入时校验 Json。
  s.ledger.append("decision", {
    decision: "dedup",
    round,
    rule: "jevSameChange",
    ruleVersion: 1,
    params: { jevModel: model },
    seed: null,
    input: { pairs } as unknown as Json,
    output: output as unknown as Json,
  });
}

/**
 * 判为相同的连成一组（监督者定：A 和 B 相同、B 和 C 相同，A、B、C 就是一个候选；方法没写这种情况）。
 * 一组里有留池候选，整组并进它（方法.md 4.3 ③）；一组连到两个留池候选按内部错误停下（真实运行里没出现过）。
 * 其余每组是一个新候选，按组里编号最小的原提案排序，编号用 candidateId。每个留池候选都列出，没并进提案的是空数组。
 */
export function groupSameChange(round: number, proposals: readonly string[], carried: readonly string[], pairs: readonly SamePair[]): DedupOutput {
  const parent = new Map([...proposals, ...carried].map((x) => [x, x]));
  const root = (x: string): string => {
    while (parent.get(x) !== x) x = parent.get(x)!;
    return x;
  };
  for (const p of pairs) if (p.judgment.yes) parent.set(root(p.a), root(p.b));
  const groups = new Map<string, string[]>();
  for (const x of [...proposals, ...carried]) groups.set(root(x), [...(groups.get(root(x)) ?? []), x]);
  const into = new Map(carried.map((c) => [c, [] as string[]]));
  const fresh: string[][] = [];
  for (const members of groups.values()) {
    const old = members.filter((m) => into.has(m));
    const added = members.filter((m) => !into.has(m)).sort(compareIds);
    if (old.length > 1) {
      throw new Error(`判重合并：提案 ${added.join("、")} 把留池候选 ${old.join("、")} 连成了一组，方法没写该并进哪一个（真实运行里没出现过），按内部错误停下`);
    }
    if (old.length === 1) into.get(old[0]!)!.push(...added);
    else fresh.push(added);
  }
  fresh.sort((a, b) => compareIds(a[0]!, b[0]!));
  return {
    fresh: fresh.map((ids, i) => ({ candidate: candidateId(round, i + 1), proposals: ids })),
    carried: [...into].sort(([a], [b]) => compareIds(a, b)).map(([candidate, ids]) => ({ candidate, proposals: ids })),
  };
}

/**
 * ⑥ 冻结协议：回应阶段关闭后，每个新候选取它最后一份回应或统一，用 protocolBody 拼主体，freezeProtocol 补上候选编号、
 * 轮次、本批基线、合并了哪些原提案。全部候选在同一个事务里写：materials 把冻结协议的出现当作"已公开"。
 * 留池候选沿用原来的冻结协议（方法.md 4.3 ①），不重写。
 */
export function freezeProtocols(s: Study, round: number): void {
  const board = new Board(s.ledger.records());
  const baseline = board.baselineOf(round);
  s.ledger.transaction(() => {
    for (const { candidate, proposals } of freshOf(board, round)) {
      const answer = board.answerOf(candidate);
      const body =
        answer.kind === "response"
          ? protocolBody({ kind: "response", content: answer.content, original: board.proposal(proposals[0]!).content })
          : protocolBody({ kind: "unification", content: answer.content });
      // SAFETY: FrozenProtocol 是纯 Json 值；账本在写入时校验 Json。
      s.ledger.append("protocol", { candidate, protocol: freezeProtocol(body, candidate, round, baseline, proposals) as unknown as Json });
    }
  });
}

/**
 * ⑧ 选批（方法.md 6.4）：池里每个候选、各人最后一份第二次预测，加本批基线分数、方向、ε、λ、每批上限、余额和种子，
 * 交给步骤清单里 selectBatch 那一步的规则，写一条选批决定。
 */
export function selectBatch(s: Study, round: number): void {
  const board = new Board(s.ledger.records());
  const { config, facts, steps } = s.frozen;
  const stageId = stageIdOf(round, "secondPrediction");
  const predictions = board.published("secondPrediction").filter((p) => p.stageId === stageId);
  const range = (researcher: string, candidate: string) => predictions.find((p) => p.author === researcher && p.target === candidate)!.content.range;
  const seed = decisionSeed({ studyId: config.studyId, round, version: SEED_VERSIONS.batch });
  const input: BatchInput = {
    candidates: poolOf(board, round).map((id) => ({ id, ranges: board.stage(stageId).roster.map((researcher) => ({ researcher, range: range(researcher, id) })) })),
    baseline: board.baselineScore(board.baselineOf(round)),
    direction: facts.finalMetric.direction,
    epsilon: facts.epsilon,
    lambda: config.flow.openReserve,
    batchCap: config.batchLimit,
    balance: runBalance(config.runs, board.runStates()).balance,
    seed,
  };
  const rule = steps.round.find((x) => x.id === "selectBatch")!.rule!.name as BatchRuleName;
  const { rule: used, ...output } = runBatchRule(rule, input);
  // SAFETY: rules 是纯计算，输入输出都是纯 Json 值；账本在写入时校验 Json。
  s.ledger.append("decision", {
    decision: "batch",
    round,
    rule: used.name,
    ruleVersion: used.version,
    params: output.params,
    seed,
    input: input as unknown as Json,
    output: output as unknown as Json,
  });
}