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

const root = resolve(import.meta.dirname, "..");
const mainPath = resolve(
  root,
  "src/data/deepseek-v2-lite-routing-history-boundary-token-control.json",
);
const reproPath = resolve(
  root,
  "src/data/deepseek-v2-lite-routing-history-boundary-token-control-repro.json",
);
const outputPath = resolve(
  root,
  "src/data/deepseek-v2-lite-routing-history-boundary-token-control-compact.json",
);

const sha256 = (path) => createHash("sha256")
  .update(readFileSync(path))
  .digest("hex");

const mainSha256 = sha256(mainPath);
const reproSha256 = sha256(reproPath);
const mainBytes = statSync(mainPath).size;
const reproBytes = statSync(reproPath).size;
const exact = mainSha256 === reproSha256 && mainBytes === reproBytes;
if (!exact) {
  throw new Error("boundary-token formal run and rerun are not byte-exact");
}

const boundary = JSON.parse(readFileSync(mainPath, "utf8"));
const boundaryEdges = [
  "system_eos",
  "system_x",
  "system_period",
  "system_newline",
  "x_at_s0",
  "x_at_s1",
  "period_at_s0",
  "period_at_s1",
  "newline_at_s0",
  "newline_at_s1",
];

const aggregateAlignment = (layer, domain, edge) => {
  const rows = layer.prompts
    .filter((prompt) => prompt.domain === domain)
    .map((prompt) => prompt.alignments[edge]);
  const aligned = rows.reduce(
    (sum, row) => sum + row.aligned_tokens,
    0,
  );
  const setExact = rows.reduce(
    (sum, row) => sum + row.set_topk_exact,
    0,
  );
  const orderedExact = rows.reduce(
    (sum, row) => sum + row.ordered_topk_exact,
    0,
  );
  const weightedJaccard = rows.reduce(
    (sum, row) => sum + row.mean_jaccard * row.aligned_tokens,
    0,
  );
  return {
    aligned,
    setExactRate: setExact / aligned,
    orderedExactRate: orderedExact / aligned,
    meanJaccard: weightedJaccard / aligned,
  };
};

const compact = {
  schemaVersion: 1,
  source: {
    mainSha256,
    reproSha256,
    mainBytes,
    reproBytes,
    exact,
  },
  domains: boundary.corpus_contract.domains,
  labels: boundary.corpus_contract.domain_labels,
  inference: boundary.inference_contract,
  contract: {
    tokenIds: boundary.history_boundary_token_contract.boundary_token_ids,
    validation: boundary.history_boundary_token_contract.render_validation,
    official: boundary.boundary.official_serialization_by_boundary,
  },
  layers: boundary.layers.slice(1).map((layer) => ({
    layer: layer.layer,
    alignment: Object.fromEntries(
      boundary.corpus_contract.domains.map((domain) => [
        domain,
        Object.fromEntries(
          boundaryEdges.map((edge) => [
            edge,
            aggregateAlignment(layer, domain, edge),
          ]),
        ),
      ]),
    ),
    scopes: Object.fromEntries(
      ["target_content", "full_input"].map((scope) => [
        scope,
        {
          modes: Object.fromEntries(
            ["prompt_balanced", "token_weighted"].map((mode) => {
              const control = (
                layer.statistics[scope].modes[mode].boundary_control
              );
              return [
                mode,
                Object.fromEntries(
                  boundary.corpus_contract.domains.map((domain) => [
                    domain,
                    {
                      distances: control[domain].system_edge_distances,
                      contrasts: (
                        control[domain].system_edge_distance_contrasts
                      ),
                      cvEdges: control[domain].metric_system_edges.cv,
                      cvContrasts: (
                        control[domain].metric_system_edge_contrasts.cv
                      ),
                      direct: control[domain].direct_substitutions,
                    },
                  ]),
                ),
              ];
            }),
          ),
        },
      ]),
    ),
  })),
};

writeFileSync(
  outputPath,
  `${JSON.stringify(compact, null, 2)}\n`,
  "utf8",
);
process.stdout.write(
  `${outputPath}\n${mainSha256}\n${mainBytes} bytes source → `
  + `${statSync(outputPath).size} bytes compact\n`,
);