#!/usr/bin/env python3
"""Exact finite-degree Fuchs optimization over the integer lattice.

The rational LDL^T decomposition uses fractions.Fraction; a common
denominator then makes every pruning decision an integer comparison.
No floating point value is used to include or exclude a lattice point.  A node limit,
if requested, is reported as an incomplete search, never as a minimum.

The objective is divided by 2*pi*c_0(lambda), so its Toeplitz coefficients
are rational for a rational lambda.  A complete search at exact degree n
also covers every degree <= n, by multiplication by a monomial.

Run ``python exact_toeplitz.py --help`` for the command line interface.
This program uses only the Python standard library.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import platform
import sys
import time
from fractions import Fraction
from pathlib import Path
from typing import Iterator

F = Fraction


def normalized_fourier(lam: F, degree: int) -> list[F]:
    """Return c_j(lambda)/c_0(lambda), 0 <= j <= degree, exactly."""
    values = [F(1)]
    for j in range(degree):
        values.append(-values[-1] * (lam - j) / (lam + j + 1))
    return values


def polynomial_product(left: list[int], right: list[int]) -> list[int]:
    result = [0] * (len(left) + len(right) - 1)
    for i, a in enumerate(left):
        for j, b in enumerate(right):
            result[i + j] += a * b
    return result


def dirichlet_product(lengths: list[int]) -> list[int]:
    result = [1]
    for length in lengths:
        if length < 1:
            raise ValueError("Dirichlet lengths must be positive integers")
        result = polynomial_product(result, [1] * length)
    return result


def normalized_objective(coefficients: list[int], fourier: list[F]) -> F:
    """Return J_lambda(P)/c_0(lambda) from an exact Fourier vector."""
    result = F(sum(a * a for a in coefficients))
    for lag in range(1, len(coefficients)):
        autocorrelation = sum(
            coefficients[i] * coefficients[i + lag]
            for i in range(len(coefficients) - lag)
        )
        result += 2 * autocorrelation * fourier[lag]
    return result


def canonical_polynomial(coefficients: list[int]) -> list[int]:
    first = 0
    while first < len(coefficients) - 1 and coefficients[first] == 0:
        first += 1
    return coefficients[first:]


def ldlt(matrix: list[list[F]]) -> tuple[list[list[F]], list[F]]:
    """Exact positive-definite LDL^T decomposition; L has unit diagonal."""
    n = len(matrix)
    lower = [[F(0) for _ in range(n)] for _ in range(n)]
    diagonal = [F(0) for _ in range(n)]
    for j in range(n):
        lower[j][j] = F(1)
        diagonal[j] = matrix[j][j] - sum(
            (lower[j][k] ** 2) * diagonal[k] for k in range(j)
        )
        if diagonal[j] <= 0:
            raise ArithmeticError("The exact Gram matrix is not positive definite")
        for i in range(j + 1, n):
            lower[i][j] = (
                matrix[i][j]
                - sum(lower[i][k] * lower[j][k] * diagonal[k] for k in range(j))
            ) / diagonal[j]
    return lower, diagonal


def solve_ldlt(lower: list[list[F]], diagonal: list[F], rhs: list[F]) -> list[F]:
    n = len(diagonal)
    forward = [F(0)] * n
    for i in range(n):
        forward[i] = rhs[i] - sum(lower[i][j] * forward[j] for j in range(i))
    middle = [forward[i] / diagonal[i] for i in range(n)]
    result = [F(0)] * n
    for i in range(n - 1, -1, -1):
        result[i] = middle[i] - sum(lower[j][i] * result[j] for j in range(i + 1, n))
    return result


def exact_integer_interval(center: F, squared_radius: F) -> tuple[int, int]:
    """Exact integer bounds for (x-center)^2 <= squared_radius.

    If center = p/q, an admissible x satisfies the integral inequality
    (x*q-p)^2 <= floor(squared_radius*q*q).  An integer square root
    therefore supplies exact endpoints, without a rounded square root.
    """
    if squared_radius < 0:
        return 1, 0
    p, q = center.numerator, center.denominator
    square_bound = (squared_radius.numerator * q * q) // squared_radius.denominator
    radius_integer = math.isqrt(square_bound)
    low = -(-(p - radius_integer) // q)
    high = (p + radius_integer) // q
    return low, high


def nearest_integer(value: F) -> int:
    return (2 * value.numerator + value.denominator) // (2 * value.denominator)


def ordered_integers(low: int, high: int, center: F) -> Iterator[int]:
    if low > high:
        return
    initial = min(high, max(low, nearest_integer(center)))
    yield initial
    distance = 1
    while initial - distance >= low or initial + distance <= high:
        left, right = initial - distance, initial + distance
        if left >= low and right <= high:
            if abs(F(left) - center) <= abs(F(right) - center):
                yield left
                yield right
            else:
                yield right
                yield left
        elif left >= low:
            yield left
        elif right <= high:
            yield right
        distance += 1


def ordered_integers_scaled(low: int, high: int, numerator: int, denominator: int) -> Iterator[int]:
    """Nearest-center order using integers only, with the same exact ties."""
    if low > high:
        return
    initial = min(high, max(low, (2 * numerator + denominator) // (2 * denominator)))
    yield initial
    distance = 1
    while initial - distance >= low or initial + distance <= high:
        left, right = initial - distance, initial + distance
        if left >= low and right <= high:
            if abs(left * denominator - numerator) <= abs(right * denominator - numerator):
                yield left
                yield right
            else:
                yield right
                yield left
        elif left >= low:
            yield left
        elif right <= high:
            yield right
        distance += 1


class NodeLimit(Exception):
    pass


def search(
    lam: F,
    degree: int,
    incumbent: list[int] | None = None,
    node_limit: int | None = None,
    progress_every: int = 0,
    checkpoint: Path | None = None,
) -> dict:
    if lam <= 0 or degree < 0:
        raise ValueError("lambda must be positive and degree nonnegative")
    if incumbent is None:
        incumbent = [1]
    if len(incumbent) > degree + 1 or incumbent[-1] != 1:
        raise ValueError("incumbent must be monic and have degree <= the search degree")
    fourier = normalized_fourier(lam, degree)
    best = normalized_objective(incumbent, fourier)
    best_polynomial = [0] * (degree + 1 - len(incumbent)) + incumbent
    initial_bound = best
    start = time.perf_counter()
    if degree == 0:
        return {
            "lambda": str(lam), "degree_bound": 0, "status": "complete",
            "normalized_minimum": str(best), "normalized_best": str(best), "polynomial": [1],
            "nodes": 1, "leaves": 1, "elapsed_seconds": time.perf_counter() - start,
        }
    matrix = [[fourier[abs(i-j)] for j in range(degree)] for i in range(degree)]
    linear = [fourier[degree-i] for i in range(degree)]
    lower, diagonal = ldlt(matrix)
    center = solve_ldlt(lower, diagonal, [-x for x in linear])
    real_minimum = 1 + sum(linear[i] * center[i] for i in range(degree))

    # At coordinate i the conditional center is
    # gamma_i - sum_{j>i} L[j,i]*a_j.  Choose a fixed denominator q_i
    # for all these rational coefficients.  The contribution to the
    # squared norm becomes D_i/q_i^2 * (q_i*a_i-p_i)^2.  A single common
    # denominator makes the entire recursive search integer arithmetic.
    center_denominators: list[int] = []
    center_constants: list[int] = []
    triangular_columns: list[list[int]] = []
    rational_weights: list[F] = []
    for i in range(degree):
        gamma = center[i] + sum(lower[j][i] * center[j] for j in range(i + 1, degree))
        q = math.lcm(gamma.denominator, *(lower[j][i].denominator for j in range(i + 1, degree)))
        center_denominators.append(q)
        center_constants.append(gamma.numerator * (q // gamma.denominator))
        triangular_columns.append([
            lower[j][i].numerator * (q // lower[j][i].denominator)
            for j in range(i + 1, degree)
        ])
        rational_weights.append(diagonal[i] / (q * q))
    common_denominator = math.lcm(real_minimum.denominator, best.denominator,
                                  *(weight.denominator for weight in rational_weights))
    weights = [weight.numerator * (common_denominator // weight.denominator)
               for weight in rational_weights]
    best_integer = best.numerator * (common_denominator // best.denominator)
    real_integer = real_minimum.numerator * (common_denominator // real_minimum.denominator)
    decomposition_seconds = time.perf_counter() - start
    assigned = [0] * degree
    nodes = leaves = improvements = 0
    canonical_ties: set[tuple[int, ...]] = {tuple(canonical_polynomial(best_polynomial))}
    best_history = [{"nodes": 0, "normalized_value": str(best),
                     "polynomial": canonical_polynomial(best_polynomial)}]

    def save_incumbent() -> None:
        if checkpoint is None:
            return
        checkpoint.parent.mkdir(parents=True, exist_ok=True)
        record = {"lambda": str(lam), "degree_bound": degree,
                  "status": "incumbent_only", "normalized_best": str(best),
                  "polynomial": canonical_polynomial(best_polynomial),
                  "nodes": nodes, "elapsed_seconds": time.perf_counter() - start,
                  "best_history": best_history}
        temporary = checkpoint.with_suffix(checkpoint.suffix + ".tmp")
        temporary.write_text(json.dumps(record, indent=2, sort_keys=True) + "\n")
        temporary.replace(checkpoint)

    save_incumbent()

    def visit(index: int, tail_cost: int) -> None:
        nonlocal nodes, leaves, improvements, best, best_integer, best_polynomial, canonical_ties
        if node_limit is not None and nodes >= node_limit:
            raise NodeLimit
        nodes += 1
        if progress_every and nodes % progress_every == 0:
            print(json.dumps({"lambda": str(lam), "degree": degree, "nodes": nodes,
                              "best": str(best),
                              "seconds": round(time.perf_counter() - start, 3)}),
                  file=sys.stderr, flush=True)
        if index < 0:
            leaves += 1
            integer_value = real_integer + tail_cost
            value = F(integer_value, common_denominator)
            polynomial = assigned[:] + [1]
            # An independent exact evaluation checks the triangular form.
            if value != normalized_objective(polynomial, fourier):
                raise ArithmeticError("triangular objective disagrees with the Toeplitz form")
            if integer_value < best_integer:
                best_integer = integer_value
                best = value
                best_polynomial = polynomial
                improvements += 1
                canonical_ties = {tuple(canonical_polynomial(polynomial))}
                best_history.append({"nodes": nodes, "normalized_value": str(best),
                                     "polynomial": canonical_polynomial(polynomial)})
                save_incumbent()
            elif integer_value == best_integer:
                canonical_ties.add(tuple(canonical_polynomial(polynomial)))
            return
        remaining = best_integer - real_integer - tail_cost
        if remaining < 0:
            return
        p = center_constants[index] - sum(
            coefficient * assigned[j]
            for j, coefficient in enumerate(triangular_columns[index], index + 1)
        )
        q = center_denominators[index]
        radius = math.isqrt(remaining // weights[index])
        low, high = -(-(p - radius) // q), (p + radius) // q
        for coordinate in ordered_integers_scaled(low, high, p, q):
            cost = weights[index] * (coordinate * q - p) ** 2
            # The incumbent may have improved since the initial interval.
            if real_integer + tail_cost + cost > best_integer:
                continue
            assigned[index] = coordinate
            visit(index - 1, tail_cost + cost)

    status = "complete"
    try:
        visit(degree-1, 0)
    except NodeLimit:
        status = "node_limit"
    elapsed = time.perf_counter() - start
    # Decimal values are for display only; they played no role in the search.
    c0_display = math.gamma(2 * float(lam) + 1) / math.gamma(float(lam) + 1) ** 2
    return {
        "lambda": str(lam), "degree_bound": degree, "status": status,
        "arithmetic": "rational LDL^T; integer comparisons and exact integer-square-root branch bounds",
        "search_denominator_bits": common_denominator.bit_length(),
        "normalized_objective": "integral / (2*pi*c0(lambda))",
        "initial_bound": str(initial_bound), "normalized_best": str(best),
        "normalized_minimum": str(best) if status == "complete" else None,
        "integral_over_2pi_display": float(best) * c0_display,
        "real_relaxation_minimum": str(real_minimum),
        "polynomial": canonical_polynomial(best_polynomial),
        "polynomial_degree": len(canonical_polynomial(best_polynomial)) - 1,
        "canonical_tie_count": len(canonical_ties),
        "canonical_ties": [list(p) for p in sorted(canonical_ties)],
        "nodes": nodes, "leaves": leaves, "improvements": improvements,
        "node_limit": node_limit, "decomposition_seconds": decomposition_seconds,
        "elapsed_seconds": elapsed, "best_history": best_history,
        "python_version": platform.python_version(),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--lambda", dest="lam", required=True, help="positive rational, e.g. 5/2")
    parser.add_argument("--degree", type=int, required=True)
    candidate = parser.add_mutually_exclusive_group()
    candidate.add_argument("--lengths", help="comma-separated Dirichlet factor lengths")
    candidate.add_argument("--coefficients", help="comma-separated coefficients in ascending order")
    candidate.add_argument("--incumbent-json", type=Path, help="JSON file containing a polynomial coefficient list")
    parser.add_argument("--node-limit", type=int)
    parser.add_argument("--progress-every", type=int, default=0)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--checkpoint", type=Path, help="save the initial and every improved incumbent atomically")
    args = parser.parse_args()
    if args.lengths:
        incumbent = dirichlet_product([int(x) for x in args.lengths.split(",")])
    elif args.coefficients:
        incumbent = [int(x) for x in args.coefficients.split(",")]
    elif args.incumbent_json:
        incumbent = json.loads(args.incumbent_json.read_text())["polynomial"]
    else:
        incumbent = [1]
    result = search(F(args.lam), args.degree, incumbent,
                    args.node_limit, args.progress_every, args.checkpoint)
    result["script_sha256"] = hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
    serialized = json.dumps(result, indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(serialized)
    print(serialized, end="")


if __name__ == "__main__":
    main()