#!/usr/bin/env python3
"""Check supplied result integrity without rerunning large enumerations."""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from fractions import Fraction
from pathlib import Path

import exact_products
import exact_toeplitz
from exact_toeplitz import canonical_polynomial, dirichlet_product, normalized_fourier, normalized_objective

F = Fraction


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--results-dir", type=Path, default=Path(__file__).resolve().parent / "results")
    args = parser.parse_args()
    directory = args.results_dir
    core_hash = hashlib.sha256(Path(exact_toeplitz.__file__).read_bytes()).hexdigest()
    products_hash = hashlib.sha256(Path(exact_products.__file__).read_bytes()).hexdigest()
    records = {}
    for path in sorted(directory.glob("toeplitz_lambda_*.json")):
        if not re.fullmatch(r"toeplitz_lambda_\d+_\d+_n\d+\.json", path.name):
            continue
        row = json.loads(path.read_text())
        if row["script_sha256"] != core_hash:
            raise AssertionError(("source hash mismatch", path.name))
        lam, degree = F(row["lambda"]), row["degree_bound"]
        polynomial = row["polynomial"]
        if not polynomial or polynomial[-1] != 1 or len(polynomial)-1 > degree:
            raise AssertionError(("invalid monic polynomial", path.name))
        if polynomial != canonical_polynomial(polynomial):
            raise AssertionError(("noncanonical stored polynomial", path.name))
        fourier = normalized_fourier(lam, degree)
        value = normalized_objective(polynomial, fourier)
        if value != F(row["normalized_best"]):
            raise AssertionError(("stored objective mismatch", path.name))
        if row["status"] == "complete":
            if F(row["normalized_minimum"]) != value:
                raise AssertionError(("minimum field mismatch", path.name))
        elif row["status"] == "node_limit":
            if row["normalized_minimum"] is not None or row["nodes"] != row["node_limit"]:
                raise AssertionError(("incomplete status mismatch", path.name))
        else:
            raise AssertionError(("unexpected run status", path.name, row["status"]))
        ties = row["canonical_ties"]
        if len({tuple(p) for p in ties}) != row["canonical_tie_count"]:
            raise AssertionError(("tie count mismatch", path.name))
        for p in ties:
            if p != canonical_polynomial(p) or normalized_objective(p, fourier) != value:
                raise AssertionError(("tie value mismatch", path.name))
        records[(str(lam), degree)] = row
    summary = json.loads((directory / "toeplitz_summary.json").read_text())["results"]
    if {(row["lambda"], row["degree_bound"]) for row in summary} != set(records):
        raise AssertionError("The combined summary does not cover exactly the individual runs")
    for row in summary:
        if row != records[(row["lambda"], row["degree_bound"])]:
            raise AssertionError(("summary record mismatch", row["lambda"], row["degree_bound"]))
    product_rows = 0
    for degree in (34, 80):
        data = json.loads((directory / f"products_n{degree}.json").read_text())
        if data["script_sha256"] != products_hash:
            raise AssertionError(("product script hash mismatch", degree))
        for row in data["results"]:
            lam = F(row["lambda"])
            polynomial = dirichlet_product(row["lengths"])
            if row["status"] != "complete_product_family" or polynomial != row["polynomial"]:
                raise AssertionError(("product-family record mismatch", degree, lam))
            if normalized_objective(polynomial, normalized_fourier(lam, degree)) != F(row["ratio_to_c0_lambda"]):
                raise AssertionError(("product-family value mismatch", degree, lam))
            product_rows += 1
    grid = json.loads((directory / "degree34_quarter_grid.json").read_text())["results"]
    if len(grid) != 24:
        raise AssertionError("The quarter-grid table must contain 24 rows")
    for row in grid:
        lam = F(row["lambda"])
        if normalized_objective(row["polynomial"], normalized_fourier(lam, 34)) != F(row["normalized_minimum"]):
            raise AssertionError(("quarter-grid value mismatch", lam))
        if row["proof_source"] == "exact_lattice_enumeration":
            result = records[(str(lam), 34)]
            if result["status"] != "complete" or result["normalized_minimum"] != row["normalized_minimum"]:
                raise AssertionError(("quarter-grid enumeration source mismatch", lam))
    witness = json.loads((directory / "nonproduct_17_4.json").read_text())
    difference = F(witness["normalized_product_minimum"]) - F(witness["normalized_value"])
    if difference <= 0 or difference != F(witness["positive_normalized_gap"]):
        raise AssertionError("The nonproduct witness must have a positive exact comparison gap")
    report = {"status": "passed", "lattice_run_records": len(records),
              "product_family_rows": product_rows, "quarter_grid_rows": len(grid),
              "exact_toeplitz_sha256": core_hash, "exact_products_sha256": products_hash}
    (directory / "result_integrity.json").write_text(json.dumps(report, indent=2, sort_keys=True)+"\n")
    print(json.dumps(report))


if __name__ == "__main__":
    main()