import json
import random
from pathlib import Path
import numpy as np


def simulate(theta0, theta1, t_end, dt, kappa, theta_c, slope=1.0, mode="baseline"):
    n = int(round(t_end / dt)) + 1
    ts = np.arange(n) * dt
    desired = theta0 + (theta1 - theta0) * np.minimum(ts / t_end, 1.0)
    r = (theta1 - theta0) / t_end
    if mode == "baseline":
        command = desired.copy()
    elif mode == "minus":
        command = desired - r / kappa
    elif mode == "plus":
        command = desired + r / kappa
    else:
        raise ValueError(mode)
    eff = np.empty(n)
    eff[0] = theta0
    for i in range(1, n):
        eff[i] = eff[i - 1] + dt * kappa * (command[i - 1] - eff[i - 1])
    lam_eff = slope * (eff - theta_c)
    lam_cmd = slope * (command - theta_c)

    def first_cross(x):
        ix = np.flatnonzero(x >= 0)
        if len(ix) == 0:
            return float("nan")
        i = int(ix[0])
        if i == 0:
            return float(ts[0])
        # Linear interpolation makes timing less dependent on dt.
        return float(ts[i - 1] + dt * (-x[i - 1]) / (x[i] - x[i - 1]))

    target_cross = (theta_c - theta0) / r
    return {
        "target_cross": target_cross,
        "command_cross": first_cross(lam_cmd),
        "effective_cross": first_cross(lam_eff),
        "delay": first_cross(lam_eff) - target_cross,
        "ts": ts, "desired": desired, "eff": eff,
    }


def linear_fit(x, y):
    x, y = np.asarray(x), np.asarray(y)
    coef = np.polyfit(x, y, 1)
    pred = np.polyval(coef, x)
    ssr = np.sum((y - pred) ** 2)
    sst = np.sum((y - y.mean()) ** 2)
    return float(coef[0]), float(coef[1]), float(1 - ssr / sst)


def math_verification():
    # Prediction 1: after ramp transients, theta_eff = theta_desired-r/kappa.
    kappa = 4.0
    t_end, dt = 20.0, 0.002
    out = simulate(0, 2, t_end, dt, kappa, 99, mode="baseline")
    r = 2 / t_end
    tail_error = np.mean((out["desired"][-1000:] - out["eff"][-1000:]))
    predicted_offset = r / kappa

    # Prediction 2: threshold delay is (approximately) r/(kappa*slope).
    kappas = np.array([1., 2., 4., 8., 16., 32.])
    delays = np.array([simulate(0, 2, 20, .002, k, 1, mode="baseline")["delay"] for k in kappas])
    slope, intercept, r2 = linear_fit(1 / kappas, delays)
    predicted_slope = 1.0  # crossing delay=(r/kappa)/(|d theta_c/dt|)=1/kappa

    # Prediction 3: inverse feed-forward reduces delay; printed proposal sign
    # is tested separately because for an increasing ramp it has the wrong sign.
    base = np.array([simulate(0, 2, 20, .002, k, 1, mode="baseline")["delay"] for k in kappas])
    minus = np.array([simulate(0, 2, 20, .002, k, 1, mode="minus")["delay"] for k in kappas])
    plus = np.array([simulate(0, 2, 20, .002, k, 1, mode="plus")["delay"] for k in kappas])
    return {
        "offset_observed": float(tail_error), "offset_predicted": predicted_offset,
        "delay_fit_slope": slope, "delay_fit_intercept": intercept,
        "delay_predicted_slope": predicted_slope, "delay_r2": r2,
        "kappas": kappas.tolist(), "baseline_delays": base.tolist(),
        "minus_delays": minus.tolist(), "plus_delays": plus.tolist(),
        "minus_reduction_fraction": float(1 - np.nanmean(np.abs(minus)) / np.nanmean(np.abs(base))),
        "plus_reduction_fraction": float(1 - np.nanmean(np.abs(plus)) / np.nanmean(np.abs(base))),
    }


def train_toy(seed=7, kappa=4.0, mode="baseline", steps=500):
    # 2-D binary classification; learning rate itself is EMA-filtered.
    rng = np.random.default_rng(seed)
    x0 = rng.normal([-0.7, -0.7], .8, (96, 2))
    x1 = rng.normal([0.7, 0.7], .8, (96, 2))
    X = np.vstack([x0, x1]); y = np.r_[np.zeros(96), np.ones(96)]
    w = np.zeros(2); b = 0.0; lr_eff = .01; lr0, lr1 = .01, .08
    losses = []
    dt = 1.0
    for step in range(steps):
        frac = min(step / 150, 1.0)
        desired = lr0 + (lr1 - lr0) * frac
        rate = (lr1 - lr0) / 150 if step < 150 else 0.0
        if mode == "baseline": command = desired
        elif mode == "minus": command = max(.0001, desired - rate / kappa)
        elif mode == "plus": command = desired + rate / kappa
        alpha = min(0.9, 0.1 * kappa)
        lr_eff += alpha * (command - lr_eff)
        z = X @ w + b
        p = 1 / (1 + np.exp(-np.clip(z, -30, 30)))
        loss = -np.mean(y*np.log(p+1e-8) + (1-y)*np.log(1-p+1e-8))
        losses.append(float(loss))
        gw = X.T @ (p-y) / len(y); gb = float(np.mean(p-y))
        w -= lr_eff * gw; b -= lr_eff * gb
        if not np.isfinite(loss) or np.linalg.norm(w) > 1e6:
            return {"final_loss": float("inf"), "best_loss": float(np.min(losses)), "exploded": True}
    return {"final_loss": losses[-1], "best_loss": float(np.min(losses)), "exploded": False}


def main():
    random.seed(7); np.random.seed(7)
    math = math_verification()
    training = {m: train_toy(mode=m, steps=300) for m in ["baseline", "minus", "plus"]}
    result = {"math_verification": math, "training": training}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()