# Iterative Entropic Renormalization (IER) — Complete Replicable Example
# Author: Mark Lindenhayn

import numpy as np
import matplotlib.pyplot as plt

def normalize_density(p, x):
    z = np.trapz(p, x)
    if z <= 0:
        raise ValueError("Non-positive probability mass.")
    return p / z

def shannon_entropy(p, x):
    p_safe = p + 1e-14
    return -np.trapz(p_safe * np.log(p_safe), x)

def fisher_information(p, x):
    p_safe = p + 1e-14
    dp = np.gradient(p_safe, x, edge_order=2)
    integrand = dp**2 / p_safe
    return np.trapz(integrand, x)

def mean_and_variance(p, x):
    mu = np.trapz(x * p, x)
    var = np.trapz((x - mu)**2 * p, x)
    return mu, var

def gaussian_pdf(x, mu=0.0, sigma=1.0):
    sigma = max(sigma, 1e-14)
    return np.exp(-0.5 * ((x - mu) / sigma)**2) / (np.sqrt(2*np.pi) * sigma)

def kl_divergence_to_gaussian(p, x, mu=0.0, sigma=1.0):
    p = normalize_density(p, x)
    g = gaussian_pdf(x, mu, sigma) + 1e-300
    p_safe = p + 1e-300
    return np.trapz(p_safe * np.log(p_safe / g), x)

def ier_step(p, x, tau=None, k=None):
    p = normalize_density(p, x)
    mu, var = mean_and_variance(p, x)
    sigma = np.sqrt(var + 1e-18)

    if tau is None:
        if k is None:
            raise ValueError("Specify tau or k.")
        tau_eff = k * sigma
    else:
        tau_eff = tau

    mask = np.abs(x) <= tau_eff
    if not np.any(mask):
        raise ValueError("Truncation removed all mass.")

    x_trunc = x[mask]
    p_trunc = normalize_density(p[mask], x_trunc)

    mu_t, var_t = mean_and_variance(p_trunc, x_trunc)
    sigma_t = np.sqrt(var_t + 1e-18)

    x_new = (x_trunc - mu_t) / sigma_t
    p_new = p_trunc * sigma_t

    idx = np.argsort(x_new)
    x_new = x_new[idx]
    p_new = normalize_density(p_new[idx], x_new)

    return p_new, x_new, tau_eff

def run_ier_with_fixed_taus(p0, x0, tau_list):
    p_current = normalize_density(p0, x0)
    x_current = x0.copy()

    ps, xs = [p_current], [x_current]
    entropies = [shannon_entropy(p_current, x_current)]
    fishers = [fisher_information(p_current, x_current)]
    taus_used = [None]

    for tau in tau_list:
        p_current, x_current, tau_eff = ier_step(p_current, x_current, tau=tau)
        ps.append(p_current)
        xs.append(x_current)
        entropies.append(shannon_entropy(p_current, x_current))
        fishers.append(fisher_information(p_current, x_current))
        taus_used.append(tau_eff)

    return ps, xs, entropies, fishers, taus_used

def main():
    x = np.linspace(-20, 20, 40001)

    GAMMA_CAUCHY = 0.2840168571
    p0 = 1/(np.pi * GAMMA_CAUCHY * (1+(x/GAMMA_CAUCHY)**2))
    p0 = normalize_density(p0, x)

    tau_list = [1.0, 0.85, 0.75, 0.65]
    ps, xs, entropies, fishers, taus = run_ier_with_fixed_taus(p0, x, tau_list)
    kl_vals = [kl_divergence_to_gaussian(p, xx) for p, xx in zip(ps, xs)]

    with open("ier_results.txt","w") as f:
        f.write("Iter | tau | Entropy | Fisher | KL
")
        for n in range(len(ps)):
            f.write(f"{n} | {taus[n]} | {entropies[n]} | {fishers[n]} | {kl_vals[n]}
")

    plt.figure(figsize=(10,6))
    for n,(xx,p) in enumerate(zip(xs,ps)):
        plt.plot(xx,p,label=f"n={n}")
    plt.legend(); plt.grid(); plt.title("IER Density Evolution")
    plt.savefig("ier_plot.png",dpi=200)

if __name__ == "__main__":
    main()