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-special-token-family-control.json",
);
const reproPath = resolve(
  root,
  "src/data/deepseek-v2-lite-routing-special-token-family-control-repro.json",
);
const boundaryPath = resolve(
  root,
  "src/data/deepseek-v2-lite-routing-history-boundary-token-control.json",
);
const outputPath = resolve(
  root,
  "src/data/deepseek-v2-lite-routing-special-token-family-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("special-token formal run and rerun are not byte-exact");
}

const family = JSON.parse(readFileSync(mainPath, "utf8"));
const boundary = JSON.parse(readFileSync(boundaryPath, "utf8"));
const familyEdges = [
  "system_eos",
  "system_bos",
  "system_x",
  "system_period",
  "bos_at_s0",
  "bos_at_s1",
  "x_at_s0",
  "x_at_s1",
  "period_at_s0",
  "period_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 sharedLevels = ["eos", "x", "period"];
const sourceById = new Map(
  boundary.corpus_contract.selected.map((source) => [source.id, source]),
);
const tokenContract = {
  compared: 0,
  messageHashExact: 0,
  renderedHashExact: 0,
  tokenIdHashExact: 0,
  targetContractExact: 0,
};
for (const source of family.corpus_contract.selected) {
  const previousSource = sourceById.get(source.id);
  if (!previousSource) throw new Error(`boundary source missing: ${source.id}`);
  for (const system of [0, 1]) {
    for (const level of sharedLevels) {
      const current = source.conditions[`s${system}_${level}`];
      const previous = previousSource.conditions[`s${system}_${level}`];
      tokenContract.compared += 1;
      tokenContract.messageHashExact += (
        current.messages_sha256 === previous.messages_sha256
      );
      tokenContract.renderedHashExact += (
        current.rendered_sha256 === previous.rendered_sha256
      );
      tokenContract.tokenIdHashExact += (
        current.token_ids_sha256 === previous.token_ids_sha256
      );
      tokenContract.targetContractExact += (
        current.tokens === previous.tokens
        && current.content_tokens === previous.content_tokens
        && current.aligned_content_tokens
          === previous.aligned_content_tokens
      );
    }
  }
}

const crossBatch = family.layers.slice(1).flatMap((layer) => {
  const previousLayer = boundary.layers.find(
    (candidate) => candidate.layer === layer.layer,
  );
  const previousById = new Map(
    previousLayer.prompts.map((prompt) => [prompt.id, prompt]),
  );
  return sharedLevels.map((level) => {
    const counts = {
      compared: 0,
      fullRouteHashExact: 0,
      targetRouteHashExact: 0,
      fullLoadExact: 0,
      targetLoadExact: 0,
    };
    for (const prompt of layer.prompts) {
      const previousPrompt = previousById.get(prompt.id);
      if (!previousPrompt) {
        throw new Error(`boundary prompt missing: ${prompt.id}`);
      }
      for (const system of [0, 1]) {
        const condition = `s${system}_${level}`;
        const current = prompt.conditions[condition];
        const previous = previousPrompt.conditions[condition];
        counts.compared += 1;
        counts.fullRouteHashExact += (
          current.topk_sha256 === previous.topk_sha256
        );
        counts.targetRouteHashExact += (
          current.content_topk_sha256 === previous.content_topk_sha256
        );
        counts.fullLoadExact += (
          JSON.stringify(current.full_load)
          === JSON.stringify(previous.full_load)
        );
        counts.targetLoadExact += (
          JSON.stringify(current.content_load)
          === JSON.stringify(previous.content_load)
        );
      }
    }
    return { layer: layer.layer, level, ...counts };
  });
});

const compact = {
  schemaVersion: 1,
  source: {
    mainSha256,
    reproSha256,
    mainBytes,
    reproBytes,
    exact,
  },
  domains: family.corpus_contract.domains,
  labels: family.corpus_contract.domain_labels,
  inference: family.inference_contract,
  contract: {
    tokenIds: family.special_token_family_contract.boundary_token_ids,
    inventory: (
      family.special_token_family_contract.tokenizer_special_inventory
    ),
    validation: family.special_token_family_contract.render_validation,
    official: family.boundary.official_serialization_by_boundary,
    classBoundary: (
      family.special_token_family_contract.class_comparison_boundary
    ),
    tokenContractAgainstBoundaryRun: tokenContract,
  },
  crossBatch,
  layers: family.layers.slice(1).map((layer) => ({
    layer: layer.layer,
    alignment: Object.fromEntries(
      family.corpus_contract.domains.map((domain) => [
        domain,
        Object.fromEntries(
          familyEdges.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(
                  family.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,
                      family: control[domain].descriptive_family_summary,
                    },
                  ]),
                ),
              ];
            }),
          ),
        },
      ]),
    ),
  })),
};

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`,
);