#!/usr/bin/env python3
"""Finite checks of information-constrained translation recovery.

Units: k_B T = 1. The program independently integrates truncated Gaussian
means and minimizes conditional quadratic work, then compares with the
closed-form theorem. This tests a specified recovery stage, not a full device.
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import numpy as np
from scipy.integrate import quad
from scipy.optimize import minimize_scalar
from scipy.special import ndtr, xlogy

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

def h(p: float) -> float:
    return float(-xlogy(p,p)-xlogy(1-p,1-p))

def run() -> dict:
    sigma=math.sqrt(.5)
    records=[]
    max_mean_residual=0.; max_opt_residual=0.; min_ep=math.inf
    for u in [0., .01, .1, .5, 1., 2., 4., 9., 16.]:
        a=math.sqrt(u); z=a/sigma
        p=float(ndtr(-z)); q=float(ndtr(z))
        phi=math.exp(-z*z/2)/math.sqrt(2*math.pi)
        mu=np.array([a-sigma*phi/p, a+sigma*phi/q])
        probs=np.array([p,q])
        density=lambda x: math.exp(-(x-a)**2)/(math.sqrt(math.pi))
        integrated=np.array([
            quad(lambda x:x*density(x),-np.inf,0,epsabs=1e-14)[0]/p,
            quad(lambda x:x*density(x),0,np.inf,epsabs=1e-13)[0]/q])
        residual=float(np.max(np.abs(integrated-mu)))
        max_mean_residual=max(max_mean_residual,residual)
        assert residual<1e-7,(u,integrated,mu)
        base=u+h(p)
        for radius in [0., .05, .2, .5, 1., 2., 4., 8.]:
            d=-np.clip(mu,-radius,radius)
            improvement=float(np.dot(probs,mu**2-np.maximum(np.abs(mu)-radius,0)**2))
            actual=base+float(np.dot(probs,2*d*mu+d*d))
            analytic=base-improvement
            numeric=base
            for weight,mean in zip(probs,mu):
                if radius>0:
                    result=minimize_scalar(lambda shift:2*shift*mean+shift*shift,
                                           bounds=(-radius,radius),method='bounded',
                                           options={'xatol':1e-12})
                    value=min(result.fun, -2*radius*abs(mean)+radius**2, 0.)
                else:value=0.
                numeric+=float(weight)*value
            label=np.maximum(a-radius,0)**2+h(p)
            max_opt_residual=max(max_opt_residual,abs(numeric-analytic),abs(actual-analytic))
            min_ep=min(min_ep,analytic)
            assert abs(actual-analytic)<1e-12
            assert abs(numeric-analytic)<1e-9
            assert analytic>=-1e-12 and analytic<=label+1e-12
            records.append(dict(signal_energy=u, displacement_budget=radius,
                                no_recovery_ep=base,label_only_ep=float(label),
                                sample_controlled_ep=analytic,
                                optimizer_residual=float(numeric-analytic)))
        unlimited=h(p)-phi*phi/(2*p*q)
        assert abs(unlimited-(base-float(np.dot(probs,mu*mu))))<1e-12
    # A finite-state conditional-expectation test independent of Gaussianity.
    rng=np.random.default_rng(20261028+271)
    trials=20_000
    weights=rng.dirichlet(np.ones(6),size=trials)
    means=rng.normal(0,2,size=(trials,6))
    radii=rng.uniform(0,3,size=(trials,1))
    fine=np.sum(weights*(means**2-np.maximum(np.abs(means)-radii,0)**2),axis=1)
    w=weights.reshape(trials,3,2).sum(axis=2)
    coarse=(weights*means).reshape(trials,3,2).sum(axis=2)/w
    coarse_gain=np.sum(w*(coarse**2-np.maximum(np.abs(coarse)-radii,0)**2),axis=1)
    assert np.min(fine-coarse_gain)>-1e-12
    return dict(units='k_B T = 1',gaussian_cases=len(records),
                conditional_information_refinements=trials,
                maximum_conditional_mean_quadrature_residual=max_mean_residual,
                maximum_quadratic_optimization_residual=max_opt_residual,
                minimum_observed_reset_ep=min_ep,
                zero_signal=dict(label_only_ep=math.log(2),
                                 sample_controlled_ep=math.log(2)-1/math.pi,
                                 recovered_work=1/math.pi,
                                 required_displacement=1/math.sqrt(math.pi)),
                cases=records)

if __name__=='__main__':
    report=run();(ROOT/'results').mkdir(exist_ok=True)
    (ROOT/'results'/'recovery.json').write_text(json.dumps(report,indent=2),encoding='utf-8')
    print(json.dumps({k:v for k,v in report.items() if k!='cases'},indent=2))