#!/usr/bin/env python3
"""Reproduce the exact lattice searches at degrees 34, 61, and 80.

The listed products supply initial upper bounds only.  Every search ranges
over all monic integer polynomials up to the stated degree, with exact
arithmetic.  A search stopped at the requested node limit is explicitly
marked incomplete in both JSON and CSV.
"""

from __future__ import annotations

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

import exact_toeplitz
from exact_toeplitz import dirichlet_product, search
from verify_nonproduct_witness import witness_polynomial


EXPERIMENTS = {
    34: {"3/2": [35], "5/2": [17, 18],
         "13/4": [8, 9, 17], "7/2": [8, 9, 17], "15/4": [8, 9, 17],
         "17/4": [5, 6, 11, 16], "9/2": [2, 3, 5, 7], "19/4": [2, 3, 5, 7],
         "21/4": [3, 4, 7, 10, 11], "11/2": [3, 4, 7, 10, 11],
         "23/4": [3, 4, 7, 10, 11]},
    61: {"17/4": None},
    80: {"3/2": [81], "5/2": [40, 41], "7/2": [20, 21, 41],
         "17/4": None, "9/2": [2, 3, 5, 7], "11/2": [8, 9, 17, 25, 26]},
}


def summarize_results(directory: Path) -> list[dict]:
    rows = []
    for degree, parameters in EXPERIMENTS.items():
        for lam in parameters:
            path = directory / f"toeplitz_lambda_{lam.replace('/', '_')}_n{degree}.json"
            if path.exists():
                rows.append(json.loads(path.read_text()))
    rows.sort(key=lambda row: (row["degree_bound"], Fraction(row["lambda"])))
    (directory / "toeplitz_summary.json").write_text(json.dumps({"results": rows}, indent=2, sort_keys=True) + "\n")
    columns = ["lambda", "degree_bound", "status", "normalized_minimum", "normalized_best",
               "integral_over_2pi_display", "polynomial_degree", "canonical_tie_count",
               "nodes", "leaves", "elapsed_seconds", "search_denominator_bits", "script_sha256"]
    with (directory / "toeplitz_summary.csv").open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=columns)
        writer.writeheader()
        for row in rows:
            writer.writerow({column: row.get(column, "") for column in columns})
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--degrees", default="34,61,80", help="comma-separated choice of 34, 61, and 80")
    parser.add_argument("--parameters", help="optional comma-separated subset of the selected parameters")
    parser.add_argument("--node-limit", type=int, default=10_000_000)
    parser.add_argument("--large-case-node-limit", type=int, default=100_000_000,
                        help="separate limit for lambda=17/4 at degree 80")
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parent / "results")
    parser.add_argument("--summarize-only", action="store_true")
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    if not args.summarize_only:
        source_hash = hashlib.sha256(Path(exact_toeplitz.__file__).read_bytes()).hexdigest()
        selected = {Fraction(x) for x in args.parameters.split(",")} if args.parameters else None
        for degree in [int(x) for x in args.degrees.split(",")]:
            for lam, lengths in EXPERIMENTS[degree].items():
                if selected is not None and Fraction(lam) not in selected:
                    continue
                polynomial = witness_polynomial() if lengths is None else dirichlet_product(lengths)
                limit = args.large_case_node_limit if degree == 80 and lam == "17/4" else args.node_limit
                checkpoint = args.output_dir / f"toeplitz_lambda_{lam.replace('/', '_')}_n{degree}_incumbent.json"
                result = search(Fraction(lam), degree, polynomial,
                                limit, progress_every=5_000_000, checkpoint=checkpoint)
                result["script_sha256"] = source_hash
                result["initial_dirichlet_lengths"] = lengths
                result["initial_candidate_kind"] = "PTE witness" if lengths is None else "Dirichlet product"
                path = args.output_dir / f"toeplitz_lambda_{lam.replace('/', '_')}_n{degree}.json"
                path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
                print(json.dumps({key: result[key] for key in
                                  ("lambda", "degree_bound", "status", "nodes", "elapsed_seconds")}), flush=True)
    rows = summarize_results(args.output_dir)
    print(json.dumps({"summary_rows": len(rows), "summary": str(args.output_dir / "toeplitz_summary.csv")}))


if __name__ == "__main__":
    main()