#!/usr/bin/env python3
"""Exact wider Bayes-risk family for the FedDPO heterogeneity lower bound.

The hidden per-coordinate preference mean is theta=kappa*(p-1/2), with p
uniform on a finite grid.  A participating client supplies one preference
bit X~Bernoulli(p), encoded as Y=kappa*(X-1/2).  The posterior mean is the
Bayes-optimal squared-loss estimator.  Because the likelihood depends only on
the number of positive bits, the complete posterior risk is an exact
Fraction-valued dynamic program over S+1 observation counts, not Monte Carlo.
"""

from __future__ import annotations

import json
from fractions import Fraction
from math import comb


GRID_SIZES = (17, 33, 65)
POPULATIONS = (16, 32, 64)
E_VALUES = (1, 2, 4, 8, 16)
KAPPAS = (Fraction(1, 4), Fraction(1, 2), Fraction(1), Fraction(2))


def normalized_bayes_risk(grid_size: int, sampled: int) -> Fraction:
    """Exact risk for kappa=1 and one coordinate."""
    risk = Fraction(0)
    for positives in range(sampled + 1):
        joint: list[tuple[Fraction, Fraction]] = []
        for index in range(1, grid_size + 1):
            p = Fraction(index, grid_size + 1)
            theta = p - Fraction(1, 2)
            mass = Fraction(1, grid_size) * comb(sampled, positives)
            mass *= p**positives * (1 - p) ** (sampled - positives)
            joint.append((mass, theta))
        total = sum(mass for mass, _ in joint)
        posterior_mean = sum(mass * theta for mass, theta in joint) / total
        risk += sum(mass * (theta - posterior_mean) ** 2 for mass, theta in joint)
    return risk


def main() -> None:
    rows: list[dict[str, object]] = []
    for grid_size in GRID_SIZES:
        for population in POPULATIONS:
            for sampled in range(1, population + 1):
                unit = normalized_bayes_risk(grid_size, sampled)
                for e_local in E_VALUES:
                    for kappa in KAPPAS:
                        risk = e_local * kappa * kappa * unit
                        target = Fraction(e_local) * kappa * kappa / sampled
                        rows.append({
                            "prior_grid": grid_size,
                            "N": population,
                            "S": sampled,
                            "E": e_local,
                            "kappa": str(kappa),
                            "bayes_risk": str(risk),
                            "target_E_kappa2_over_S": str(target),
                            "ratio": str(risk / target),
                            "ratio_decimal": float(risk / target),
                        })

    ratios = [Fraction(row["ratio"]) for row in rows]
    result = {
        "construction": "finite-grid Bernoulli preference family with exact posterior mean",
        "prior_grids": list(GRID_SIZES),
        "populations": list(POPULATIONS),
        "all_S_values": {str(n): list(range(1, n + 1)) for n in POPULATIONS},
        "E_values": list(E_VALUES),
        "kappa_values": [str(k) for k in KAPPAS],
        "base_population_grid_cells": len(GRID_SIZES) * sum(POPULATIONS),
        "executed_parameter_cells": len(rows),
        "min_ratio_risk_over_E_kappa2_over_S": str(min(ratios)),
        "min_ratio_decimal": float(min(ratios)),
        "max_ratio_decimal": float(max(ratios)),
        "rows": rows,
        "bayes_optimality": "posterior mean minimizes conditional squared risk, so every estimator has at least this Bayes risk",
    }
    print(json.dumps(result, indent=2, sort_keys=True))
    assert result["min_ratio_decimal"] > 0.05


if __name__ == "__main__":
    main()