#!/usr/bin/env python3
"""Independent small-case and analytic checks of the exact lattice search.

This does not rerun the degree-34/80 experiments.  It compares the solver
with a separate Cartesian-product search in small dimensions and with the
proved formulas on the first two parameter intervals.
"""

from __future__ import annotations

import argparse
import itertools
import json
from fractions import Fraction
from pathlib import Path

from exact_toeplitz import normalized_fourier, normalized_objective, search

F = Fraction


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    checks = []

    # This independent implementation sums all matrix entries directly.
    def direct_value(coefficients: tuple[int, ...], fourier: list[F]) -> F:
        return sum(a * b * fourier[abs(i-j)]
                   for i, a in enumerate(coefficients)
                   for j, b in enumerate(coefficients))

    for numerator in range(1, 25):
        lam = F(numerator, 4)
        for degree in (1, 2, 3):
            fourier = normalized_fourier(lam, degree)
            boxed_best = min(
                direct_value(coefficients + (1,), fourier)
                for coefficients in itertools.product(range(-3, 4), repeat=degree)
            )
            result = search(lam, degree)
            if result["status"] != "complete" or F(result["normalized_minimum"]) != boxed_best:
                raise AssertionError(("small box comparison", lam, degree, boxed_best, result))
            checks.append({"kind": "independent_box_comparison", "lambda": str(lam),
                           "degree": degree, "box": [-3, 3], "status": "passed"})

    for lam in (F(1, 4), F(1, 2), F(3, 4)):
        for degree in (1, 4, 8):
            result = search(lam, degree)
            if F(result["normalized_minimum"]) != 1 or result["polynomial"] != [1]:
                raise AssertionError(("base interval", lam, degree, result))
            checks.append({"kind": "exact_base_interval", "lambda": str(lam),
                           "degree": degree, "status": "passed"})

    for degree in (1, 4, 8):
        result = search(F(1), degree)
        expected = {tuple([1] * length) for length in range(1, degree + 2)}
        observed = {tuple(polynomial) for polynomial in result["canonical_ties"]}
        if F(result["normalized_minimum"]) != 1 or observed != expected:
            raise AssertionError(("lambda=1 ties", degree, result))
        checks.append({"kind": "lambda_one_all_minimizers", "lambda": "1",
                       "degree": degree, "status": "passed"})

    for alpha in (F(1, 4), F(1, 2), F(3, 4)):
        for degree in (1, 4, 8, 13):
            lam = 1 + alpha
            r = normalized_fourier(alpha, degree + 1)[degree + 1]
            expected = (alpha + 1) / (2 * alpha + 1) * (1 - r)
            result = search(lam, degree, [1] * (degree + 1))
            if F(result["normalized_minimum"]) != expected or result["polynomial"] != [1] * (degree + 1):
                raise AssertionError(("first shifted interval", lam, degree, expected, result))
            checks.append({"kind": "exact_first_shifted_interval", "lambda": str(lam),
                           "degree": degree, "status": "passed"})

    # Deliberate early termination must never emit a certified minimum.
    limited = search(F(5, 2), 13, node_limit=1)
    if limited["status"] != "node_limit" or limited["normalized_minimum"] is not None:
        raise AssertionError(("node limit status", limited))
    checks.append({"kind": "node_limit_is_incomplete", "status": "passed"})

    report = {"status": "passed", "checks_passed": len(checks), "checks": checks}
    serialized = json.dumps(report, indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(serialized)
    print(json.dumps({"status": "passed", "checks_passed": len(checks)}))


if __name__ == "__main__":
    main()