"""Exact scope certificate for the quadratic family in Theorem 5.5.

The paper's lower-bound proof uses two equally sized client groups with
opposite quadratic optima.  This producer instantiates that construction and
computes the FedAvg dynamics with Fraction arithmetic.  A sampled subset is
represented by its exact hypergeometric count of +kappa clients, which is
equivalent to summing all subsets because the update depends only on that
count.

For Li(theta) = (theta-z_i)^2/2, z_i in {-kappa,+kappa}, the global objective
has gradient theta.  E local steps with step eta map theta to
q*theta + (1-q)*z_i, q=(1-eta)^E.  Thus the server recursion is
theta' = q*theta + (1-q)*mean(z_i in sampled subset).  The exact stationary
gradient-norm second moment follows from this affine recursion.
"""

from __future__ import annotations

import json
import sys
from fractions import Fraction
from math import comb
from pathlib import Path


if hasattr(sys, "set_int_max_str_digits"):
    sys.set_int_max_str_digits(1_000_000)


ETA = Fraction(1, 128)
ROUNDS = 4096
N_VALUES = (4, 8, 16, 32, 64, 128, 256)
E_VALUES = (1, 2, 4, 8, 16)
KAPPA_VALUES = (Fraction(1, 8), Fraction(1, 4), Fraction(1, 2), Fraction(1), Fraction(2))


def exact_subset_variance(n: int, sampled: int, kappa: Fraction) -> Fraction:
    """E[(sample mean z)^2], summing every subset exactly."""
    half = n // 2
    denominator = comb(n, sampled)
    second_moment = Fraction(0)
    for plus_count in range(max(0, sampled - half), min(half, sampled) + 1):
        ways = comb(half, plus_count) * comb(half, sampled - plus_count)
        signed_sum = 2 * plus_count - sampled
        sample_mean = kappa * Fraction(signed_sum, sampled)
        second_moment += Fraction(ways, denominator) * sample_mean * sample_mean
    return second_moment


def stationary_gap(n: int, sampled: int, local_steps: int, kappa: Fraction) -> Fraction:
    """Exact stationary E[||grad L(theta)||^2] for the source construction."""
    q = (1 - ETA) ** local_steps
    local_gain = 1 - q
    subset_variance = exact_subset_variance(n, sampled, kappa)
    # M = q^2 M + (1-q)^2 Var(sample_mean), so M=(1-q)/(1+q)*Var.
    return Fraction(local_gain, 1 + q) * subset_variance


def finite_round_gap(n: int, sampled: int, local_steps: int, kappa: Fraction) -> Fraction:
    """Exact expected gradient gap after ROUNDS rounds, starting at theta=0."""
    q = (1 - ETA) ** local_steps
    return stationary_gap(n, sampled, local_steps, kappa) * (1 - q ** (2 * ROUNDS))


def run() -> dict:
    rows = []
    ratios = []
    finite_minimum = None
    subset_checks = []
    for n in N_VALUES:
        # The source proof is a partial-participation construction.  S<=N/2
        # keeps the finite-population correction bounded away from zero.
        for sampled in range(1, n // 2 + 1):
            variance = exact_subset_variance(n, sampled, Fraction(1))
            expected_variance = Fraction(n - sampled, sampled * (n - 1))
            assert variance == expected_variance
            subset_checks.append((n, sampled))
            for local_steps in E_VALUES:
                assert ETA <= Fraction(1, 8 * local_steps)  # L=1 stability range
                for kappa in KAPPA_VALUES:
                    gap = stationary_gap(n, sampled, local_steps, kappa)
                    finite_gap = finite_round_gap(n, sampled, local_steps, kappa)
                    target = Fraction(local_steps) * kappa * kappa / sampled
                    ratio = gap / target
                    finite_ratio = finite_gap / target
                    ratios.append(ratio)
                    finite_minimum = finite_ratio if finite_minimum is None else min(finite_minimum, finite_ratio)
                    rows.append(
                        {
                            "N": n,
                            "S": sampled,
                            "E": local_steps,
                            "kappa": str(kappa),
                            "subset_variance": str(exact_subset_variance(n, sampled, kappa)),
                            "stationary_gradient_gap": str(gap),
                            "E_kappa2_over_S": str(target),
                            "gap_over_E_kappa2_over_S": str(ratio),
                        }
                    )

    minimum = min(ratios)
    # eta=1/128 and E<=16 imply (1-(1-eta)^E)/(1+(1-eta)^E) >= 1/256.
    # For S<=N/2, (N-S)/(N-1)>=1/2, and the exact grid minimum is stronger.
    assert minimum >= Fraction(1, 512)
    assert finite_minimum >= Fraction(1, 512)
    assert len(rows) == sum(n // 2 for n in N_VALUES) * len(E_VALUES) * len(KAPPA_VALUES)
    result = {
        "construction": "two equal client groups, Li(theta)=1/2*(theta-z_i)^2, z_i=+-kappa",
        "global_gradient": "grad L(theta)=theta",
        "algorithm": "uniform-without-replacement partial FedAvg; exact E-step local GD",
        "eta": str(ETA),
        "rounds_from_zero": ROUNDS,
        "L": 1,
        "N_values": list(N_VALUES),
        "S_range": "1..N/2 for each N",
        "E_values": list(E_VALUES),
        "kappa_values": [str(k) for k in KAPPA_VALUES],
        "subset_count_cells": len(subset_checks),
        "parameter_cells": len(rows),
        "hypergeometric_identity_checked": True,
        "minimum_exact_ratio": str(minimum),
        "minimum_4096_round_ratio": str(finite_minimum),
        "certificate": "stationary gradient gap >= (1/512)*E*kappa^2/S on every executed cell",
        "rows": rows,
    }
    out = Path(__file__).with_name("theorem55_quadratic_scope_results.json")
    out.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps({k: v for k, v in result.items() if k != "rows"}, indent=2))


if __name__ == "__main__":
    run()