#!/usr/bin/env python3
"""Exact toy verification of Fisher-geodesic finite-step annealing.

Path: p_theta = Normal(0, q(theta)), q(theta)=q0*exp(log_ratio*theta).
For this path, I(theta)=0.5*log_ratio**2 and equal Fisher arc length is
therefore uniform theta.  To make the comparison nontrivial, the baseline
uses a linear variance path q(theta)=q0+(q1-q0)*theta, whose Fisher speed
varies strongly; the proposed schedule is the inverse-arc schedule on that
same path (equivalently, geometric spacing of q).
"""
import csv
import json
import math
from pathlib import Path
import numpy as np


def q_linear(theta, q0, q1):
    return q0 + (q1-q0)*np.asarray(theta)


def fisher_speed(theta, q0, q1):
    q = q_linear(theta, q0, q1)
    dq = q1-q0
    return abs(dq)/(math.sqrt(2.0)*q)


def fisher_length(q0, q1):
    return abs(math.log(q1/q0))/math.sqrt(2.0)


def geodesic_schedule(n, q0, q1):
    # Equal arc length: log(q_k/q0)=k*log(q1/q0)/n.
    k = np.arange(n+1, dtype=float)
    q = q0 * (q1/q0)**(k/n)
    return (q-q0)/(q1-q0)


def linear_schedule(n):
    return np.linspace(0.0, 1.0, n+1)


def kl_zero_gaussians(q_from, q_to):
    # KL(N(0,q_from) || N(0,q_to)).
    r = q_from/q_to
    return 0.5*(r - 1.0 - math.log(r))


def local_kls(schedule, q0, q1):
    q = q_linear(schedule, q0, q1)
    return np.array([kl_zero_gaussians(q[i], q[i+1]) for i in range(len(q)-1)])


def arc_steps(schedule, q0, q1):
    q = q_linear(schedule, q0, q1)
    return np.abs(np.log(q[1:]/q[:-1]))/math.sqrt(2.0)


def summarize(schedule, q0, q1):
    kls = local_kls(schedule, q0, q1)
    return {
        "cumulative_kl": float(kls.sum()),
        "local_kl_cv": float(kls.std()/kls.mean()),
        "arc_step_cv": float(arc_steps(schedule,q0,q1).std()/arc_steps(schedule,q0,q1).mean()),
        "local_kls": kls.tolist(),
    }


def fit_slope(ns, vals):
    return float(np.polyfit(np.log(ns), np.log(vals), 1)[0])


def run():
    # Exact Fisher identity check against Monte Carlo score estimates.
    rng = np.random.default_rng(2733)
    q0, q1 = 0.05, 20.0
    theta = 0.37
    q = q_linear(theta,q0,q1)
    dq = q1-q0
    x = rng.normal(0.0, math.sqrt(q), size=1_000_000)
    # d/dtheta log N(x;0,q(theta)) = q'*(x^2-q)/(2q^2)
    score = dq*(x*x-q)/(2*q*q)
    mc_fisher = float(np.mean(score*score))
    exact_fisher = fisher_speed(theta,q0,q1)**2
    fisher_relerr = abs(mc_fisher-exact_fisher)/exact_fisher

    ns = [4,8,16,32,64,128]
    ranges = [(1.0,2.0),(1.0,10.0),(0.05,20.0)]
    rows=[]
    for a,b in ranges:
        L=fisher_length(a,b)
        for n in ns:
            base=summarize(linear_schedule(n),a,b)
            geo=summarize(geodesic_schedule(n,a,b),a,b)
            rows.append({"q0":a,"q1":b,"N":n,"L":L,
                         "baseline_cumulative_kl":base["cumulative_kl"],
                         "geodesic_cumulative_kl":geo["cumulative_kl"],
                         "baseline_local_kl_cv":base["local_kl_cv"],
                         "geodesic_local_kl_cv":geo["local_kl_cv"],
                         "baseline_arc_cv":base["arc_step_cv"],
                         "geodesic_arc_cv":geo["arc_step_cv"],
                         "predicted_leading_kl":L*L/(2*n)})

    # Fit the asymptotic exponent on the finer four grids for the wide range.
    wide=[r for r in rows if r["q0"]==0.05 and r["N"]>=16]
    slope_geo=fit_slope(np.array([r["N"] for r in wide]),np.array([r["geodesic_cumulative_kl"] for r in wide]))
    slope_base=fit_slope(np.array([r["N"] for r in wide]),np.array([r["baseline_cumulative_kl"] for r in wide]))

    # Mechanism summaries at N=32, including the predicted leading coefficient.
    mech=[]
    for a,b in ranges:
        r=[z for z in rows if z["q0"]==a and z["q1"]==b and z["N"]==32][0]
        mech.append({"range":f"{a:g}->{b:g}","fisher_length":r["L"],
                     "predicted_equal_arc_cv":0.0,
                     "observed_geodesic_arc_cv":r["geodesic_arc_cv"],
                     "observed_geodesic_local_kl_cv":r["geodesic_local_kl_cv"],
                     "observed_linear_local_kl_cv":r["baseline_local_kl_cv"],
                     "predicted_kl":r["predicted_leading_kl"],
                     "observed_kl":r["geodesic_cumulative_kl"],
                     "ratio_linear_over_geodesic":r["baseline_cumulative_kl"]/r["geodesic_cumulative_kl"]})

    out={"fisher_check":{"mc_fisher":mc_fisher,"exact_fisher":exact_fisher,"relative_error":fisher_relerr},
         "predictions":{"equal_arc_arc_cv":mech,
                         "wide_range_loglog_slope_geodesic":slope_geo,
                         "wide_range_loglog_slope_linear_baseline":slope_base,
                         "asymptotic_target_slope":-1.0},
         "rows":rows}
    Path("toy_results.json").write_text(json.dumps(out,indent=2))
    with open("toy_results.csv","w",newline="") as f:
        w=csv.DictWriter(f,fieldnames=rows[0].keys()); w.writeheader(); w.writerows(rows)
    print(json.dumps({"fisher_check":out["fisher_check"],"predictions":out["predictions"]},indent=2))

if __name__ == "__main__": run()