"""Validate Vortex Exact Attention's retained reference and design models."""

from __future__ import annotations

import hashlib
import json
import math
from pathlib import Path


PROJECT = Path(__file__).resolve().parents[1]
ROOT = PROJECT.parents[1]


def load_models() -> tuple[dict, dict]:
    inventory = json.loads((PROJECT / "benchmarks/reference_inventory.json").read_text())
    model = json.loads((PROJECT / "benchmarks/architecture_model.json").read_text())
    return inventory, model


def validate_documents(inventory: dict, model: dict, root: Path = ROOT) -> None:
    if len(inventory["artifacts"]) != 10:
        raise RuntimeError("the imported reference inventory must contain ten artifacts")
    for artifact in inventory["artifacts"]:
        path = root / artifact["path"]
        if path.stat().st_size != artifact["size_bytes"]:
            raise RuntimeError(f"size mismatch: {path}")
        if hashlib.sha256(path.read_bytes()).hexdigest() != artifact["sha256"]:
            raise RuntimeError(f"SHA-256 mismatch: {path}")

    if model["status"] != "projected_not_measured":
        raise RuntimeError("design estimates must not be labeled as measurements")
    if len(model["candidates"]) < 3:
        raise RuntimeError("at least three state-ownership designs are required")
    preparation_components = model["model_basis"]["exact_preparation_components_ms"]
    preparation_range = model["model_basis"]["exact_preparation_ms"]
    if preparation_range[0] != math.ceil(sum(preparation_components.values())):
        raise RuntimeError("exact-preparation lower bound is inconsistent")
    candidate_ids = [candidate["id"] for candidate in model["candidates"]]
    if set(candidate_ids) != {"VEA-A", "VEA-B", "VEA-C"} or len(candidate_ids) != 3:
        raise RuntimeError("the three reviewed architecture IDs must be unique and retained")
    for candidate in model["candidates"]:
        ownership = candidate["ownership"]
        if ownership["qk_fragments"] == ownership["output_accumulation"]:
            raise RuntimeError(f"{candidate['id']} does not separate QK and RO ownership")
        if ownership["m_d"] == ownership["output_accumulation"]:
            raise RuntimeError(f"{candidate['id']} does not separate m/d and RO ownership")
        if candidate["handoff"]["same_warp"]:
            raise RuntimeError(f"{candidate['id']} retains QK and PV state in one warp")
        if not candidate["handoff"]["bounded_tile_local"]:
            raise RuntimeError(f"{candidate['id']} uses a material handoff")
        if candidate["handoff"]["consumer_order"] != "strict_epoch_order":
            raise RuntimeError(f"{candidate['id']} does not preserve tile consumption order")
        if candidate["mainloop_ms_projected"][0] >= 220:
            raise RuntimeError(f"{candidate['id']} cannot plausibly beat the model gate")
        recovery = candidate["recovered_no_eligible_fraction"]
        synchronization = candidate["synchronization_ms_projected"]
        baseline = model["baseline"]["mainloop_p50_ms"]
        no_eligible = model["baseline"]["no_eligible_cycles_percent"] / 100.0
        expected_mainloop = [
            round(baseline * (1.0 - no_eligible * recovery[1]) + synchronization[0]),
            round(baseline * (1.0 - no_eligible * recovery[0]) + synchronization[1]),
        ]
        if candidate["mainloop_ms_projected"] != expected_mainloop:
            raise RuntimeError(f"{candidate['id']} screening arithmetic is inconsistent")
        expected_complete = [value + margin for value, margin in zip(
            candidate["mainloop_ms_projected"], preparation_range, strict=True,
        )]
        if candidate["complete_attention_ms_projected"] != expected_complete:
            raise RuntimeError(f"{candidate['id']} complete-attention arithmetic is inconsistent")
        if not candidate["capability_dependencies"]:
            raise RuntimeError(f"{candidate['id']} omits unresolved capability dependencies")

    by_id = {candidate["id"]: candidate for candidate in model["candidates"]}
    a_shared = by_id["VEA-A"]["shared_kib"]
    if a_shared["total"] != [
        a_shared["q"] + a_shared["k_double"] + a_shared["v_double"]
        + a_shared["score_double"] + value for value in a_shared["metadata"]
    ]:
        raise RuntimeError("VEA-A shared-memory arithmetic is inconsistent")
    b_shared = by_id["VEA-B"]["shared_kib"]
    if b_shared["total"] != [value - b_shared["alias_saving"] for value in b_shared["unaliased"]]:
        raise RuntimeError("VEA-B shared-memory alias arithmetic is inconsistent")
    c_shared = by_id["VEA-C"]["shared_kib"]
    if c_shared["cluster_total"] != c_shared["producer"] + c_shared["consumer"]:
        raise RuntimeError("VEA-C shared-memory arithmetic is inconsistent")
    decision = model["decision"]
    if decision["kernel_implementation_started"] or decision["production_dispatch_changed"]:
        raise RuntimeError("Phase 1 must remain design-only and isolated")


def validate() -> None:
    validate_documents(*load_models())


if __name__ == "__main__":
    validate()
    print("Vortex exact reference and architecture models are valid")