#!/usr/bin/env python3
"""Validate, aggregate, and publish K3 AttnRes Round 05 experiment data."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import shutil
import statistics
from pathlib import Path
from typing import Any, Iterable


PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1"
ARCHITECTURES = ("baseline", "block")
DEPTHS = (16, 32)
SEEDS = (2026073001, 2026073002, 2026073003)
STEPS = (0, 100, 500, 2000, 4000, 8000)
FORMAL_STEPS = 8000
FORMAL_BATCH = 32
TARGET_BYTES_PER_RUN = 65_536_000
EXPECTED_TOTAL_TARGET_BYTES = 786_432_000
SMOKE_COMPARE_FIELDS = (
    "protocol_id",
    "run_kind",
    "architecture",
    "depth",
    "seed",
    "steps",
    "batch_size",
    "target_bytes_seen",
    "manifest",
    "model",
    "optimizer",
    "hashes",
    "evaluations",
    "diagnostics",
    "training_history",
    "gradient_gate",
    "environment",
)
REPLAY_COMPARE_FIELDS = tuple(
    field for field in SMOKE_COMPARE_FIELDS if field != "run_kind"
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--formal-dir", type=Path, required=True)
    parser.add_argument("--smoke-a-dir", type=Path, required=True)
    parser.add_argument("--smoke-b-dir", type=Path, required=True)
    parser.add_argument("--replay", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--raw-output-dir", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--compact-output", type=Path, required=True)
    parser.add_argument("--reproduction-output", type=Path, required=True)
    return parser.parse_args()


def read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text())


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def canonical_sha256(value: Any) -> str:
    payload = json.dumps(
        value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
    ).encode()
    return hashlib.sha256(payload).hexdigest()


def atomic_json(path: Path, value: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
    )
    os.replace(temporary, path)


def mean(values: Iterable[float]) -> float:
    return statistics.fmean(values)


def require_finite(value: Any, path: str = "root") -> None:
    if isinstance(value, float):
        if not math.isfinite(value):
            raise ValueError(f"non-finite float at {path}")
    elif isinstance(value, dict):
        for key, child in value.items():
            require_finite(child, f"{path}.{key}")
    elif isinstance(value, list):
        for index, child in enumerate(value):
            require_finite(child, f"{path}[{index}]")


def selected(run: dict[str, Any], fields: tuple[str, ...]) -> dict[str, Any]:
    return {field: run[field] for field in fields}


def final_diagnostic(run: dict[str, Any]) -> dict[str, Any]:
    diagnostic = run["diagnostics"][-1]
    if diagnostic["step"] != FORMAL_STEPS:
        raise ValueError("final diagnostic is not step 8000")
    return diagnostic


def validate_run(
    run: dict[str, Any],
    *,
    path: Path,
    depth: int,
    architecture: str,
    seed: int,
    manifest: dict[str, Any],
    manifest_hash: str,
) -> None:
    if run["protocol_id"] != PROTOCOL_ID or run["run_kind"] != "formal":
        raise ValueError(f"formal protocol/kind mismatch: {path}")
    if (
        run["depth"] != depth
        or run["architecture"] != architecture
        or run["seed"] != seed
    ):
        raise ValueError(f"formal identity mismatch: {path}")
    if (
        run["steps"] != FORMAL_STEPS
        or run["batch_size"] != FORMAL_BATCH
        or run["target_bytes_seen"] != TARGET_BYTES_PER_RUN
    ):
        raise ValueError(f"formal budget mismatch: {path}")
    if run["manifest"]["file_sha256"] != manifest_hash:
        raise ValueError(f"manifest file hash mismatch: {path}")
    for key in (
        "formal_schedule_sha256",
        "validation_tensor_sha256",
        "diagnostic_tensor_sha256",
    ):
        if run["manifest"][key] != manifest["windows"][key]:
            raise ValueError(f"manifest {key} mismatch: {path}")
    if [row["step"] for row in run["evaluations"]] != list(STEPS):
        raise ValueError(f"evaluation steps mismatch: {path}")
    if [row["step"] for row in run["diagnostics"]] != list(STEPS):
        raise ValueError(f"diagnostic steps mismatch: {path}")
    if run["model"]["layers"] != depth:
        raise ValueError(f"model depth mismatch: {path}")

    for diagnostic in run["diagnostics"]:
        capture = diagnostic["capture"]
        if (
            capture["count"] != depth
            or capture["shape"] != [16, 256, 192]
            or set(capture["dtypes"]) != {"torch.float32"}
            or not capture["all_gradients_finite"]
            or not capture["all_gradients_present"]
            or not capture["storage_unique"]
        ):
            raise ValueError(f"activation capture gate mismatch: {path}")
        for field in (
            "activation_grad_rms_by_block",
            "activation_output_rms_by_block",
            "core_parameter_grad_rms_by_block",
        ):
            if len(diagnostic[field]) != depth:
                raise ValueError(f"{field} length mismatch: {path}")
        for field in (
            "layer_input_rms_by_sublayer",
            "branch_output_rms_by_sublayer",
            "stream_state_rms_by_sublayer",
        ):
            if len(diagnostic[field]) != depth * 2:
                raise ValueError(f"{field} length mismatch: {path}")
        if architecture == "baseline":
            if diagnostic["depth_weights"] or diagnostic["output_weights"] is not None:
                raise ValueError(f"unexpected Baseline mixer trace: {path}")
        else:
            if (
                len(diagnostic["depth_weights"]) != depth * 2
                or diagnostic["output_weights"]["sources"] != 9
            ):
                raise ValueError(f"Block mixer trace mismatch: {path}")
    require_finite(run, path.name)


def relative_reduction(baseline: float, block: float) -> float:
    return (baseline - block) / baseline


def depth_verdict(rows: list[dict[str, Any]]) -> dict[str, Any]:
    cv_reductions = [row["relative_cv_reduction"] for row in rows]
    imbalance_reductions = [
        row["relative_imbalance_reduction"] for row in rows
    ]
    mean_cv = mean(cv_reductions)
    mean_imbalance = mean(imbalance_reductions)
    support = (
        all(value > 0 for value in cv_reductions)
        and mean_cv >= 0.20
        and all(value > 0 for value in imbalance_reductions)
        and mean_imbalance >= 0.20
    )
    concern = (
        all(value < 0 for value in cv_reductions)
        and -mean_cv >= 0.20
        and all(value < 0 for value in imbalance_reductions)
        and -mean_imbalance >= 0.20
    )
    if support:
        label = "joint directional support at this depth"
    elif concern:
        label = "joint directional concern at this depth"
    else:
        label = "mixed / inconclusive at this depth"
    return {
        "label": label,
        "threshold_relative": 0.20,
        "cv_reductions": cv_reductions,
        "mean_cv_reduction": mean_cv,
        "imbalance_reductions": imbalance_reductions,
        "mean_imbalance_reduction": mean_imbalance,
        "all_cv_improve": all(value > 0 for value in cv_reductions),
        "all_imbalance_improve": all(
            value > 0 for value in imbalance_reductions
        ),
    }


def summarize_depth(
    depth: int, runs: dict[tuple[int, str, int], dict[str, Any]]
) -> dict[str, Any]:
    rows = []
    for seed in SEEDS:
        baseline = runs[(depth, "baseline", seed)]
        block = runs[(depth, "block", seed)]
        baseline_diagnostic = final_diagnostic(baseline)
        block_diagnostic = final_diagnostic(block)
        baseline_activation = baseline_diagnostic[
            "activation_grad_statistics"
        ]
        block_activation = block_diagnostic["activation_grad_statistics"]
        baseline_bpc = baseline["evaluations"][-1]["bits_per_byte"]
        block_bpc = block["evaluations"][-1]["bits_per_byte"]
        rows.append(
            {
                "seed": seed,
                "baseline_bpc": baseline_bpc,
                "block_bpc": block_bpc,
                "block_minus_baseline_bpc": block_bpc - baseline_bpc,
                "baseline_activation_grad_mean": baseline_activation["mean"],
                "block_activation_grad_mean": block_activation["mean"],
                "block_to_baseline_activation_grad_mean": (
                    block_activation["mean"] / baseline_activation["mean"]
                ),
                "baseline_cv": baseline_activation["population_cv"],
                "block_cv": block_activation["population_cv"],
                "relative_cv_reduction": relative_reduction(
                    baseline_activation["population_cv"],
                    block_activation["population_cv"],
                ),
                "baseline_first_to_last_ratio": baseline_activation[
                    "first_to_last_ratio"
                ],
                "block_first_to_last_ratio": block_activation[
                    "first_to_last_ratio"
                ],
                "baseline_imbalance": baseline_activation[
                    "imbalance_abs_log_ratio"
                ],
                "block_imbalance": block_activation[
                    "imbalance_abs_log_ratio"
                ],
                "relative_imbalance_reduction": relative_reduction(
                    baseline_activation["imbalance_abs_log_ratio"],
                    block_activation["imbalance_abs_log_ratio"],
                ),
                "baseline_parameter_grad_cv": baseline_diagnostic[
                    "core_parameter_grad_statistics"
                ]["population_cv"],
                "block_parameter_grad_cv": block_diagnostic[
                    "core_parameter_grad_statistics"
                ]["population_cv"],
            }
        )
    verdict = depth_verdict(rows)
    return {
        "depth": depth,
        "by_seed": rows,
        "means": {
            "baseline_bpc": mean(row["baseline_bpc"] for row in rows),
            "block_bpc": mean(row["block_bpc"] for row in rows),
            "block_minus_baseline_bpc": mean(
                row["block_minus_baseline_bpc"] for row in rows
            ),
            "baseline_cv": mean(row["baseline_cv"] for row in rows),
            "block_cv": mean(row["block_cv"] for row in rows),
            "relative_cv_reduction": mean(
                row["relative_cv_reduction"] for row in rows
            ),
            "baseline_imbalance": mean(
                row["baseline_imbalance"] for row in rows
            ),
            "block_imbalance": mean(
                row["block_imbalance"] for row in rows
            ),
            "relative_imbalance_reduction": mean(
                row["relative_imbalance_reduction"] for row in rows
            ),
            "block_to_baseline_activation_grad_mean": mean(
                row["block_to_baseline_activation_grad_mean"] for row in rows
            ),
            "baseline_parameter_grad_cv": mean(
                row["baseline_parameter_grad_cv"] for row in rows
            ),
            "block_parameter_grad_cv": mean(
                row["block_parameter_grad_cv"] for row in rows
            ),
        },
        "verdict": verdict,
    }


def compact_cell(run: dict[str, Any]) -> dict[str, Any]:
    diagnostics = []
    for row in run["diagnostics"]:
        compact = {
            "step": row["step"],
            "loss_nats": row["loss_nats"],
            "bits_per_byte": row["bits_per_byte"],
            "activation_grad_rms_by_block": row[
                "activation_grad_rms_by_block"
            ],
            "activation_grad_statistics": row[
                "activation_grad_statistics"
            ],
            "activation_output_rms_by_block": row[
                "activation_output_rms_by_block"
            ],
            "activation_output_statistics": row[
                "activation_output_statistics"
            ],
            "core_parameter_grad_rms_by_block": row[
                "core_parameter_grad_rms_by_block"
            ],
            "core_parameter_grad_statistics": row[
                "core_parameter_grad_statistics"
            ],
        }
        if row["depth_weights"]:
            compact["depth_weights"] = row["depth_weights"]
            compact["output_weights"] = row["output_weights"]
        diagnostics.append(compact)
    return {
        "architecture": run["architecture"],
        "depth": run["depth"],
        "seed": run["seed"],
        "evaluations": run["evaluations"],
        "diagnostics": diagnostics,
        "timing": run["timing"],
        "parameters": run["model"]["parameters"],
        "hashes": run["hashes"],
    }


def main() -> None:
    args = parse_args()
    manifest = read_json(args.manifest)
    manifest_hash = file_sha256(args.manifest)
    if manifest["protocol_id"] != PROTOCOL_ID:
        raise ValueError("manifest protocol mismatch")
    if (
        manifest["windows"]["formal_schedule_cells"] != 768_000
        or manifest["windows"]["formal_steps"] != FORMAL_STEPS
        or manifest["windows"]["formal_batch"] != FORMAL_BATCH
    ):
        raise ValueError("manifest schedule budget mismatch")

    runs: dict[tuple[int, str, int], dict[str, Any]] = {}
    source_paths: dict[str, Path] = {}
    formal_hashes: dict[str, str] = {}
    for depth in DEPTHS:
        for architecture in ARCHITECTURES:
            for seed in SEEDS:
                name = f"depth-{depth}-{architecture}-seed-{seed}.json"
                path = args.formal_dir / name
                run = read_json(path)
                validate_run(
                    run,
                    path=path,
                    depth=depth,
                    architecture=architecture,
                    seed=seed,
                    manifest=manifest,
                    manifest_hash=manifest_hash,
                )
                runs[(depth, architecture, seed)] = run
                public_name = f"formal-{name}"
                source_paths[public_name] = path
                formal_hashes[public_name] = file_sha256(path)

    if sum(run["target_bytes_seen"] for run in runs.values()) != (
        EXPECTED_TOTAL_TARGET_BYTES
    ):
        raise ValueError("formal total target-byte budget mismatch")

    common_initial_exact: dict[str, Any] = {}
    input_gate_exact: dict[str, Any] = {}
    for depth in DEPTHS:
        for seed in SEEDS:
            baseline = runs[(depth, "baseline", seed)]
            block = runs[(depth, "block", seed)]
            public_fields = (
                "initial_public_parameter_structure",
                "initial_public_parameter_tensors",
                "initial_public_parameter_elements",
                "initial_public_parameters",
            )
            exact = all(
                baseline["hashes"][field] == block["hashes"][field]
                for field in public_fields
            )
            gate_exact = (
                baseline["manifest"]["input_gate_tensor_hashes"]
                == block["manifest"]["input_gate_tensor_hashes"]
            )
            key = f"depth-{depth}-seed-{seed}"
            common_initial_exact[key] = {
                "exact": exact,
                "baseline": {
                    field: baseline["hashes"][field] for field in public_fields
                },
                "block": {
                    field: block["hashes"][field] for field in public_fields
                },
            }
            input_gate_exact[key] = {
                "exact": gate_exact,
                "hashes": baseline["manifest"]["input_gate_tensor_hashes"],
            }
            if not exact or not gate_exact:
                raise ValueError(f"paired equality gate failed: {key}")

    smoke_exact: dict[str, Any] = {}
    smoke_hashes: dict[str, str] = {}
    for depth in DEPTHS:
        for architecture in ARCHITECTURES:
            name = f"depth-{depth}-{architecture}.json"
            left_path = args.smoke_a_dir / name
            right_path = args.smoke_b_dir / name
            left = read_json(left_path)
            right = read_json(right_path)
            left_selected = selected(left, SMOKE_COMPARE_FIELDS)
            right_selected = selected(right, SMOKE_COMPARE_FIELDS)
            exact = left_selected == right_selected
            if (
                not exact
                or left["run_kind"] != "smoke"
                or left["steps"] != 20
                or not left["gradient_gate"]["passed"]
            ):
                raise ValueError(f"smoke gate failed: {name}")
            key = f"depth-{depth}-{architecture}"
            smoke_exact[key] = {
                "exact": exact,
                "compare_sha256": canonical_sha256(left_selected),
                "gradient_gate": left["gradient_gate"],
            }
            for label, path in (("a", left_path), ("b", right_path)):
                public_name = f"smoke-{label}-{name}"
                source_paths[public_name] = path
                smoke_hashes[public_name] = file_sha256(path)

    replay = read_json(args.replay)
    replay_formal = runs[(32, "block", 2026073001)]
    replay_left = selected(replay_formal, REPLAY_COMPARE_FIELDS)
    replay_right = selected(replay, REPLAY_COMPARE_FIELDS)
    replay_exact = replay_left == replay_right
    if (
        replay["run_kind"] != "replay"
        or replay["depth"] != 32
        or replay["architecture"] != "block"
        or replay["seed"] != 2026073001
        or not replay_exact
    ):
        raise ValueError("formal replay gate failed")
    replay_public_name = "replay-depth-32-block-seed-2026073001.json"
    source_paths[replay_public_name] = args.replay

    depth_summaries = {
        str(depth): summarize_depth(depth, runs) for depth in DEPTHS
    }
    depth_labels = [
        depth_summaries[str(depth)]["verdict"]["label"] for depth in DEPTHS
    ]
    if all(
        label == "joint directional support at this depth"
        for label in depth_labels
    ):
        overall_verdict = (
            "scale-consistent directional support in this operationalization"
        )
    elif all(
        label == "joint directional concern at this depth"
        for label in depth_labels
    ):
        overall_verdict = (
            "scale-consistent directional concern in this operationalization"
        )
    else:
        overall_verdict = "depth-dependent or inconclusive"

    full = {
        "schema_version": 1,
        "protocol_id": PROTOCOL_ID,
        "manifest": manifest,
        "study": {
            "architectures": list(ARCHITECTURES),
            "depths": list(DEPTHS),
            "seeds": list(SEEDS),
            "diagnostic_steps": list(STEPS),
            "formal_runs": len(runs),
            "formal_target_bytes": EXPECTED_TOTAL_TARGET_BYTES,
            "replay_target_bytes": TARGET_BYTES_PER_RUN,
            "gradient_object": (
                "RMS of d(mean token CE)/d(post-MLP Transformer-block output) "
                "over batch×time×channel"
            ),
        },
        "depth_summaries": depth_summaries,
        "overall_verdict": overall_verdict,
        "gates": {
            "common_initial_parameters": common_initial_exact,
            "paired_input_tensors": input_gate_exact,
            "smoke_exact": smoke_exact,
            "replay": {
                "exact": replay_exact,
                "compare_fields": list(REPLAY_COMPARE_FIELDS),
                "formal_compare_sha256": canonical_sha256(replay_left),
                "replay_compare_sha256": canonical_sha256(replay_right),
                "formal_final_model_state": replay_formal["hashes"][
                    "final_model_state"
                ],
                "replay_final_model_state": replay["hashes"][
                    "final_model_state"
                ],
                "formal_final_optimizer_state": replay_formal["hashes"][
                    "final_optimizer_state"
                ],
                "replay_final_optimizer_state": replay["hashes"][
                    "final_optimizer_state"
                ],
            },
        },
        "runs": {
            f"depth-{depth}-{architecture}-seed-{seed}": run
            for (depth, architecture, seed), run in sorted(runs.items())
        },
    }
    full["canonical_sha256_without_self"] = canonical_sha256(full)

    compact = {
        "schema_version": 1,
        "protocol_id": PROTOCOL_ID,
        "study": full["study"],
        "manifest_summary": {
            "file_sha256": manifest_hash,
            "dataset_revision": manifest["dataset"]["revision"],
            "train_bytes_sha256": manifest["dataset"]["splits"]["train"][
                "concatenated_sha256"
            ],
            "formal_schedule_sha256": manifest["windows"][
                "formal_schedule_sha256"
            ],
            "validation_tensor_sha256": manifest["windows"][
                "validation_tensor_sha256"
            ],
            "diagnostic_tensor_sha256": manifest["windows"][
                "diagnostic_tensor_sha256"
            ],
        },
        "depth_summaries": depth_summaries,
        "overall_verdict": overall_verdict,
        "replay_exact": replay_exact,
        "cells": [
            compact_cell(runs[(depth, architecture, seed)])
            for depth in DEPTHS
            for architecture in ARCHITECTURES
            for seed in SEEDS
        ],
    }
    compact["canonical_sha256_without_self"] = canonical_sha256(compact)

    args.raw_output_dir.mkdir(parents=True, exist_ok=True)
    for public_name, source_path in sorted(source_paths.items()):
        target = args.raw_output_dir / public_name
        temporary = target.with_suffix(target.suffix + ".tmp")
        shutil.copyfile(source_path, temporary)
        os.replace(temporary, target)

    reproduction = {
        "schema_version": 1,
        "protocol_id": PROTOCOL_ID,
        "manifest": {
            "path": str(args.manifest),
            "sha256": manifest_hash,
        },
        "protocol_sha256": file_sha256(
            Path("research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md")
        ),
        "definition_audit_sha256": file_sha256(
            Path("research/K3_ATTNRES_GRADIENT_DEFINITION_AUDIT.md")
        ),
        "runner_sha256": file_sha256(
            Path("experiments/k3/attnres_gradient/train.py")
        ),
        "analyzer_sha256": file_sha256(Path(__file__)),
        "formal_raw_sha256": formal_hashes,
        "smoke_raw_sha256": smoke_hashes,
        "replay_raw_sha256": {
            replay_public_name: file_sha256(args.replay)
        },
        "formal_target_bytes": EXPECTED_TOTAL_TARGET_BYTES,
        "replay_target_bytes": TARGET_BYTES_PER_RUN,
        "smoke_exact": smoke_exact,
        "replay_exact": {
            "exact": replay_exact,
            "compare_sha256": canonical_sha256(replay_left),
            "final_model_state": replay["hashes"]["final_model_state"],
            "final_optimizer_state": replay["hashes"][
                "final_optimizer_state"
            ],
        },
        "aggregate_sha256": full["canonical_sha256_without_self"],
        "compact_sha256": compact["canonical_sha256_without_self"],
        "overall_verdict": overall_verdict,
    }
    reproduction["canonical_sha256_without_self"] = canonical_sha256(
        reproduction
    )

    atomic_json(args.output, full)
    atomic_json(args.compact_output, compact)
    atomic_json(args.reproduction_output, reproduction)
    print(
        json.dumps(
            {
                "formal_runs": len(runs),
                "formal_target_bytes": EXPECTED_TOTAL_TARGET_BYTES,
                "smoke_exact": all(
                    row["exact"] for row in smoke_exact.values()
                ),
                "replay_exact": replay_exact,
                "depth_verdicts": {
                    depth: depth_summaries[str(depth)]["verdict"]["label"]
                    for depth in DEPTHS
                },
                "overall_verdict": overall_verdict,
                "aggregate_sha256": full[
                    "canonical_sha256_without_self"
                ],
                "compact_sha256": compact[
                    "canonical_sha256_without_self"
                ],
                "reproduction_sha256": reproduction[
                    "canonical_sha256_without_self"
                ],
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()