#!/usr/bin/env python3
"""Build the full degree-34 quarter-grid table from exact sources.

Rows at lambda <= 3 use the analytic theorems in the paper; integer rows
use explicit PTE witnesses and the integer lower bound.  All noninteger
rows above 3 require a completed exact lattice search.  This script checks
the candidate values again using rational Toeplitz coefficients.
"""

from __future__ import annotations

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

from exact_toeplitz import dirichlet_product, normalized_fourier, normalized_objective, polynomial_product

F = Fraction
INTEGER_LENGTHS = {1: [], 2: [2], 3: [2, 3], 4: [2, 3, 5],
                   5: [2, 3, 5, 7], 6: [2, 3, 4, 5, 7]}


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
    directory.mkdir(parents=True, exist_ok=True)
    rows = []
    for numerator in range(1, 25):
        lam = F(numerator, 4)
        nodes = 0
        if lam.denominator == 1:
            m = int(lam)
            lengths = INTEGER_LENGTHS[m]
            polynomial = dirichlet_product(lengths)
            q = polynomial[:]
            for _ in range(m):
                q = polynomial_product(q, [1, -1])
            if sum(x*x for x in q) != 2*m:
                raise AssertionError(("integer PTE witness", m))
            expected = F(2*m, math.comb(2*m, m))
            source = "integer_PTE_theorem"
            relation = "equals_global_infimum"
        elif lam < 1:
            lengths, polynomial, expected = [], [1], F(1)
            source = "analytic_base_interval"
            relation = "equals_global_infimum"
        elif lam < 2:
            alpha = lam - 1
            lengths, polynomial = [35], [1]*35
            expected = (alpha + 1)/(2*alpha + 1) * (1 - normalized_fourier(alpha, 35)[35])
            source = "analytic_first_shifted_finite_degree_theorem"
            relation = "strictly_above_global_infimum"
        elif lam < 3:
            alpha = lam - 2
            lengths, polynomial = [17, 18], dirichlet_product([17, 18])
            r = normalized_fourier(alpha, 35)
            product_value = 4 + 2*r[35] + 2*r[1] - 4*r[17] - 4*r[18]
            c0_ratio = (2*(2*alpha+1)/(alpha+1)) * (2*(2*alpha+3)/(alpha+2))
            expected = product_value/c0_ratio
            source = "analytic_second_shifted_finite_degree_theorem"
            relation = "strictly_above_global_infimum"
        else:
            path = directory / f"toeplitz_lambda_{str(lam).replace('/', '_')}_n34.json"
            result = json.loads(path.read_text())
            if result["status"] != "complete":
                raise AssertionError(f"An incomplete search cannot enter the exact table: {path}")
            polynomial = result["polynomial"]
            lengths = result.get("initial_dirichlet_lengths", [])
            expected = F(result["normalized_minimum"])
            nodes = result["nodes"]
            source = "exact_lattice_enumeration"
            relation = "global_infimum_unknown"
        value = normalized_objective(polynomial, normalized_fourier(lam, 34))
        if value != expected:
            raise AssertionError(("objective mismatch", lam, value, expected))
        c0_display = math.gamma(2*float(lam)+1) / math.gamma(float(lam)+1)**2
        rows.append({"lambda": str(lam), "degree_bound": 34, "status": "exact",
                     "proof_source": source, "normalized_minimum": str(value),
                     "integral_over_2pi_display": float(value)*c0_display,
                     "polynomial_degree": len(polynomial)-1, "polynomial": polynomial,
                     "representative_lengths": lengths, "global_relation": relation,
                     "enumeration_nodes": nodes})
    report = {"normalization": "normalized_minimum = F_le_34(lambda)/(2*pi*c0(lambda))",
              "results": rows}
    (directory / "degree34_quarter_grid.json").write_text(json.dumps(report, indent=2, sort_keys=True)+"\n")
    fields = ["lambda", "degree_bound", "status", "proof_source", "normalized_minimum",
              "integral_over_2pi_display", "polynomial_degree", "representative_lengths",
              "global_relation", "enumeration_nodes"]
    with (directory / "degree34_quarter_grid.csv").open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        for row in rows:
            writer.writerow({key: row[key] for key in fields})
    print(json.dumps({"status": "passed", "exact_quarter_grid_rows": len(rows)}))


if __name__ == "__main__":
    main()