import json
from pathlib import Path
import numpy as np

# Convex information-coordinate toy for I_q(Y;Z|U)>=H.
# A binary symmetric channel supplies the attainable range, while r is a
# differentiable information-capacity coordinate in [I0,Imax].
def entropy_binary(q):
    q = np.clip(q, 1e-12, 1-1e-12)
    return -(q*np.log(q)+(1-q)*np.log(1-q))

def channel_info(p):
    return float(np.log(2.)-entropy_binary(1-p))

def run_dual(target, r0, rmax, k=2., eta_l=.05, lam_max=8., steps=3000):
    # F(r)=k/2(r-r0)^2; exact inner minimizer is r=min(rmax,r0+lambda/k).
    lam = 0.; ema = r0; trace=[]
    for t in range(steps):
        r = float(np.clip(r0 + lam/k, r0, rmax))
        ema = .9*ema + .1*r
        lam = float(np.clip(lam + eta_l*(target-ema), 0., lam_max))
        if t % 10 == 0: trace.append((t,r,lam))
    return r,lam,trace

def fixed_information_drift(target, fixed, eta=.07, steps=100, clip=8.):
    lam=0.
    for _ in range(steps): lam=np.clip(lam+eta*(target-fixed),0.,clip)
    return float(lam), float(steps*eta*(target-fixed))

def main():
    # p=.55 is the unconstrained channel; pmax defines attainable information.
    r0, rmax = channel_info(.55), channel_info(.995)
    targets=[.5*r0, r0+.35*(rmax-r0), rmax+.10]
    rows=[]
    for name,h in zip(('inactive','binding','saturated'),targets):
        r,lam,_=run_dual(h,r0,rmax)
        rows.append({'regime':name,'H':h,
          'predicted_I':r0 if name=='inactive' else (h if name=='binding' else rmax),
          'observed_I':r,'abs_I_error':abs(r-(r0 if name=='inactive' else (h if name=='binding' else rmax))),
          'predicted_lambda':0. if name=='inactive' else (None if name=='binding' else 8.),
          'observed_lambda':lam})
    sweep=[]
    for frac in (.15,.35,.55,.75,.90):
        h=r0+frac*(rmax-r0); r,lam,_=run_dual(h,r0,rmax)
        sweep.append({'H':h,'observed_I':r,'abs_error':abs(r-h),'lambda':lam})
    drifts=[]
    for delta in (-.08,.05,.30):
        obs,pred=fixed_information_drift(r0+delta,r0)
        drifts.append({'delta_H_minus_I':delta,'predicted_unclipped_lambda':pred,'observed_lambda':obs})
    # Standard fixed coefficient: same primal with lambda=1, cannot track changing H.
    fixed_r=min(rmax,r0+1/2.)
    out={'channel':{'I0':r0,'Imax':rmax},
      'predictions':{'P1':'H<I0 => lambda=0 and I=I0','P2':'I0<H<Imax => I tracks H','P3':'H>Imax => lambda clips and I=Imax','P4':'fixed-I lambda slope=eta*(H-I)'},
      'regimes':rows,'feasible_target_sweep':sweep,'fixed_I_drift':drifts,
      'baseline_fixed_lambda':{'lambda':1.,'I':fixed_r}}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()