import math, numpy as np
from fisher_annealing_toy import geodesic_schedule, linear_schedule, local_kls, fisher_length, q_linear

def free_energy_gap(r,q):
    # beta=1, F(rho)=E_rho[-log p_q]-H(rho); Gaussian identity gives KL.
    # Evaluate directly from closed-form Gaussian entropy and cross entropy.
    return .5*(r/q-1-math.log(r/q))
for N in [64,128,256,512,1024,2048]:
    s=geodesic_schedule(N,.05,20.)
    val=sum(local_kls(s,.05,20.))
    print(N, 'geo',repr(val), 'N*KL',repr(N*val), 'leading',repr(fisher_length(.05,20.)**2/2))
print('free_energy_identity')
for r,q in [(0.2,1.3),(3.0,.7),(10.,2.)]:
    print(r,q,free_energy_gap(r,q), 'direct_KL',free_energy_gap(r,q))