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 = {
  sampling: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-cross-source-sampling.json",
  ),
  evaluation: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-cross-source-sampling-eval.json",
  ),
  rerun: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-cross-source-sampling-repro-r0.json",
  ),
  reproduction: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-cross-source-sampling-reproduction.json",
  ),
  analysis: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-cross-source-sampling-analysis.json",
  ),
  output: resolve(
    root,
    "src/data/deepseek-v2-lite-chat-cross-source-sampling-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 sampling = readJson(paths.sampling);
const evaluation = readJson(paths.evaluation);
const reproduction = readJson(paths.reproduction);
const analysis = readJson(paths.analysis);
const samplingArtifact = artifact(paths.sampling);
const evaluationArtifact = artifact(paths.evaluation);
const rerunArtifact = artifact(paths.rerun);
const reproductionArtifact = artifact(paths.reproduction);

if (evaluation.input.sampling_sha256 !== samplingArtifact.sha256) {
  throw new Error("evaluation → sampling hash contract failed");
}
if (reproduction.formal.sha256 !== samplingArtifact.sha256) {
  throw new Error("reproduction → formal hash contract failed");
}
if (reproduction.rerun.sha256 !== rerunArtifact.sha256) {
  throw new Error("reproduction → rerun hash contract failed");
}
if (
  reproduction.summary.cells !== 64
  || reproduction.summary.all_preregistered_fields_exact !== 64
  || Object.values(reproduction.summary.by_field).some(
    (value) => value !== 64,
  )
) {
  throw new Error("64-cell reproduction gate failed");
}
if (
  analysis.inputs.sampling.sha256 !== samplingArtifact.sha256
  || analysis.inputs.evaluation.sha256 !== evaluationArtifact.sha256
  || analysis.inputs.reproduction.sha256 !== reproductionArtifact.sha256
) {
  throw new Error("analysis input hash chain failed");
}
if (
  analysis.contract.sources !== 16
  || analysis.contract.sources_per_domain !== 4
  || analysis.contract.seeds_per_source_condition !== 4
  || analysis.contract.outputs !== 256
) {
  throw new Error("source-blocked analysis grid contract failed");
}

const conditions = analysis.contract.conditions;
const conditionLabels = {
  s0_eos: "S0 · EOS",
  s1_eos: "S1 · EOS",
  s0_period: "S0 · 句点",
  s1_period: "S1 · 句点",
};
const domainLabels = {
  english: "English · WikiText-2",
  chinese: "中文 · TNEWS",
  code: "Code · HumanEval",
  math: "Math · GSM8K",
};
const sourceLabel = (sourceId) => {
  if (sourceId.startsWith("HumanEval/")) {
    return sourceId;
  }
  if (sourceId.startsWith("gsm8k/")) {
    return `GSM8K/${sourceId.split("/").at(-1)}`;
  }
  if (sourceId.startsWith("tnews/")) {
    return `TNEWS/${sourceId.split("/").at(-1)}`;
  }
  if (sourceId.startsWith("wikitext2/")) {
    return `WikiText/${sourceId.split("/").at(-1)}`;
  }
  return sourceId;
};

const sourceMeta = Object.fromEntries(
  Object.entries(analysis.source_metadata).map(([sourceId, row]) => [
    sourceId,
    {
      ...row,
      label: sourceLabel(sourceId),
    },
  ]),
);

const evalRows = evaluation.rows.map((row) => {
  let task = null;
  if (row.domain === "math") {
    task = {
      predicted: row.task_evaluation.predicted_final,
      gold: row.task_evaluation.gold_final,
      covered: row.task_evaluation.evaluator_covered,
      fixedPass: row.task_evaluation.fixed_budget_numeric_exact,
      strictPass: row.task_evaluation.strict_complete_numeric_exact,
      method: row.task_evaluation.extraction_method,
    };
  } else if (row.domain === "code") {
    task = {
      ast: row.task_evaluation.python_ast_parse,
      executed: row.task_evaluation.execution.status !== "not_run",
      status: row.task_evaluation.execution.status,
      fixedPass: row.task_evaluation.fixed_budget_tests_pass,
      strictPass: row.task_evaluation.strict_complete_tests_pass,
      candidateHash: row.task_evaluation.candidate_sha256,
      cacheHit: row.task_evaluation.execution_cache_hit,
    };
  }
  return {
    sourceId: row.source_id,
    domain: row.domain,
    condition: row.condition,
    replicate: row.replicate_label,
    baseSeed: row.base_seed,
    generatedTokens: row.generated_tokens,
    hitEos: row.hit_eos,
    truncated: row.stopped_at_max_new_tokens,
    trajectoryHash: row.generated_token_ids_sha256,
    completionClass: row.completion_class,
    task,
  };
});

const taskFailures = evalRows
  .filter((row) => row.task && !row.task.strictPass)
  .map((row) => ({
    sourceId: row.sourceId,
    sourceLabel: sourceLabel(row.sourceId),
    domain: row.domain,
    condition: row.condition,
    replicate: row.replicate,
    generatedTokens: row.generatedTokens,
    hitEos: row.hitEos,
    completionClass: row.completionClass,
    detail: row.domain === "math"
      ? {
          predicted: row.task.predicted,
          gold: row.task.gold,
          failure: "wrong_numeric_answer",
        }
      : {
          status: row.task.status,
          candidateHash: row.task.candidateHash,
          failure: row.task.status,
        },
  }));

const taskTotals = Object.fromEntries(
  Object.entries(analysis.task_matrix).map(([domain, tasks]) => [
    domain,
    tasks.map((task) => ({
      sourceId: task.source_id,
      label: sourceLabel(task.source_id),
      withinDomainIndex: task.within_domain_index,
      conditions: Object.fromEntries(
        Object.entries(task.conditions).map(([condition, cell]) => [
          condition,
          {
            pass: cell.strict_complete_success,
            outputs: 4,
            anyPass: cell.observed_any_strict_complete_pass,
            naturalEos: cell.natural_eos,
            meanGeneratedTokens: cell.generated_tokens.mean,
            minGeneratedTokens: cell.generated_tokens.min,
            maxGeneratedTokens: cell.generated_tokens.max,
          },
        ]),
      ),
      totalPass: Object.values(task.conditions).reduce(
        (total, cell) => total + cell.strict_complete_success,
        0,
      ),
      totalOutputs: 16,
    })),
  ]),
);

const periodDirection = Object.fromEntries(
  Object.entries(analysis.prior_direction_check).map(([domain, row]) => {
    const values = Object.values(
      row.period_shortens_mean_tokens.by_source,
    );
    return [domain, {
      shorter: values.filter(Boolean).length,
      sources: values.length,
      bySource: row.period_shortens_mean_tokens.by_source,
    }];
  }),
);
const systemEosDirection = Object.fromEntries(
  Object.entries(analysis.prior_direction_check).map(([domain, row]) => {
    const values = Object.values(
      row.system_on_raises_natural_eos_rate.by_source,
    );
    return [domain, {
      raises: values.filter(Boolean).length,
      sources: values.length,
      bySource: row.system_on_raises_natural_eos_rate.by_source,
    }];
  }),
);

const result = {
  schemaVersion: 1,
  capturedAt: sampling.captured_at,
  contract: {
    protocolId: sampling.protocol_id,
    model: sampling.model.repo,
    revision: sampling.model.revision,
    conditions,
    conditionLabels,
    domains: Object.keys(analysis.contract.sources_by_domain),
    domainLabels,
    sources: analysis.contract.sources,
    sourcesPerDomain: analysis.contract.sources_per_domain,
    seedsPerCell: analysis.contract.seeds_per_source_condition,
    outputs: analysis.contract.outputs,
    baseSeeds: sampling.seed_contract.executed_base_seeds,
    decode: {
      doSample: sampling.generation_contract.do_sample,
      temperature: sampling.generation_contract.temperature,
      topP: sampling.generation_contract.top_p,
      topK: sampling.generation_contract.top_k,
      maxNewTokens: sampling.generation_contract.max_new_tokens,
    },
    sourceIsPrimaryCoverageUnit: (
      analysis.contract.source_is_primary_coverage_unit
    ),
    seedIsWithinSourceRepeat: (
      analysis.contract.seed_is_within_source_repeat
    ),
    noPopulationConfidenceIntervals: (
      analysis.contract.no_population_confidence_intervals
    ),
    noPValues: analysis.contract.no_p_values,
    batchSeedAlignedNotCommonRandomNumbers: (
      sampling.seed_contract.batch_seed_aligned_not_common_random_numbers
    ),
  },
  headline: {
    ...analysis.overall,
    promptHashExact: sampling.source_contract.prompt_hash_audit.exact,
    promptHashCells: sampling.source_contract.prompt_hash_audit.cells,
    firstTwoSeedComparableCells: (
      sampling.summary.first_two_seeds.comparable_cells
    ),
    firstTwoSeedDifferentTrajectories: (
      sampling.summary.first_two_seeds.different_trajectories
    ),
    reproducedCells: (
      reproduction.summary.all_preregistered_fields_exact
    ),
    reproductionCells: reproduction.summary.cells,
    codeAstParse: evaluation.summary.code.ast_parse,
    codeExecuted: evaluation.summary.code.executed,
    codeStatuses: evaluation.summary.code.execution_statuses,
    codeUniqueExecutionKeys: (
      evaluation.sandbox.unique_code_cache_entries
    ),
  },
  sourceMeta,
  sourcesByDomain: analysis.contract.sources_by_domain,
  sourceCells: analysis.source_condition_cells,
  domainConditions: analysis.domain_condition_summary,
  sourceContrasts: analysis.source_contrasts,
  domainContrasts: analysis.domain_contrasts,
  periodDirection,
  systemEosDirection,
  taskTotals,
  taskFailures,
  evalRows,
  reproduction: reproduction.summary,
  artifacts: {
    sampling: samplingArtifact,
    evaluation: evaluationArtifact,
    rerun: rerunArtifact,
    reproduction: reproductionArtifact,
    analysis: artifact(paths.analysis),
  },
  execution: {
    generationSeconds: sampling.sources.reduce(
      (total, source) => total + source.runs.reduce(
        (subtotal, run) => subtotal + run.generation_seconds,
        0,
      ),
      0,
    ),
    peakCudaMemoryAllocatedBytes: (
      sampling.execution.peak_cuda_memory_allocated_bytes
    ),
    torch: sampling.execution.torch,
    transformers: sampling.execution.transformers,
    deviceMap: sampling.execution.device_map,
    sandbox: evaluation.sandbox,
  },
  claimBoundary: [
    ...sampling.claim_boundary,
    ...evaluation.claim_boundary,
    ...analysis.claim_boundary,
    ...reproduction.claim_boundary,
  ],
};

writeFileSync(paths.output, `${JSON.stringify(result, null, 2)}\n`);
console.log(JSON.stringify({
  output: paths.output,
  ...artifact(paths.output),
  headline: result.headline,
  periodDirection: result.periodDirection,
}, null, 2));