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-role-marker-head-control.json",
);
const reproPath = resolve(
  root,
  "src/data/deepseek-v2-lite-routing-role-marker-head-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-role-marker-head-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("role-marker formal run and rerun are not byte-exact");
}

const role = JSON.parse(readFileSync(mainPath, "utf8"));
const boundary = JSON.parse(readFileSync(boundaryPath, "utf8"));
const roleEdges = [
  "system_official",
  "system_target_assistant",
  "system_target_x",
  "system_suffix_user",
  "target_assistant_at_s0",
  "target_assistant_at_s1",
  "target_x_at_s0",
  "target_x_at_s1",
  "suffix_user_at_s0",
  "suffix_user_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 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 role.corpus_contract.selected) {
  const priorSource = sourceById.get(source.id);
  if (!priorSource) throw new Error(`boundary source missing: ${source.id}`);
  for (const system of [0, 1]) {
    const current = source.conditions[`s${system}_official`];
    const previous = priorSource.conditions[`s${system}_eos`];
    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 = role.layers.slice(1).map((layer) => {
  const previousLayer = boundary.layers.find(
    (candidate) => candidate.layer === layer.layer,
  );
  const previousById = new Map(
    previousLayer.prompts.map((prompt) => [prompt.id, prompt]),
  );
  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 current = prompt.conditions[`s${system}_official`];
      const previous = previousPrompt.conditions[`s${system}_eos`];
      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, ...counts };
});

const compact = {
  schemaVersion: 1,
  source: {
    mainSha256,
    reproSha256,
    mainBytes,
    reproBytes,
    exact,
  },
  domains: role.corpus_contract.domains,
  labels: role.corpus_contract.domain_labels,
  inference: role.inference_contract,
  contract: {
    tokenIds: role.role_marker_head_contract.role_token_ids,
    validation: role.role_marker_head_contract.render_validation,
    official: role.boundary.official_serialization_by_role_head,
    tokenContractAgainstBoundaryRun: tokenContract,
  },
  crossBatch,
  layers: role.layers.slice(1).map((layer) => ({
    layer: layer.layer,
    alignment: Object.fromEntries(
      role.corpus_contract.domains.map((domain) => [
        domain,
        Object.fromEntries(
          roleEdges.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].role_marker_control
              );
              return [
                mode,
                Object.fromEntries(
                  role.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`,
);