"""Maintainer CPU replay: three examples, no fitting or benchmark evaluation."""
from __future__ import annotations

import argparse
import ast
from datetime import datetime, timezone
import gc
import importlib.util
import json
from pathlib import Path

import numpy as np
import torch

import inference as api


def import_file(name, path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def run(workspace, package, output):
    workspace, package, output = Path(workspace).resolve(), Path(package).resolve(), Path(output).resolve()
    api.require(not output.exists(), "Refusing to replace CPU replay evidence")
    torch.set_num_threads(4)
    torch.set_float32_matmul_precision("highest")
    catalog_path = package / "model_catalog.json"
    catalog_sha = api.file_sha(catalog_path)
    catalog = api.read_catalog(catalog_path, catalog_sha)
    tokenizer = api.load_tokenizer(catalog, package)
    train_root = workspace / "bio2nl/review/new_data_full_training_v1_2026-10-02"
    reference_path = train_root / "code_snapshot/model.py"
    api.require(api.file_sha(reference_path) == catalog["frozen_model_implementation_sha256"], "Reference code changed")
    reference = import_file("original_frozen_model", reference_path)
    source_path = workspace / "bio2nl/review/new_data_training_smoke_v1_2026-10-01/prepared/source/train.jsonl"
    with source_path.open() as stream:
        row = json.loads(next(stream))
    pairs = [{key: row[key] for key in ("sentence1", "sentence2")},
             {"sentence1": "A small bird sits beside the window.", "sentence2": "A bird is sitting next to a window."},
             {"sentence1": "The protein sequence is an experimental input.", "sentence2": "The train leaves at noon."}]
    arrays, details = api.encode_pairs(tokenizer, pairs)
    encoding_path = workspace / "bio2nl/review/new_data_fixed_transfer_v1_2026-10-02/prepare_data.py"
    module = ast.parse(encoding_path.read_text())
    functions = [node for node in module.body if isinstance(node, ast.FunctionDef) and node.name in {"encode_content", "pack"}]
    api.require(len(functions) == 2, "Reference encoding functions missing")
    namespace = {"np": np, "require": api.require}
    exec(compile(ast.Module(body=functions, type_ignores=[]), str(encoding_path), "exec"), namespace)
    for index, pair in enumerate(pairs):
        original_ids, original_mask = namespace["pack"](*[namespace["encode_content"](tokenizer, pair[k]) for k in ("sentence1", "sentence2")])
        api.require(np.array_equal(original_ids, arrays["input_ids"][index]) and
                    np.array_equal(original_mask, arrays["attention_mask"][index]), "Reference encoding differs")
    array_path = source_path.with_suffix(".npz")
    with np.load(array_path, allow_pickle=False) as values:
        api.require(np.array_equal(values["input_ids"][0], arrays["input_ids"][0]) and
                    np.array_equal(values["attention_mask"][0], arrays["attention_mask"][0]), "Original training encoding differs")
    selection = json.loads((train_root / "results/source_selection.json").read_text())
    chosen = next(x for x in selection["selected"] if x["job"] == {"condition": "EP", "pt_seed": 0, "ft_seed": 0, "smoke": False})
    classifier_path = train_root / "results" / chosen["checkpoint"]["file"]
    portable = api.load_model(catalog, "EP-pt0-ft0", package, classifier_path)
    original, payload = reference.load_checkpoint(classifier_path,
        expected_sha256=chosen["checkpoint"]["sha256"], expected_kind="source_classifier", expected_job=chosen["job"],
        expected_state_sha256=chosen["checkpoint"]["state_sha256"], expected_protocol_sha256=selection["protocol_sha256"],
        expected_manifest_sha256=selection["prepared_manifest_sha256"], expected_epoch=chosen["best_epoch"],
        expected_config=api.config_for_core(catalog["config"]))
    del payload
    original_output = reference.predict_source(original, arrays, batch_size=3)
    portable_output = api.core.predict_source(portable, arrays, batch_size=3)
    classifier_diff = float(np.max(np.abs(original_output["log_probabilities"] - portable_output["log_probabilities"])))
    api.require(classifier_diff == 0 and np.array_equal(original_output["predictions"], portable_output["predictions"]),
                "Original/portable classifier replay differs")
    single_output = api.core.predict_source(portable, arrays, batch_size=1)
    padding_diff = float(np.max(np.abs(single_output["log_probabilities"] - portable_output["log_probabilities"])))
    api.require(padding_diff <= 1e-5, "CPU batch/padding replay exceeds fixed tolerance")
    del original, portable
    gc.collect()
    lm_path = train_root / "results/full/pretrain/EP/pt0/model.pt"
    meta = json.loads(lm_path.with_name("metadata.json").read_text())
    portable_lm = api.load_model(catalog, "EP-pt0", package, lm_path)
    original_lm, payload = reference.load_checkpoint(lm_path,
        expected_sha256=meta["checkpoint"]["sha256"], expected_kind="causal_lm", expected_job=meta["job"],
        expected_state_sha256=meta["checkpoint"]["state_sha256"], expected_protocol_sha256=meta["protocol_sha256"],
        expected_manifest_sha256=meta["prepared_manifest_sha256"],
        expected_config=api.config_for_core(catalog["config"]))
    del payload
    with torch.inference_mode():
        ids = torch.as_tensor(arrays["input_ids"][:, :32])
        mask = torch.as_tensor(arrays["attention_mask"][:, :32])
        a = original_lm(input_ids=ids, attention_mask=mask, use_cache=False).logits
        b = portable_lm(input_ids=ids, attention_mask=mask, use_cache=False).logits
        lm_diff = float((a - b).abs().max())
    api.require(lm_diff == 0, "Original/portable pretraining replay differs")
    evidence = {"status": "passed", "completed_at_utc": datetime.now(timezone.utc).isoformat(),
        "catalog_sha256": catalog_sha, "cpu_threads": 4, "device": "cpu", "precision": "float32",
        "examples": {"source_training_protein_pairs": 1, "source_training_row_id": row["row_id"], "handwritten_english_pairs": 2},
        "reference_model_code_sha256": api.file_sha(reference_path), "portable_core_sha256": api.file_sha(package / "portable_core.py"),
        "reference_encoding_code_sha256": api.file_sha(encoding_path), "all_encodings_exact": True,
        "original_training_first_row_input_exact": True, "classifier_id": "EP-pt0-ft0", "pretraining_id": "EP-pt0",
        "classifier_max_abs_log_probability_difference": classifier_diff, "classifier_predictions_exact": True,
        "batch_size_1_vs_3_max_abs_log_probability_difference": padding_diff, "fixed_batch_padding_tolerance": 1e-5,
        "pretraining_max_abs_logit_difference": lm_diff,
        "gpu_used": False, "optimizer_updates": 0, "benchmark_scores_computed": False,
        "limitation": "CPU FP32 replay validates portable loading/encoding. It does not promise bitwise equality with prior CUDA BF16 benchmark predictions.",
        "code_sha256": {p.name: api.file_sha(p) for p in (package / "inference.py", package / "portable_core.py", package / "validate_cpu_replay.py")}}
    output.write_text(json.dumps(evidence, indent=2, sort_keys=True) + "\n")
    print(json.dumps(evidence, indent=2))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--workspace", type=Path, required=True)
    parser.add_argument("--package", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    run(args.workspace, args.package, args.output)