import { createHash } from "node:crypto";
import { readFileSync, statSync, writeFileSync } from "node:fs";
import { resolve } from "node:path";

const root = resolve(import.meta.dirname, "..");
const paths = {
  sampling: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-sampling.json",
  ),
  evaluation: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-sampling-eval.json",
  ),
  rerun: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-sampling-repro-r0r1.json",
  ),
  reproduction: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-sampling-reproduction.json",
  ),
  output: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-sampling-compact.json",
  ),
};

const readJson = (path) => JSON.parse(readFileSync(path, "utf8"));
const sha256 = (path) => createHash("sha256")
  .update(readFileSync(path))
  .digest("hex");
const artifact = (path) => ({
  bytes: statSync(path).size,
  sha256: sha256(path),
});
const mean = (values) => (
  values.length
    ? values.reduce((total, value) => total + value, 0) / values.length
    : null
);

const sampling = readJson(paths.sampling);
const evaluation = readJson(paths.evaluation);
const reproduction = readJson(paths.reproduction);
const samplingArtifact = artifact(paths.sampling);
const evaluationArtifact = artifact(paths.evaluation);
const rerunArtifact = artifact(paths.rerun);

if (evaluation.input.sampling_sha256 !== samplingArtifact.sha256) {
  throw new Error("sampling evaluator input hash contract failed");
}
if (reproduction.formal.sha256 !== samplingArtifact.sha256) {
  throw new Error("reproduction formal hash contract failed");
}
if (reproduction.rerun.sha256 !== rerunArtifact.sha256) {
  throw new Error("reproduction rerun hash contract failed");
}
if (
  reproduction.summary.all_preregistered_fields_exact
  !== reproduction.summary.cells
) {
  throw new Error("sampling reproduction is not exact");
}

const conditions = sampling.seed_contract.condition_row_order;
const edges = Object.keys(
  sampling.summary.by_source_edge[sampling.sources[0].id],
);
const edgeLabels = {
  system_eos: "System on − off · EOS",
  system_bos: "System on − off · BOS",
  system_x: "System on − off · x",
  system_period: "System on − off · 句点",
  bos_at_s0: "BOS − EOS · system off",
  bos_at_s1: "BOS − EOS · system on",
  x_at_s0: "x − EOS · system off",
  x_at_s1: "x − EOS · system on",
  period_at_s0: "句点 − EOS · system off",
  period_at_s1: "句点 − EOS · system on",
};

const evalKey = (sourceId, baseSeed, condition) => (
  `${sourceId}\0${baseSeed}\0${condition}`
);
const evaluationByKey = new Map(
  evaluation.rows.map((row) => [
    evalKey(row.source_id, row.base_seed, row.condition),
    row,
  ]),
);

const sourceRows = sampling.sources.map((source) => {
  const conditionRows = Object.fromEntries(conditions.map((condition) => {
    const rawSummary = sampling.summary.by_source_condition[
      source.id
    ][condition];
    const evalSummary = evaluation.summary.by_source_condition[
      source.id
    ][condition];
    const samples = source.runs.map((run) => {
      const output = run.outputs.find(
        (candidate) => candidate.condition === condition,
      );
      const assessed = evaluationByKey.get(
        evalKey(source.id, run.base_seed, condition),
      );
      if (!output || !assessed) {
        throw new Error(
          `sample/evaluation row missing: ${source.id}/${run.base_seed}/${condition}`,
        );
      }
      return {
        replicate: run.replicate_label,
        baseSeed: run.base_seed,
        runSeed: run.run_seed,
        generatedTokens: output.generated_tokens,
        hitEos: output.hit_eos,
        truncated: output.stopped_at_max_new_tokens,
        trajectoryHash: output.generated_token_ids_sha256,
        textHash: output.text_sha256,
        preview: output.text.replace(/\s+/g, " ").trim().slice(0, 220),
        completionClass: assessed.completion_class,
        math: assessed.task_evaluation && source.domain === "math"
          ? {
              predicted: assessed.task_evaluation.predicted_final,
              gold: assessed.task_evaluation.gold_final,
              exact: assessed.task_evaluation.fixed_budget_numeric_exact,
              method: assessed.task_evaluation.extraction_method,
            }
          : null,
        code: assessed.task_evaluation && source.domain === "code"
          ? {
              ast: assessed.task_evaluation.python_ast_parse,
              status: assessed.task_evaluation.execution.status,
              passed: assessed.task_evaluation.fixed_budget_tests_pass,
              cacheHit: assessed.task_evaluation.execution_cache_hit,
            }
          : null,
      };
    });
    return [condition, {
      samples: rawSummary.samples,
      naturalEos: rawSummary.natural_eos,
      truncated: rawSummary.budget_truncated,
      uniqueTrajectories: rawSummary.unique_generated_token_hashes,
      greedyInSamples: rawSummary.greedy_full_trajectory_in_sample_set,
      pairwiseSimilarity: rawSummary.pairwise_token_similarity,
      generatedTokens: rawSummary.generated_tokens,
      evaluation: evalSummary,
      trajectories: samples,
    }];
  }));
  return {
    id: source.id,
    domain: source.domain,
    label: source.label,
    sourceCharacters: source.source_characters,
    sourceTokens: source.source_tokens,
    conditions: conditionRows,
    edges: Object.fromEntries(edges.map((edge) => [
      edge,
      {
        ...sampling.summary.by_source_edge[source.id][edge],
        evaluation: evaluation.summary.by_source_edge[source.id][edge],
      },
    ])),
  };
});

const conditionSummary = Object.fromEntries(conditions.map((condition) => {
  const sourceConditions = sourceRows.map(
    (source) => source.conditions[condition],
  );
  const assessed = evaluation.summary.by_condition[condition];
  return [condition, {
    outputs: assessed.outputs,
    naturalEos: assessed.natural_eos,
    truncated: assessed.budget_truncated,
    meanGeneratedTokens: assessed.mean_generated_tokens,
    uniqueTrajectoriesAcrossSourceSets: sourceConditions.reduce(
      (total, row) => total + row.uniqueTrajectories,
      0,
    ),
    fullEightWayDiversitySets: sourceConditions.filter(
      (row) => row.uniqueTrajectories === 8,
    ).length,
    greedyIncludedSourceSets: sourceConditions.filter(
      (row) => row.greedyInSamples,
    ).length,
    math: assessed.math,
    code: assessed.code,
  }];
}));

const edgeSummary = Object.fromEntries(edges.map((edge) => {
  const rows = sourceRows.map((source) => source.edges[edge]);
  return [edge, {
    label: edgeLabels[edge],
    sources: rows.length,
    meanAlignedSimilarity: mean(rows.map(
      (row) => row.batch_seed_aligned_similarity.mean,
    )),
    meanSymmetricNearestSimilarity: mean(rows.map(
      (row) => row.symmetric_mean_nearest_neighbor_similarity,
    )),
    exactHashIntersections: rows.reduce(
      (total, row) => total + row.generated_hash_set_intersection,
      0,
    ),
    exactHashUnion: rows.reduce(
      (total, row) => total + row.generated_hash_set_union,
      0,
    ),
    naturalEosDeltaRightMinusLeft: rows.reduce(
      (total, row) => (
        total + row.natural_eos_count_difference_right_minus_left
      ),
      0,
    ),
  }];
}));

const result = {
  schemaVersion: 1,
  generatedAt: new Date().toISOString(),
  contract: {
    protocolId: sampling.protocol_id,
    model: sampling.model.repo,
    revision: sampling.model.revision,
    checkpointIdentity: sampling.model.checkpoint_identity,
    conditions,
    conditionFactors: sampling.generation_contract.conditions,
    edges,
    edgeLabels,
    sources: sampling.sources.length,
    replicates: sampling.seed_contract.executed_base_seeds.length,
    baseSeeds: sampling.seed_contract.executed_base_seeds,
    decode: {
      doSample: sampling.generation_contract.do_sample,
      temperature: sampling.generation_contract.temperature,
      topP: sampling.generation_contract.top_p,
      topK: sampling.generation_contract.top_k,
      maxNewTokens: sampling.generation_contract.max_new_tokens,
    },
    batchSeedAlignedNotCommonRandomNumbers: (
      sampling.seed_contract.batch_seed_aligned_not_common_random_numbers
    ),
    counterfactualBoundary: (
      sampling.generation_contract.counterfactual_boundary
    ),
  },
  headline: {
    outputs: sampling.summary.outputs,
    naturalEos: sampling.summary.natural_eos,
    truncated: sampling.summary.budget_truncated,
    uniqueTrajectoryHashes: (
      sampling.summary.unique_generated_token_hashes
    ),
    firstTwoSeedComparableCells: (
      sampling.summary.first_two_seeds.comparable_cells
    ),
    firstTwoSeedDifferentTrajectories: (
      sampling.summary.first_two_seeds.different_trajectories
    ),
    promptHashExact: (
      sampling.source_contract.prompt_hash_audit.exact
    ),
    promptHashCells: (
      sampling.source_contract.prompt_hash_audit.cells
    ),
    greedyIncludedSourceConditionSets: sourceRows.reduce(
      (total, source) => total + conditions.filter(
        (condition) => source.conditions[condition].greedyInSamples,
      ).length,
      0,
    ),
    sourceConditionSets: sourceRows.length * conditions.length,
    mathExact: evaluation.summary.math.fixed_budget_exact,
    mathOutputs: evaluation.summary.math.outputs,
    mathMajority: evaluation.summary.math.unique_absolute_majority,
    mathGold: evaluation.summary.math.gold,
    codePassed: evaluation.summary.code.tests_pass,
    codeOutputs: evaluation.summary.code.outputs,
    codeUniqueExecutionKeys: (
      evaluation.sandbox.unique_code_cache_entries
    ),
    reproducedCells: (
      reproduction.summary.all_preregistered_fields_exact
    ),
    reproductionCells: reproduction.summary.cells,
  },
  conditions: conditionSummary,
  edges: edgeSummary,
  sources: sourceRows,
  reproduction: reproduction.summary,
  artifacts: {
    sampling: samplingArtifact,
    evaluation: evaluationArtifact,
    rerun: rerunArtifact,
    reproduction: artifact(paths.reproduction),
  },
  execution: {
    generationSeconds: sampling.sources.reduce(
      (total, source) => total + source.runs.reduce(
        (sourceTotal, run) => sourceTotal + run.generation_seconds,
        0,
      ),
      0,
    ),
    peakCudaMemoryAllocatedBytes: (
      sampling.execution.peak_cuda_memory_allocated_bytes
    ),
    deviceMap: sampling.execution.device_map,
    transformers: sampling.execution.transformers,
    torch: sampling.execution.torch,
    sandbox: evaluation.sandbox,
  },
  claimBoundary: [
    ...sampling.claim_boundary,
    ...evaluation.claim_boundary,
    ...reproduction.claim_boundary,
  ],
};

writeFileSync(
  paths.output,
  `${JSON.stringify(result, null, 2)}\n`,
);
console.log(JSON.stringify({
  output: paths.output,
  ...artifact(paths.output),
  headline: result.headline,
}, null, 2));