import assert from "node:assert/strict";
import { test } from "node:test";
import {
  adopt,
  adoptionLine,
  batchValue,
  decisionSeed,
  below,
  density,
  forwardMergedRange,
  forwardMixtures,
  knotsMean,
  ablationCellMixture,
  mixtureBelow,
  mixtureQuantile,
  mixtureRange,
  personKnots,
  selectBestExpected,
  selectUniformRandom,
  toImprovement,
  type Knots,
  type Mixture,
} from "../src/rules/index.ts";
import { close, closeArray, loadFixture, range, scenarioPicker } from "./rules-fixtures-helpers.ts";

// forward.json：方法.md 6.1–6.5 与约定 1–16，对照独立 Python 参考实现算出的预期值。
const fixture = loadFixture("forward.json");
const pick = scenarioPicker(fixture);
const TOL: number = fixture.tolerance.abs;
const SEED = "fixture-seed";

const xs = (k: Knots): number[] => k.map(([x]) => x);
const ys = (k: Knots): number[] => k.map(([, y]) => y);
const mean = (m: Mixture): number => m.reduce((s, c) => s + c.weight * knotsMean(c.knots), 0);
/** 正向一个候选的各人范围：[{researcher, range}]。 */
const rangesOf = (rs: any[]) => rs.map((r) => ({ researcher: r.researcher as string, range: range(r.range) }));
const candidatesOf = (input: Record<string, any>) =>
  Object.entries(input.candidates as Record<string, any[]>).map(([id, rs]) => ({ id, ranges: rangesOf(rs) }));

for (const name of ["cdf_method_example_2.44_2.52_2.62", "cdf_asymmetric_narrow"]) {
  test(`forward.json ${name}`, () => {
    const s = pick.get(name);
    const k = personKnots(range(s.input.range));
    closeArray(xs(k), s.expected.points.map((p: number[]) => p[0]), TOL, "折点横坐标");
    closeArray(ys(k), s.expected.points.map((p: number[]) => p[1]), TOL, "折点纵坐标");
    closeArray(s.input.x.map((x: number) => below(k, x)), s.expected.cdf, TOL, "F");
    closeArray(s.input.densityAt.map((x: number) => density(k, x)), s.expected.density, TOL, "密度");
    close(knotsMean(k), s.expected.mean, TOL, "期望");
  });
}

test("forward.json minimize_to_improvement", () => {
  const s = pick.get("minimize_to_improvement");
  const imp = toImprovement([{ weight: 1, knots: personKnots(range(s.input.range)) }], s.input.b, "minimize");
  const k = imp[0]!.knots;
  closeArray(xs(k), s.expected.Ypoints.map((p: number[]) => p[0]), TOL, "Y 折点横坐标");
  closeArray(ys(k), s.expected.Ypoints.map((p: number[]) => p[1]), TOL, "Y 折点纵坐标");
  closeArray([k[1]![0], k[2]![0], k[3]![0]], [s.expected.Yrange.q1, s.expected.Yrange.q5, s.expected.Yrange.q9], TOL, "Y 三个数");
  closeArray(s.input.y.map((y: number) => mixtureBelow(imp, y)), s.expected.Ycdf, TOL, "Y 的 F");
});

for (const name of ["forward_merge_r1_real", "forward_merge_minimize"]) {
  test(`forward.json ${name}`, () => {
    const s = pick.get(name);
    const { b } = s.input;
    const direction = s.params.direction;
    const cands = candidatesOf(s.input);
    const raws = forwardMixtures(cands, s.params.lambda);
    const imps = raws.map((m) => toImprovement(m, b, direction));
    closeArray(xs(raws[0]!.at(-1)!.knots), s.expected.reserveRaw, TOL, "原值上的开放预留");
    closeArray(xs(imps[0]!.at(-1)!.knots), s.expected.reserveY, TOL, "改进量上的开放预留");
    cands.forEach((c, i) => {
      const e = s.expected.candidates[c.id];
      const imp = imps[i]!;
      closeArray(s.input.Yprobe.map((y: number) => mixtureBelow(imp, y)), e.Ycdf, TOL, `${c.id} Y 的 F`);
      const yq = [0.1, 0.5, 0.9].map((p) => mixtureQuantile(imp, p));
      closeArray(yq, [e.Yquantiles.q1, e.Yquantiles.q5, e.Yquantiles.q9], TOL, `${c.id} Y 分位数`);
      const r = forwardMergedRange(imp, b, direction);
      closeArray([r.q1, r.q5, r.q9], [e.rawQuantiles.q1, e.rawQuantiles.q5, e.rawQuantiles.q9], TOL, `${c.id} 原值分位数`);
      close(mean(imp), e.Ymean, TOL, `${c.id} Y 期望`);
    });
  });
}

test("forward.json ablation_cell_merge_r1_p1", () => {
  const s = pick.get("ablation_cell_merge_r1_p1");
  const m = ablationCellMixture(
    s.input.explanations.map((e: any) => ({ id: e.id, author: e.author, range: range(e.range) })),
    s.params.lambda,
  );
  closeArray(xs(m.at(-1)!.knots), s.expected.reserve, TOL, "这一格的开放预留");
  closeArray(m.map((c) => c.weight), s.expected.weights, TOL, "成分权重");
  closeArray(s.input.x.map((x: number) => mixtureBelow(m, x)), s.expected.cdf, TOL, "F");
  const q = mixtureRange(m);
  closeArray([q.q1, q.q5, q.q9], [s.expected.quantiles.q1, s.expected.quantiles.q5, s.expected.quantiles.q9], TOL, "分位数");
});

test("forward.json quantile_flat_left_end", () => {
  const s = pick.get("quantile_flat_left_end");
  const m: Mixture = s.input.mixture.map((c: any) => ({ weight: c.weight, knots: personKnots(range(c.range)) }));
  closeArray(s.input.p.map((p: number) => mixtureQuantile(m, p)), s.expected.quantiles, TOL, "分位数");
});

/** 入池走选批规则的正式路径：候选的入池把握和是否入池都在 BatchDecision.candidates 里。 */
function poolOf(s: { params: Record<string, any>; input: Record<string, any> }, candidates: Record<string, any[]>) {
  return selectBestExpected({
    candidates: Object.entries(candidates).map(([id, rs]) => ({ id, ranges: rangesOf(rs) })),
    baseline: s.input.b,
    direction: s.params.direction,
    epsilon: s.params.eps,
    lambda: s.params.lambda,
    batchCap: 1,
    balance: 1,
    seed: SEED,
  }).candidates;
}

for (const name of ["pool_r1_real", "pool_minimize"]) {
  test(`forward.json ${name}`, () => {
    const s = pick.get(name);
    const got = poolOf(s, s.input.candidates);
    assert.deepEqual(got.map((c) => c.id), Object.keys(s.expected).sort());
    for (const c of got) {
      close(c.lossProbability, s.expected[c.id].probBelowMinusEps, TOL, `${c.id} P(Y < −ε)`);
      assert.equal(c.inPool, s.expected[c.id].inPool, `${c.id} 是否入池`);
    }
  });
}

test("forward.json pool_exactly_half", () => {
  const s = pick.get("pool_exactly_half");
  s.input.cases.forEach((cs: any, i: number) => {
    const [c] = poolOf(s, cs.candidates);
    close(c!.lossProbability, s.expected.cases[i].probBelowMinusEps, TOL, `第 ${i} 例 P(Y < −ε)`);
    assert.equal(c!.inPool, s.expected.cases[i].inPool, `第 ${i} 例是否入池`);
  });
});

test("forward.json V_method_example_X_Y", () => {
  const s = pick.get("V_method_example_X_Y");
  const d = (name: string): Mixture =>
    s.input.Ymixtures[name].map((c: any) => ({ weight: c.weight, knots: personKnots(range(c.range)) }));
  close(batchValue([d("X")]), s.expected.V.X, TOL, "V(X)");
  close(batchValue([d("Y")]), s.expected.V.Y, TOL, "V(Y)");
  close(batchValue([d("X"), d("Y")]), s.expected.V["X,Y"], TOL, "V(X,Y)");
});

for (const name of ["batch_r1_real_k2", "batch_minimize_out_of_pool", "batch_tie_by_id"]) {
  test(`forward.json ${name}`, () => {
    const s = pick.get(name);
    const d = selectBestExpected({
      candidates: candidatesOf(s.input),
      baseline: s.input.b,
      direction: s.params.direction,
      epsilon: s.params.eps,
      lambda: s.params.lambda,
      batchCap: s.params.batchCap,
      balance: s.input.balance,
      seed: SEED,
    });
    assert.deepEqual(d.pool, s.expected.pool, "候选池");
    assert.equal(d.k, s.expected.k, "批量");
    assert.deepEqual(d.combos.map((c) => c.ids), s.expected.combos.map((c: any) => c.combo), "列出的组合");
    closeArray(d.combos.map((c) => c.value), s.expected.combos.map((c: any) => c.V), TOL, "各组合的 V");
    assert.deepEqual(d.chosen, s.expected.selected, "选中的组合");
    close(d.chosenValue, s.expected.selectedV, TOL, "V(S*)");
    close(d.poolValue, s.expected.poolAllV, TOL, "V(候选池)");
    close(d.scaleReading, s.expected.rulerReading, TOL, "分数尺读数");
  });
}

test("forward.json batch_uniform_random_minimize", () => {
  const s = pick.get("batch_uniform_random_minimize");
  const { studyId, round, version } = s.input.seedFrom;
  assert.equal(decisionSeed({ studyId, round, version }), s.input.seed, "种子");
  const d = selectUniformRandom({
    candidates: candidatesOf(s.input),
    baseline: s.input.b,
    direction: s.params.direction,
    epsilon: s.params.eps,
    lambda: s.params.lambda,
    batchCap: s.params.batchCap,
    balance: s.input.balance,
    seed: s.input.seed,
  });
  assert.deepEqual(d.pool, s.expected.pool, "候选池");
  assert.equal(d.k, s.expected.k, "批量");
  assert.deepEqual(d.combos, [], "均匀随机不列组合");
  assert.deepEqual(d.chosen, s.expected.selected, "抽中的组合");
  assert.ok(d.chosen.some((id) => !d.pool.includes(id)), "池外候选可以被抽中");
  close(d.chosenValue, s.expected.selectedV, TOL, "V(S*)");
  close(d.poolValue, s.expected.poolAllV, TOL, "V(候选池)");
  close(d.scaleReading, s.expected.rulerReading, TOL, "分数尺读数（按候选池）");
});

test("forward.json adoption_line_cases", () => {
  const s = pick.get("adoption_line_cases");
  s.input.cases.forEach((c: any, i: number) => {
    const { line } = adoptionLine({ baseline: c.B, knownBest: c.knownBest, epsilon: s.params.eps, direction: c.direction });
    close(line, s.expected.lines[i], TOL, `第 ${i} 例采纳线（${c.note}）`);
  });
});

test("forward.json adoption_decision_cases", () => {
  const s = pick.get("adoption_decision_cases");
  s.input.cases.forEach((c: any, i: number) => {
    const e = s.expected.cases[i];
    const r = adopt({
      baseline: c.B,
      knownBest: c.knownBest,
      epsilon: s.params.eps,
      direction: c.direction,
      // 只收有有效成绩的候选（由调用方筛）。
      candidates: Object.entries(c.results as Record<string, number | null>).flatMap(([id, score]) =>
        score === null ? [] : [{ id, score }],
      ),
    });
    close(r.line, e.line, TOL, `第 ${i} 例采纳线`);
    assert.equal(r.best?.id ?? null, e.best, `第 ${i} 例最好的候选`);
    if (e.best !== null) close(r.best!.improvement, e.gain, TOL, `第 ${i} 例改进`);
    assert.equal(r.adopted !== null, e.adopt, `第 ${i} 例是否采纳`);
    if (e.adopt) assert.equal(r.adopted, e.best);
  });
});

test("forward.json 每个场景都已覆盖", () => {
  assert.deepEqual(pick.unused(), []);
});