#!/usr/bin/env node

import { createHash } from "node:crypto";
import { readFile, writeFile } from "node:fs/promises";
import process from "node:process";

const CONDITIONS = ["s0_eos", "s1_eos", "s0_period", "s1_period"];
const DOMAINS = ["code", "math"];
const CONTRASTS = [
  "period_at_s0",
  "period_at_s1",
  "system_at_eos",
  "system_at_period",
];
const PATHS = {
  sampling: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn.json",
  evaluation: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn-eval.json",
  reproduction: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn-reproduction.json",
  analysis: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn-analysis.json",
  output: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn-compact.json",
};

function sha256(value) {
  return createHash("sha256").update(value).digest("hex");
}

function parseArgs(argv) {
  const result = { ...PATHS };
  for (let index = 2; index < argv.length; index += 2) {
    const key = argv[index]?.replace(/^--/, "");
    const value = argv[index + 1];
    if (!(key in result) || value === undefined) {
      throw new Error(`Unknown or incomplete argument: ${argv[index]}`);
    }
    result[key] = value;
  }
  return result;
}

function success(row) {
  const evaluation = row.task_evaluation;
  return row.domain === "code"
    ? evaluation.fixed_budget_tests_pass
    : evaluation.fixed_budget_numeric_exact;
}

const paths = parseArgs(process.argv);
const inputBytes = {};
const input = {};
for (const name of ["sampling", "evaluation", "reproduction", "analysis"]) {
  inputBytes[name] = await readFile(paths[name]);
  input[name] = JSON.parse(inputBytes[name]);
}
const { sampling, evaluation, reproduction, analysis } = input;
const evaluationIndex = new Map(
  evaluation.rows.map((row) => [
    `${row.source_id}\0${row.tape_label}\0${row.condition}`,
    row,
  ]),
);
const sampleIndex = new Map(
  sampling.sources.flatMap((source) =>
    source.runs.flatMap((run) =>
      run.outputs.map((output) => [
        `${source.id}\0${run.tape_label}\0${output.condition}`,
        output,
      ]),
    ),
  ),
);

const conditionTable = Object.fromEntries(
  DOMAINS.map((domain) => [
    domain,
    CONDITIONS.map((condition) => {
      const summary =
        evaluation.summary.main_t0_by_domain_condition[domain][condition];
      return {
        condition,
        outputs: summary.outputs,
        success: summary.fixed_budget_success,
        successRate: summary.fixed_budget_success / summary.outputs,
        naturalEos: summary.natural_eos,
        naturalEosRate: summary.natural_eos / summary.outputs,
        meanTokens: summary.mean_generated_tokens,
        outcomes: summary.task_outcomes,
      };
    }),
  ]),
);

const contrasts = Object.fromEntries(
  DOMAINS.map((domain) => [
    domain,
    Object.fromEntries(
      CONTRASTS.map((name) => {
        const row =
          analysis.main_t0_selected_task_analysis[domain].contrasts[name];
        const metrics = Object.fromEntries(
          Object.entries(row.metrics).map(([metric, value]) => [
            metric,
            {
              point: value.right_minus_left_point,
              band: value.selected_task_resampling_band,
              directions: value.direction_counts,
              transition: value.transition ?? null,
            },
          ]),
        );
        return [
          name,
          {
            left: row.left,
            right: row.right,
            metrics,
            trajectory: {
              sources: row.trajectory.sources,
              sharedUniformExact:
                row.trajectory.shared_uniform_prefix_exact,
              exactTrajectories: row.trajectory.exact_trajectories,
              commonPrefixTokens: row.trajectory.common_prefix_tokens,
            },
          },
        ];
      }),
    ),
  ]),
);

const diagnostic = Object.fromEntries(
  DOMAINS.map((domain) => [
    domain,
    {
      sourceIds: analysis.multi_tape_diagnostic[domain].source_ids,
      tapes: analysis.multi_tape_diagnostic[domain].tapes,
      contrasts: Object.fromEntries(
        CONTRASTS.map((name) => {
          const row =
            analysis.multi_tape_diagnostic[domain].contrasts[name];
          return [
            name,
            {
              left: row.left,
              right: row.right,
              success: {
                matrix:
                  row.metrics.fixed_budget_success.matrix_task_by_tape,
                tapeMeans:
                  row.metrics.fixed_budget_success.tape_means,
                taskRangeWithinTape:
                  row.metrics.fixed_budget_success.task_range_within_tape,
                tapeRangeWithinTask:
                  row.metrics.fixed_budget_success.tape_range_within_task,
              },
              tokens: {
                matrix: row.metrics.generated_tokens.matrix_task_by_tape,
                tapeMeans: row.metrics.generated_tokens.tape_means,
                taskRangeWithinTape:
                  row.metrics.generated_tokens.task_range_within_tape,
                tapeRangeWithinTask:
                  row.metrics.generated_tokens.tape_range_within_task,
              },
            },
          ];
        }),
      ),
    },
  ]),
);

const tasks = Object.fromEntries(
  DOMAINS.map((domain) => {
    const domainSources = sampling.sources
      .filter((source) => source.domain === domain)
      .sort((left, right) => left.within_domain_index - right.within_domain_index);
    return [
      domain,
      domainSources.map((source) => {
        const conditions = Object.fromEntries(
          CONDITIONS.map((condition) => {
            const evaluationRow = evaluationIndex.get(
              `${source.id}\0T0\0${condition}`,
            );
            const samplingRow = sampleIndex.get(
              `${source.id}\0T0\0${condition}`,
            );
            return [
              condition,
              {
                success: success(evaluationRow),
                outcome: evaluationRow.task_outcome,
                tokens: evaluationRow.generated_tokens,
                naturalEos: evaluationRow.hit_eos,
                truncated: evaluationRow.stopped_at_max_new_tokens,
                trajectoryHash:
                  evaluationRow.generated_token_ids_sha256.slice(0, 12),
                firstTokenIds: samplingRow.generated_token_ids.slice(0, 8),
              },
            ];
          }),
        );
        const taskContrasts = Object.fromEntries(
          CONTRASTS.map((name) => {
            const sourceRow =
              analysis.main_t0_selected_task_analysis[domain].contrasts[name];
            const successRow =
              sourceRow.metrics.fixed_budget_success.by_source.find(
                (row) => row.source_id === source.id,
              );
            const tokenRow = sourceRow.metrics.generated_tokens.by_source.find(
              (row) => row.source_id === source.id,
            );
            const trajectoryRow = sourceRow.trajectory.rows.find(
              (row) => row.source_id === source.id,
            );
            return [
              name,
              {
                successDelta: successRow.right_minus_left,
                tokenDelta: tokenRow.right_minus_left,
                commonPrefixTokens: trajectoryRow.common_prefix_tokens,
                exactTrajectory: trajectoryRow.token_ids_exact,
              },
            ];
          }),
        );
        return {
          id: source.id,
          index: source.within_domain_index,
          conditions,
          contrasts: taskContrasts,
          diagnosticTapes: source.runs.map((run) => run.tape_label),
        };
      }),
    ];
  }),
);

const exampleSource = sampling.sources.find(
  (source) => source.id === "HumanEval/31",
);
const exampleRun = exampleSource.runs.find((run) => run.tape_label === "T0");
const compact = {
  schemaVersion: 1,
  protocolId: sampling.protocol_id,
  capturedAt: sampling.captured_at,
  model: {
    repo: sampling.model.repo,
    revision: sampling.model.revision,
    checkpointIdentity: sampling.model.checkpoint_identity,
    dtype: sampling.model.dtype,
  },
  artifactHashes: Object.fromEntries(
    Object.entries(inputBytes).map(([name, bytes]) => [name, sha256(bytes)]),
  ),
  grid: {
    formalSources: sampling.summary.sources,
    formalRuns: sampling.summary.runs,
    formalOutputs: sampling.summary.outputs,
    mainT0Outputs: evaluation.summary.main_t0.outputs,
    diagnosticAdditionalOutputs:
      evaluation.summary.diagnostic_additional_t1_t3.outputs,
    naturalEos: sampling.summary.natural_eos,
    budgetTruncated: sampling.summary.budget_truncated,
    uniqueTrajectories: sampling.summary.unique_generated_token_hashes,
    promptHashesExact: sampling.source_contract.prompt_hash_audit.exact,
    torchRngUnchangedRuns: sampling.summary.torch_rng_unchanged_runs,
  },
  sampler: {
    name: sampling.generation_contract.decode,
    temperature: sampling.generation_contract.distribution_temperature,
    topP: sampling.generation_contract.distribution_top_p,
    softmaxDtype: sampling.generation_contract.softmax_dtype,
    cdfDtype: sampling.generation_contract.cdf_dtype,
    uniformDtype: sampling.generation_contract.uniform_dtype,
    transformersGenerateCalled:
      sampling.generation_contract.transformers_generate_called,
    torchMultinomialCalled:
      sampling.generation_contract.torch_multinomial_called,
    commonRandomNumbers:
      sampling.seed_contract.explicit_common_random_numbers,
  },
  conditionTable,
  contrasts,
  diagnostic,
  tasks,
  outcomes: {
    code: evaluation.summary.main_t0.by_domain.code.task_outcomes,
    math: evaluation.summary.main_t0.by_domain.math.task_outcomes,
  },
  reproduction: reproduction.summary,
  bootstrap: analysis.bootstrap_contract,
  uniformExample: {
    sourceId: exampleSource.id,
    tape: exampleRun.tape_label,
    uniformUint64FirstEightHex: exampleRun.uniform_uint64_first_eight_hex,
    uniformFloat32FirstEight: exampleRun.uniform_float32_first_eight,
    conditions: Object.fromEntries(
      exampleRun.outputs.map((output) => [
        output.condition,
        {
          generatedTokenIds: output.generated_token_ids.slice(0, 8),
          generatedTokens: output.generated_tokens,
          naturalEos: output.hit_eos,
        },
      ]),
    ),
  },
  deviations: [
    {
      id: "prompt-hash-correction",
      severity: "corrected-before-output",
      summary:
        "The first manifest mislabeled routing-probe hashes as Chat prompt hashes; all 256 Chat hashes were corrected before model output.",
    },
    {
      id: "gold-loaded-in-runner",
      severity: "reported-process-deviation",
      summary:
        "The reused runner loaded gold for a post-decode narrow task_score before generation finished. Gold never entered prompts, logits, sampling, selection, or the authoritative evaluator.",
    },
  ],
  evidenceBoundary: [
    "HumanEval and GSM8K remain separate.",
    "Selected-task bands cover only the frozen 32-task frame under T0.",
    "T1-T3 are sensitivity diagnostics, not extra independent tasks.",
    "The explicit sampler uses the official .3/.95 distribution but is not a torch.multinomial trajectory.",
    "Period prompts are counterfactual and not official-valid chats.",
  ],
};

await writeFile(paths.output, `${JSON.stringify(compact, null, 2)}\n`);
const outputBytes = await readFile(paths.output);
process.stdout.write(
  `${JSON.stringify(
    {
      output: paths.output,
      bytes: outputBytes.length,
      sha256: sha256(outputBytes),
      tasks: Object.values(tasks).reduce((sum, rows) => sum + rows.length, 0),
      formalOutputs: compact.grid.formalOutputs,
      reproduction: compact.reproduction,
    },
    null,
    2,
  )}\n`,
);