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 = {
  baseline: resolve(root, "src/data/deepseek-v2-lite-chat-behavior.json"),
  completion: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-completion-512.json",
  ),
  completionEval: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-completion-512-eval.json",
  ),
  completionRepro: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-completion-512-repro-1pd.json",
  ),
  depth: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-full-depth.json",
  ),
  depthRepro: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-full-depth-repro-1pd.json",
  ),
  output: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-completion-depth-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 baseline = readJson(paths.baseline);
const completion = readJson(paths.completion);
const completionEval = readJson(paths.completionEval);
const completionRepro = readJson(paths.completionRepro);
const depth = readJson(paths.depth);
const depthRepro = readJson(paths.depthRepro);

const completionArtifact = artifact(paths.completion);
const baselineArtifact = artifact(paths.baseline);
if (
  completionEval.input.behavior_sha256 !== completionArtifact.sha256
  || completionEval.input.baseline_sha256 !== baselineArtifact.sha256
) {
  throw new Error("completion evaluator input hash contract failed");
}

const outputKey = (sourceId, condition) => `${sourceId}\0${condition}`;
const completionByKey = new Map(
  completion.sources.flatMap((source) => source.outputs.map((output) => [
    outputKey(source.id, output.condition),
    output,
  ])),
);
const completionReproduction = {
  sources: completionRepro.sources.length,
  cells: 0,
  promptHashExact: 0,
  generatedTokenIdsExact: 0,
  generatedTextExact: 0,
  eosStateExact: 0,
  truncationStateExact: 0,
};
for (const source of completionRepro.sources) {
  for (const output of source.outputs) {
    const formal = completionByKey.get(outputKey(source.id, output.condition));
    if (!formal) {
      throw new Error(
        `completion formal output missing: ${source.id}/${output.condition}`,
      );
    }
    completionReproduction.cells += 1;
    completionReproduction.promptHashExact += (
      formal.prompt_token_ids_sha256 === output.prompt_token_ids_sha256
    );
    completionReproduction.generatedTokenIdsExact += (
      JSON.stringify(formal.generated_token_ids)
      === JSON.stringify(output.generated_token_ids)
    );
    completionReproduction.generatedTextExact += formal.text === output.text;
    completionReproduction.eosStateExact += formal.hit_eos === output.hit_eos;
    completionReproduction.truncationStateExact += (
      formal.stopped_at_max_new_tokens
      === output.stopped_at_max_new_tokens
    );
  }
}
for (const [name, value] of Object.entries(completionReproduction)) {
  if (
    !["sources", "cells"].includes(name)
    && value !== completionReproduction.cells
  ) {
    throw new Error(
      `completion reproduction mismatch: ${name}=${value}`,
    );
  }
}

const stableDepthSource = (source) => {
  const copy = structuredClone(source);
  delete copy.forward_seconds;
  delete copy.peak_cuda_memory_allocated_bytes;
  return copy;
};
const depthSourceById = new Map(
  depth.sources.map((source) => [source.id, source]),
);
const depthReproduction = {
  sources: depthRepro.sources.length,
  sourceObjectsExactAfterRuntimeStrip: 0,
  hiddenStagesPerSource: 29,
  routerLayersPerSource: 26,
  hiddenTensorHashesExact: 0,
  hiddenTensorHashesCompared: 0,
  orderedRouteHashesExact: 0,
  orderedRouteHashesCompared: 0,
  routeWeightHashesExact: 0,
  routeWeightHashesCompared: 0,
};
for (const rerunSource of depthRepro.sources) {
  const formalSource = depthSourceById.get(rerunSource.id);
  if (!formalSource) {
    throw new Error(`depth formal source missing: ${rerunSource.id}`);
  }
  depthReproduction.sourceObjectsExactAfterRuntimeStrip += (
    JSON.stringify(stableDepthSource(formalSource))
    === JSON.stringify(stableDepthSource(rerunSource))
  );
  for (const [stage, stageValue] of Object.entries(
    rerunSource.hidden_stages,
  )) {
    for (const [condition, conditionValue] of Object.entries(
      stageValue.conditions,
    )) {
      for (const scope of ["target_content", "full_input"]) {
        depthReproduction.hiddenTensorHashesCompared += 1;
        depthReproduction.hiddenTensorHashesExact += (
          formalSource.hidden_stages[stage].conditions[condition][scope]
            .tensor_sha256
          === conditionValue[scope].tensor_sha256
        );
      }
    }
  }
  for (const [layer, layerValue] of Object.entries(
    rerunSource.router_layers,
  )) {
    for (const [condition, conditionValue] of Object.entries(
      layerValue.conditions,
    )) {
      for (const scope of ["target_content", "full_input"]) {
        const formalScope = (
          formalSource.router_layers[layer].conditions[condition][scope]
        );
        depthReproduction.orderedRouteHashesCompared += 1;
        depthReproduction.orderedRouteHashesExact += (
          formalScope.ordered_route_sha256
          === conditionValue[scope].ordered_route_sha256
        );
        depthReproduction.routeWeightHashesCompared += 1;
        depthReproduction.routeWeightHashesExact += (
          formalScope.route_weight_sha256
          === conditionValue[scope].route_weight_sha256
        );
      }
    }
  }
}
if (
  depthReproduction.sourceObjectsExactAfterRuntimeStrip
    !== depthReproduction.sources
  || depthReproduction.hiddenTensorHashesExact
    !== depthReproduction.hiddenTensorHashesCompared
  || depthReproduction.orderedRouteHashesExact
    !== depthReproduction.orderedRouteHashesCompared
  || depthReproduction.routeWeightHashesExact
    !== depthReproduction.routeWeightHashesCompared
) {
  throw new Error(
    `depth reproduction mismatch: ${JSON.stringify(depthReproduction)}`,
  );
}

const conditions = Object.keys(completion.generation_contract.conditions);
const edgeOrder = Object.keys(depth.trace_contract.edge_pairs);
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 domains = ["all", "english", "chinese", "code", "math"];
const hiddenStages = [
  "embedding",
  ...Array.from({ length: 27 }, (_, index) => (
    `layer_${String(index).padStart(2, "0")}`
  )),
  "final_norm",
];
const routerLayers = Array.from({ length: 26 }, (_, index) => (
  `layer_${String(index + 1).padStart(2, "0")}`
));

const sourcesForDomain = (domain) => (
  domain === "all"
    ? depth.sources
    : depth.sources.filter((source) => source.domain === domain)
);
const mean = (values) => (
  values.reduce((sum, value) => sum + value, 0) / values.length
);
const aggregateHidden = (sources, stage, edge) => {
  const rows = sources.map(
    (source) => (
      source.hidden_stages[stage].target_comparisons[edge]
    ),
  );
  const tokens = rows.reduce((sum, row) => sum + row.tokens, 0);
  return {
    tokens,
    exactRows: rows.reduce(
      (sum, row) => sum + row.exact_hidden_rows,
      0,
    ),
    meanCosine: rows.reduce(
      (sum, row) => sum + row.mean_cosine_similarity * row.tokens,
      0,
    ) / tokens,
    meanRelativeL2: mean(rows.map((row) => row.mean_relative_l2)),
    maxAbsDelta: Math.max(...rows.map((row) => row.max_abs_delta)),
  };
};
const aggregateRouter = (sources, layer, edge) => {
  const rows = sources.map(
    (source) => (
      source.router_layers[layer].target_comparisons[edge]
    ),
  );
  const tokens = rows.reduce((sum, row) => sum + row.tokens, 0);
  const orderedExact = rows.reduce(
    (sum, row) => sum + row.ordered_topk_exact_tokens,
    0,
  );
  const setExact = rows.reduce(
    (sum, row) => sum + row.set_exact_tokens,
    0,
  );
  return {
    tokens,
    orderedExact,
    orderedExactRate: orderedExact / tokens,
    setExact,
    setExactRate: setExact / tokens,
    meanSetJaccard: rows.reduce(
      (sum, row) => sum + row.mean_set_jaccard * row.tokens,
      0,
    ) / tokens,
    meanTokenWeightedTv: mean(
      rows.map((row) => row.mean_token_weighted_tv),
    ),
    meanAggregateLoadTv: mean(
      rows.map((row) => row.aggregate_load_tv),
    ),
  };
};

const hiddenSeries = Object.fromEntries(domains.map((domain) => {
  const sources = sourcesForDomain(domain);
  return [domain, Object.fromEntries(edgeOrder.map((edge) => [
    edge,
    hiddenStages.map((stage) => ({
      stage,
      ...aggregateHidden(sources, stage, edge),
    })),
  ]))];
}));
const routerSeries = Object.fromEntries(domains.map((domain) => {
  const sources = sourcesForDomain(domain);
  return [domain, Object.fromEntries(edgeOrder.map((edge) => [
    edge,
    routerLayers.map((layer) => ({
      layer,
      ...aggregateRouter(sources, layer, edge),
    })),
  ]))];
}));

const evaluationRows = completionEval.rows.map((row) => ({
  sourceId: row.source_id,
  domain: row.domain,
  condition: row.condition,
  generatedTokens: row.generated_tokens,
  hitEos: row.hit_eos,
  completionClass: row.completion_class,
  taskEvaluation: row.task_evaluation,
}));
const incompleteRows = evaluationRows.filter((row) => !row.hitEos);
const codeRows = evaluationRows.filter((row) => row.domain === "code");
const mathRows = evaluationRows.filter((row) => row.domain === "math");
const codeStatuses = Object.groupBy(
  codeRows,
  (row) => row.taskEvaluation.execution?.status ?? "not_run",
);
const mathExtraction = Object.groupBy(
  mathRows,
  (row) => row.taskEvaluation.extraction_method,
);

const baselineNaturalEos = baseline.sources.reduce(
  (sum, source) => (
    sum + source.outputs.filter((output) => output.hit_eos).length
  ),
  0,
);
const boundaryCrossingTokens = depth.sources.reduce(
  (sum, source) => (
    sum + Object.values(source.boundary_crossing_tokens_by_condition)
      .reduce((sourceSum, value) => sourceSum + value, 0)
  ),
  0,
);
const {
  hidden: _hiddenSummary,
  routes: _routeSummary,
  ...depthHeadline
} = depth.summary;

const compact = {
  schemaVersion: 1,
  artifacts: Object.fromEntries(
    Object.entries(paths)
      .filter(([name]) => name !== "output")
      .map(([name, path]) => [name, artifact(path)]),
  ),
  model: {
    repo: completion.model.repo,
    revision: completion.model.revision,
    checkpointIdentity: completion.model.checkpoint_identity,
    architecture: completion.model.architecture,
    dtype: completion.model.dtype,
  },
  contract: {
    conditions,
    conditionFactors: completion.generation_contract.conditions,
    edgeOrder,
    edgePairs: depth.trace_contract.edge_pairs,
    edgeLabels,
    domains,
    targetScope: depth.source_contract.target_scope,
    decode: completion.generation_contract.decode,
    maxNewTokens: completion.generation_contract.max_new_tokens,
    useCacheGeneration: completion.generation_contract.use_cache,
    useCacheTrace: depth.trace_contract.use_cache,
  },
  completion: {
    budgetLadder: [
      {
        maxNewTokens: baseline.generation_contract.max_new_tokens,
        naturalEos: baselineNaturalEos,
        truncated: 128 - baselineNaturalEos,
        outputs: 128,
      },
      {
        maxNewTokens: 512,
        naturalEos: completionEval.summary.natural_eos,
        truncated: completionEval.summary.budget_truncated,
        outputs: completionEval.summary.outputs,
      },
    ],
    summary: completionEval.summary,
    incompleteRows,
    taskRows: evaluationRows.filter(
      (row) => ["code", "math"].includes(row.domain),
    ),
    taskTotals: {
      math: {
        strictCorrect: mathRows.filter(
          (row) => (
            row.taskEvaluation.strict_complete_numeric_exact
          ),
        ).length,
        total: mathRows.length,
        extractionMethods: Object.fromEntries(
          Object.entries(mathExtraction).map(([key, rows]) => [
            key,
            rows.length,
          ]),
        ),
      },
      code: {
        testsPassed: codeRows.filter(
          (row) => row.taskEvaluation.fixed_budget_tests_pass,
        ).length,
        total: codeRows.length,
        statuses: Object.fromEntries(
          Object.entries(codeStatuses).map(([key, rows]) => [
            key,
            rows.length,
          ]),
        ),
      },
    },
    sandbox: completionEval.sandbox,
    reproduction: completionReproduction,
    execution: completion.execution,
    claimBoundary: completionEval.claim_boundary,
  },
  depth: {
    summary: depthHeadline,
    hiddenSeries,
    routerSeries,
    hiddenStages,
    routerLayers,
    boundaryCrossingTokens,
    sources: depth.sources.map((source) => ({
      id: source.id,
      domain: source.domain,
      label: source.label,
      contentTokens: source.content_tokens,
      contentTokenIdsSha256: source.content_token_ids_sha256,
      boundaryCrossingTokensByCondition: (
        source.boundary_crossing_tokens_by_condition
      ),
    })),
    reproduction: depthReproduction,
    execution: depth.execution,
    claimBoundary: depth.claim_boundary,
  },
};

writeFileSync(
  paths.output,
  `${JSON.stringify(compact, null, 2)}\n`,
  "utf8",
);
process.stdout.write(
  `${paths.output}\n`
  + `${statSync(paths.output).size} bytes compact\n`
  + `${completionReproduction.generatedTokenIdsExact}`
  + ` / ${completionReproduction.cells} completion cells exact\n`
  + `${depthReproduction.sourceObjectsExactAfterRuntimeStrip}`
  + ` / ${depthReproduction.sources} depth sources exact`
  + " after runtime strip\n",
);