#!/usr/bin/env python3
"""Exact finite-prior Bayes-risk certificate for the participation lower bound.

The construction is deliberately small and discrete.  For each of E independent
local coordinates, a hidden client-level signal theta is either +kappa or
-kappa.  A participating client returns theta plus an independent
Rademacher heterogeneity term, also of magnitude kappa.  The learner sees S
clients and estimates theta.  A uniform prior over theta and exhaustive
enumeration of all N-client noise assignments and all S-client subsets gives
the Bayes risk.  Bayes risk is a lower bound for every estimator under this
finite prior.

All probabilities and risks are computed with Fraction; no floating-point
calculation is used for the certificate.
"""

from __future__ import annotations

import itertools
import json
from collections import defaultdict
from fractions import Fraction
from math import comb
from pathlib import Path


def exact_unit_risk(n_clients: int, sampled: int) -> dict[str, object]:
    """Enumerate the two worlds, all noise assignments, and all sample sets.

    Values are normalized by kappa, so the returned risk is multiplied by
    kappa**2 for an arbitrary positive kappa.  The posterior-mean estimator is
    Bayes optimal for squared loss.  For each observation o with joint masses
    p_plus and p_minus, its contribution is
        4 * p_plus * p_minus / (p_plus + p_minus),
    which is the exact posterior Bayes risk contribution for theta in {-1,+1}.
    """
    if not (1 <= sampled <= n_clients):
        raise ValueError("sampled must be in [1, n_clients]")

    subsets = tuple(itertools.combinations(range(n_clients), sampled))
    subset_count = len(subsets)
    # obs -> [joint mass under theta=+1, joint mass under theta=-1]
    masses: dict[tuple[int, ...], list[Fraction]] = defaultdict(
        lambda: [Fraction(0), Fraction(0)]
    )
    per_world = Fraction(1, 2 * (1 << n_clients) * subset_count)
    for theta_index, theta in enumerate((1, -1)):
        for noise_mask in range(1 << n_clients):
            noise = tuple(1 if (noise_mask >> i) & 1 else -1 for i in range(n_clients))
            values = tuple(theta + z for z in noise)
            for subset in subsets:
                observation = tuple(values[i] for i in subset)
                masses[observation][theta_index] += per_world

    risk = Fraction(0)
    ambiguous_mass = Fraction(0)
    for p_plus, p_minus in masses.values():
        total = p_plus + p_minus
        if p_plus and p_minus:
            ambiguous_mass += total
            risk += Fraction(4) * p_plus * p_minus / total

    expected_formula = Fraction(1, 1 << sampled)
    if risk != expected_formula:
        raise AssertionError((n_clients, sampled, risk, expected_formula))

    return {
        "n_clients": n_clients,
        "sampled": sampled,
        "enumerated_noise_assignments_per_world": 1 << n_clients,
        "enumerated_subsets": subset_count,
        "joint_world_subset_cases": 2 * (1 << n_clients) * subset_count,
        "distinct_observations": len(masses),
        "ambiguous_observation_mass": str(ambiguous_mass),
        "risk_over_kappa_squared": str(risk),
        "risk_over_kappa_squared_decimal": float(risk),
        "closed_form_cross_check": str(expected_formula),
    }


def main() -> None:
    # These N values cover several finite population sizes; every S from one
    # client through full participation is enumerated for each one.
    population_sizes = (5, 8, 10)
    e_values = (1, 2, 4, 8)
    kappa_values = (Fraction(1, 4), Fraction(1, 2), Fraction(1), Fraction(2))

    base_rows: list[dict[str, object]] = []
    for n_clients in population_sizes:
        for sampled in range(1, n_clients + 1):
            base_rows.append(exact_unit_risk(n_clients, sampled))

    # Extend the exact base risks over executed E and kappa regimes.  The E
    # coordinates are independent, so squared risks add exactly; kappa scales
    # the normalized risk by kappa**2.
    cells: list[dict[str, object]] = []
    min_ratio: Fraction | None = None
    for base in base_rows:
        n_clients = int(base["n_clients"])
        sampled = int(base["sampled"])
        unit_risk = Fraction(str(base["risk_over_kappa_squared"]))
        for e_local in e_values:
            for kappa in kappa_values:
                risk = e_local * kappa * kappa * unit_risk
                target_scale = Fraction(e_local) * kappa * kappa / sampled
                ratio = risk / target_scale
                min_ratio = ratio if min_ratio is None else min(min_ratio, ratio)
                cells.append(
                    {
                        "N": n_clients,
                        "S": sampled,
                        "E": e_local,
                        "kappa": str(kappa),
                        "bayes_risk": str(risk),
                        "target_E_kappa2_over_S": str(target_scale),
                        "ratio_to_target": str(ratio),
                    }
                )

    assert min_ratio is not None
    certificate_constant = Fraction(5, 512)
    if min_ratio < certificate_constant:
        raise AssertionError((min_ratio, certificate_constant))

    result = {
        "construction": "two-world Rademacher heterogeneity, exact finite-prior Bayes risk",
        "population_sizes": list(population_sizes),
        "sample_sizes_per_population": {
            str(n): list(range(1, n + 1)) for n in population_sizes
        },
        "E_values": list(e_values),
        "kappa_values": [str(k) for k in kappa_values],
        "base_rows": base_rows,
        "executed_parameter_cells": len(cells),
        "cells": cells,
        "min_ratio_risk_over_E_kappa2_over_S": str(min_ratio),
        "finite_family_certificate": f"risk >= ({certificate_constant}) * E*kappa^2/S",
        "bayes_optimality": (
            "For squared loss, posterior mean minimizes conditional risk; "
            "therefore every estimator has expected risk at least this Bayes risk."
        ),
    }
    out = Path(__file__).with_name("bayes_lower_bound_results.json")
    out.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps({
        "output": str(out),
        "base_rows": len(base_rows),
        "executed_parameter_cells": len(cells),
        "min_ratio": str(min_ratio),
        "certificate": result["finite_family_certificate"],
    }, sort_keys=True))


if __name__ == "__main__":
    main()