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

const block = JSON.parse(readFileSync(mainPath, "utf8"));
const head = JSON.parse(readFileSync(roleHeadPath, "utf8"));
const blockEdges = [
  "system_user_colon",
  "system_assistant_colon",
  "system_user_x",
  "system_assistant_x",
  "assistant_colon_at_s0",
  "assistant_colon_at_s1",
  "user_x_at_s0",
  "user_x_at_s1",
  "assistant_x_at_s0",
  "assistant_x_at_s1",
  "head_at_x_s0",
  "head_at_x_s1",
  "delimiter_at_assistant_s0",
  "delimiter_at_assistant_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 sharedMappings = {
  user_colon: "official",
  assistant_colon: "target_assistant",
};
const sourceById = new Map(
  head.corpus_contract.selected.map((source) => [source.id, source]),
);
const tokenContract = {
  compared: 0,
  messageHashExact: 0,
  renderedHashExact: 0,
  tokenIdHashExact: 0,
  targetContractExact: 0,
};
for (const source of block.corpus_contract.selected) {
  const previousSource = sourceById.get(source.id);
  if (!previousSource) throw new Error(`role-head source missing: ${source.id}`);
  for (const system of [0, 1]) {
    for (const [level, previousLevel] of Object.entries(sharedMappings)) {
      const current = source.conditions[`s${system}_${level}`];
      const previous = (
        previousSource.conditions[`s${system}_${previousLevel}`]
      );
      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 = block.layers.slice(1).flatMap((layer) => {
  const previousLayer = head.layers.find(
    (candidate) => candidate.layer === layer.layer,
  );
  const previousById = new Map(
    previousLayer.prompts.map((prompt) => [prompt.id, prompt]),
  );
  return Object.entries(sharedMappings).map(([level, previousLevel]) => {
    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(`role-head prompt missing: ${prompt.id}`);
      }
      for (const system of [0, 1]) {
        const current = prompt.conditions[`s${system}_${level}`];
        const previous = (
          previousPrompt.conditions[`s${system}_${previousLevel}`]
        );
        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,
      previousLevel,
      ...counts,
    };
  });
});

const compact = {
  schemaVersion: 1,
  source: {
    mainSha256,
    reproSha256,
    mainBytes,
    reproBytes,
    exact,
  },
  domains: block.corpus_contract.domains,
  labels: block.corpus_contract.domain_labels,
  inference: block.inference_contract,
  contract: {
    tokenIds: block.role_marker_block_contract.role_token_ids,
    officialIds: (
      block.role_marker_block_contract.official_target_block_ids
    ),
    levels: block.role_marker_block_contract.level_factors,
    validation: block.role_marker_block_contract.render_validation,
    official: block.boundary.official_serialization_by_role_block,
    tokenContractAgainstRoleHeadRun: tokenContract,
  },
  crossBatch,
  layers: block.layers.slice(1).map((layer) => ({
    layer: layer.layer,
    alignment: Object.fromEntries(
      block.corpus_contract.domains.map((domain) => [
        domain,
        Object.fromEntries(
          blockEdges.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_block_control
              );
              return [
                mode,
                Object.fromEntries(
                  block.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,
                      factorial: {
                        factorCoding: (
                          control[domain].role_block_factorial.factor_coding
                        ),
                        systemEdgeDistanceEffects: (
                          control[domain].role_block_factorial
                            .system_edge_distance_effects
                        ),
                        metricSystemEdgeEffects: (
                          control[domain].role_block_factorial
                            .metric_system_edge_effects
                        ),
                        distributionSystemEdgeMagnitudes: (
                          Object.fromEntries(
                            Object.entries(
                              control[domain].role_block_factorial
                                .distribution_system_edge_effects,
                            ).map(([name, value]) => [
                              name,
                              value.half_l1_magnitude,
                            ]),
                          )
                        ),
                        directFactorEdges: (
                          control[domain].role_block_factorial
                            .direct_factor_edges
                        ),
                        directEffectDependencies: (
                          control[domain].role_block_factorial
                            .direct_effect_dependencies
                        ),
                      },
                    },
                  ]),
                ),
              ];
            }),
          ),
        },
      ]),
    ),
  })),
};

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