import json
from pathlib import Path
import numpy as np

# Exact discrete analogue of a 1-D double-well sampler.
# q_beta(x) proportional to exp(-beta E(x) + c beta F(x)); c=0 is equilibrium.
# The additional feature F makes the sampler nonequilibrium while retaining an
# analytically differentiable beta-conditioned distribution.


def energy(x):
    return (x * x - 1.0) ** 2 + 0.08 * x


def feature(x):
    return x  # not collinear with E on the finite state space


def moments(beta, c=0.0):
    x = np.linspace(-2.2, 2.2, 801)
    E = energy(x)
    F = feature(x)
    logits = -beta * E + c * beta * F
    logits -= logits.max()
    p = np.exp(logits)
    p /= p.sum()
    mu = np.sum(p * E)
    ec = E - mu
    fc = F - np.sum(p * F)
    # score = d_beta log q = -(E-mu) + c(F-E[F])
    score = -ec + c * fc
    V = np.sum(p * ec**2)
    I = np.sum(p * score**2)
    response = np.sum(p * ec * score)  # d E_q[E] / d beta
    return dict(x=x, E=E, p=p, mu=mu, V=V, I=I,
                response=response, score=score)


def finite_response(beta, h, c):
    return (moments(beta + h, c)['mu'] - moments(beta - h, c)['mu']) / (2*h)


def controller_step(beta, delta, dbmax, c, eps=1e-12):
    m = moments(beta, c)
    raw = delta / np.sqrt(m['V'] * m['I'] + eps)
    return min(raw, dbmax), raw, m


def exact_prediction_checks():
    betas = np.linspace(0.2, 4.0, 20)
    out = {}
    for c, label in [(0.0, 'equilibrium'), (0.65, 'nonequilibrium')]:
        rows = []
        for b in betas:
            m = moments(b, c)
            rho = abs(m['response']) / np.sqrt(m['V'] * m['I'])
            h = 1e-4
            rho_fd = abs(finite_response(b, h, c)) / np.sqrt(m['V'] * m['I'])
            rows.append((rho, rho_fd, m['V'], m['I']))
        a = np.asarray(rows)
        out[label] = {
            'rho_mean': float(a[:, 0].mean()),
            'rho_max': float(a[:, 0].max()),
            'rho_fd_max': float(a[:, 1].max()),
            'inequality_max_excess': float(max(0, a[:, 0].max() - 1.0)),
            'mean_V': float(a[:, 2].mean()),
            'mean_I': float(a[:, 3].mean()),
        }
    return out


def scaling_check(c=0.0):
    # At fixed beta, the proposed step must scale linearly with delta and
    # actual first-order energy displacement should have slope <= 1.
    beta = 1.4
    dbmax = 10.0
    deltas = np.array([0.002, 0.004, 0.008, 0.016, 0.032])
    rows = []
    for d in deltas:
        db, raw, m = controller_step(beta, d, dbmax, c)
        actual = abs(moments(beta + db, c)['mu'] - m['mu'])
        predicted = db * np.sqrt(m['V'] * m['I'])
        rows.append([d, db, actual, predicted, actual / d, actual / predicted])
    a = np.asarray(rows)
    slope = np.polyfit(np.log(deltas), np.log(a[:, 1]), 1)[0]
    return {
        'rows': a.tolist(), 'step_loglog_slope': float(slope),
        'max_actual_over_delta': float(a[:, 4].max()),
        'max_actual_over_predicted': float(a[:, 5].max()),
    }


def schedule_comparison(c=0.0):
    # Same start/end beta and number of updates. Geometric beta increments are
    # compared with CR-limited increments; CR uses the requested annealing sign.
    start, end, n = 0.2, 4.0, 40
    geometric = np.linspace(start, end, n + 1)
    results = {}
    for delta in [0.01, 0.03, 0.06]:
        beta = start
        cr_path = [beta]
        cr_jumps = []
        for _ in range(n):
            db, _, m = controller_step(beta, delta, (end-start)/n*4, c)
            # Do not overshoot the requested endpoint.
            db = min(db, end-beta)
            new_mu = moments(beta + db, c)['mu']
            cr_jumps.append(abs(new_mu - m['mu']))
            beta += db
            cr_path.append(beta)
            if beta >= end - 1e-12:
                break
        geo_mu = [moments(b, c)['mu'] for b in geometric]
        geo_jumps = np.abs(np.diff(geo_mu))
        results[str(delta)] = {
            'cr_updates_to_endpoint': len(cr_jumps),
            'cr_final_beta': float(beta),
            'cr_max_jump': float(max(cr_jumps) if cr_jumps else 0),
            'geometric_max_jump': float(geo_jumps.max()),
            'cr_mean_jump': float(np.mean(cr_jumps) if cr_jumps else 0),
            'geometric_mean_jump': float(geo_jumps.mean()),
            'cr_path_first_last': [float(cr_path[0]), float(cr_path[-1])],
        }
    return results


def main():
    checks = exact_prediction_checks()
    scaling_eq = scaling_check(0.0)
    scaling_neq = scaling_check(0.65)
    schedules = schedule_comparison(0.0)
    report = {
        'prediction_1_saturation_and_bound': checks,
        'prediction_2_delta_scaling_equilibrium': scaling_eq,
        'prediction_2_delta_scaling_nonequilibrium': scaling_neq,
        'prediction_3_schedule_jump_control': schedules,
        'interpretation': {
            'equilibrium_prediction': 'rho=1 exactly because score is -(E-mu) and I=Var(E)',
            'nonequilibrium_prediction': 'rho<1 when the beta score contains an F component not determined by E',
            'controller_prediction': 'actual energy jump is approximately <= delta for small delta and controller steps scale linearly with delta',
        }
    }
    Path('results.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()