import json
import math
import random
import numpy as np

# Barrier-controlled basin switching MVP.
# Potential U(x)=(x^2-1)^2/4 has minima +/-1 and saddle 0.
# Delta V = U(0)-U(+/-1)=1/4; U''(min)=2, U''(saddle)=-1.
# For dX=-U'(X)dt+sqrt(2 eps)dW, Kramers predicts
# k(eps) ~= sqrt(2)/(2*pi)*exp(-0.25/eps).

SEED = 1306
DT = 0.01
N_PATHS = 220
T = 1200.0
STEPS = int(T / DT)
BARRIER = 0.25
PREF = math.sqrt(2.0) / (2.0 * math.pi)


def potential(x):
    return 0.25 * (x*x - 1.0)**2


def grad(x):
    return x*x*x - x


def simulate_rate(eps, seed, n=N_PATHS, t=T):
    """Estimate committed basin switching rate using +/-0.5 hysteresis."""
    rng = np.random.default_rng(seed)
    x = np.full(n, -1.0)
    state = np.full(n, -1, dtype=np.int8)
    switches = 0
    # Count only transitions after reaching the opposite committed region.
    for _ in range(int(t / DT)):
        x += -DT * grad(x) + math.sqrt(2.0 * eps * DT) * rng.standard_normal(n)
        new = state.copy()
        new[(state == -1) & (x > 0.5)] = 1
        new[(state == 1) & (x < -0.5)] = -1
        switches += int(np.count_nonzero(new != state))
        state = new
    # Each path has an opportunity to switch in either direction; this is a
    # conservative event rate and is comparable across noise settings.
    return switches / (n * t), switches


def feedback_rate(q, seed, n=N_PATHS, t=T):
    """Use eps=DeltaV/q while in the escape-to-target mode."""
    eps = np.clip(BARRIER / q, 1e-4, 0.20)
    return simulate_rate(float(eps), seed, n, t)


def linear_fit(x, y):
    return float(np.polyfit(np.asarray(x), np.asarray(y), 1)[0])


def main():
    random.seed(SEED)
    np.random.seed(SEED)

    # Prediction 1: log transition rate against 1/eps has slope -DeltaV.
    eps_values = np.array([0.045, 0.055, 0.070, 0.090, 0.120, 0.160])
    raw = []
    for i, eps in enumerate(eps_values):
        rate, events = simulate_rate(float(eps), SEED + i)
        raw.append({"eps": float(eps), "rate": rate, "events": events})
    usable = [r for r in raw if r["events"] >= 3]
    slope = linear_fit([1/r["eps"] for r in usable], [math.log(r["rate"]) for r in usable])
    slope_error = abs(slope + BARRIER) / BARRIER

    # Prediction 2: feedback makes the exponential part exp(-q), so
    # log(rate/PREF) versus q should have slope -1.
    q_values = np.array([1.6, 2.0, 2.5, 3.0, 3.5, 4.0])
    fb = []
    for i, q in enumerate(q_values):
        rate, events = feedback_rate(float(q), SEED + 100 + i)
        fb.append({"q": float(q), "eps": BARRIER/float(q), "rate": rate, "events": events})
    usable_fb = [r for r in fb if r["events"] >= 3]
    feedback_slope = linear_fit([r["q"] for r in usable_fb],
                                [math.log(r["rate"]/PREF) for r in usable_fb])
    feedback_error = abs(feedback_slope + 1.0)

    # Prediction 3: changing q by delta q changes rate ratio by exp(-delta q).
    q_lo, q_hi = 2.0, 3.5
    lo = next(r for r in fb if r["q"] == q_lo)
    hi = next(r for r in fb if r["q"] == q_hi)
    observed_ratio = hi["rate"] / lo["rate"]
    predicted_ratio = math.exp(-(q_hi-q_lo))
    ratio_factor = observed_ratio / predicted_ratio if predicted_ratio else float("inf")

    # Small optimizer comparison: equal Euler updates, beginning in the left
    # basin and targeting the right basin. Fixed-noise SGD is compared with
    # barrier feedback (escape q=2 while left, safe q=8 after reaching right).
    def train_controller(mode, seed, steps=120000):
        rng = np.random.default_rng(seed)
        x = -1.0
        reached = None
        switches = 0
        for k in range(steps):
            if mode == "fixed":
                eps = BARRIER / 3.0
            else:
                # Known toy barrier estimate; the current basin is identified
                # by sign and target is the lower-loss right basin.
                eps = BARRIER / (2.0 if x < 0 else 8.0)
            x += -DT * grad(x) + math.sqrt(2*eps*DT) * rng.standard_normal()
            if reached is None and x > 0.5:
                reached = k + 1
            if x > 0.5 and x < 0.0:
                switches += 1
        return reached, x, switches
    baseline = train_controller("fixed", SEED + 500)
    idea = train_controller("controlled", SEED + 501)

    result = {
        "known_barrier": BARRIER,
        "kramers_prefactor": PREF,
        "prediction_1_rate_slope": {"predicted": -BARRIER, "observed": slope,
                                    "relative_error": slope_error, "data": raw},
        "prediction_2_feedback_slope": {"predicted": -1.0, "observed": feedback_slope,
                                         "absolute_error": feedback_error, "data": fb},
        "prediction_3_ratio": {"q_low": q_lo, "q_high": q_hi,
                               "predicted_rate_ratio": predicted_ratio,
                               "observed_rate_ratio": observed_ratio,
                               "observed_over_predicted": ratio_factor},
        "mini_optimizer": {"baseline_fixed_noise": baseline,
                            "barrier_controlled": idea,
                            "steps": 120000, "dt": DT},
        "criterion": {"slope_tolerance_fraction": 0.20,
                      "feedback_slope_tolerance": 0.20,
                      "ratio_factor_tolerance": [0.5, 2.0]},
        "worked": bool(slope_error <= 0.20 and feedback_error <= 0.20 and
                        0.5 <= ratio_factor <= 2.0)
    }
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()