#!/usr/bin/env python3
"""Exact counterexample to Claim 3's constant-step stochastic theorem."""

from __future__ import annotations

import json
from fractions import Fraction
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]


def main() -> None:
    # Phi(theta)=theta^2/2, g_hat=theta+xi, xi in {-1,+1} equiprobably.
    # With r=1/4, theta_{s+1}=3 theta_s/4-xi_s/4.  If m_s=E theta_s
    # and v_s=Var(theta_s), then m_{s+1}=3m_s/4 and
    # v_{s+1}=9v_s/16+1/16 exactly.
    mean = Fraction(1)
    variance = Fraction(0)
    rows = []
    for s in range(1, 4097):
        mean *= Fraction(3, 4)
        variance = Fraction(9, 16) * variance + Fraction(1, 16)
        grad_sq = mean * mean + variance
        closed_variance = Fraction(1, 7) * (1 - Fraction(9, 16) ** s)
        assert variance == closed_variance
        assert grad_sq >= Fraction(1, 16)
        rows.append((s, grad_sq))

    result = {
        "schema": "wgf-outer-loop-counterexample-v1",
        "objective": "Phi(theta)=theta^2/2",
        "smoothness_L_Phi": 1,
        "constant_step_r": "1/4",
        "gradient_estimator": "theta + Rademacher noise",
        "bias": 0,
        "variance_sigma_squared": 1,
        "inner_sampling_error": 0,
        "exact_iterations": len(rows),
        "minimum_expected_gradient_squared_after_first_update": "1/16",
        "limit_expected_gradient_squared": "1/7",
        "contradicted_epsilon": "any epsilon_opt < 1/4",
        "claim3_falsified": True,
    }
    (ROOT / "outer_loop_counterexample_results.json").write_text(
        json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()