import json
from pathlib import Path
import numpy as np
from scipy.optimize import least_squares

SEED = 1411
TRUE_A = 0.70
TRUE_B = 1.20
TRUE_TAU = 0.80
NOISE = 0.01


def disturbance(u, alpha, tau):
    return alpha * np.tanh(u / tau)


def make_data(alpha, n, u_limit, seed, tau=TRUE_TAU):
    rng = np.random.default_rng(seed)
    x = rng.uniform(-1.0, 1.0, n)
    u = rng.uniform(-u_limit, u_limit, n)
    y = TRUE_A * x + TRUE_B * (u + disturbance(u, alpha, tau))
    y += rng.normal(0.0, NOISE, n)
    return x, u, y


def fit_baseline(x, u, y):
    X = np.column_stack([x, u, np.ones_like(x)])
    p, *_ = np.linalg.lstsq(X, y, rcond=None)
    return p


def fit_joint(x, u, y, tau_hint=TRUE_TAU):
    def residual(p):
        a, b, alpha, tau, c = p
        pred = a * x + b * (u + disturbance(u, alpha, tau)) + c
        return (pred - y) / NOISE
    p0 = np.array([0.6, 1.1, 0.15, tau_hint, 0.0])
    lo = np.array([-3., -3., -2., 0.15, -1.])
    hi = np.array([3., 3., 2., 3.0, 1.])
    result = least_squares(residual, p0, bounds=(lo, hi), max_nfev=1000)
    return result.x


def mse(y, pred):
    return float(np.mean((y - pred) ** 2))


def main():
    # Prediction 1: max_u |d'(u)| = |alpha|/tau, attained at u=0.
    grid = np.linspace(-8, 8, 20001)
    alpha, tau = 0.37, 0.63
    analytic = alpha / tau / np.cosh(grid / tau) ** 2
    numeric = np.gradient(disturbance(grid, alpha, tau), grid)
    derivative_check = {
        "predicted_max": float(abs(alpha) / tau),
        "observed_max": float(np.max(np.abs(numeric))),
        "relative_error": float(abs(np.max(np.abs(numeric)) - abs(alpha) / tau) / (abs(alpha) / tau)),
        "bound_location": float(grid[np.argmax(np.abs(numeric))]),
        "analytic_numeric_max_gap": float(np.max(np.abs(analytic - numeric))),
    }

    # Prediction 2: for small alpha, omitted-variable bias in the linear action
    # coefficient is approximately B*alpha*E[u*tanh(u/tau)]/E[u^2], hence linear in alpha.
    alphas = np.array([0.0, 0.05, 0.10, 0.20, 0.40, 0.80])
    rows = []
    for i, a in enumerate(alphas):
        x, u, y = make_data(a, 5000, 1.5, SEED + i)
        pb = fit_baseline(x, u, y)
        pjoint = fit_joint(x, u, y)
        # Test on a fresh broad command range to expose saturation extrapolation.
        xt, ut, yt = make_data(a, 3000, 2.2, SEED + 100 + i)
        pred_b = pb[0] * xt + pb[1] * ut + pb[2]
        pred_j = pjoint[0] * xt + pjoint[1] * (ut + disturbance(ut, pjoint[2], pjoint[3])) + pjoint[4]
        ratio = np.mean(u * np.tanh(u / TRUE_TAU)) / np.mean(u * u)
        predicted_bias = TRUE_B * a * ratio
        rows.append({
            "alpha": float(a),
            "predicted_linear_bias": float(predicted_bias),
            "observed_baseline_b_bias": float(pb[1] - TRUE_B),
            "joint_train_mse": mse(y, pjoint[0] * x + pjoint[1] * (u + disturbance(u, pjoint[2], pjoint[3])) + pjoint[4]),
            "baseline_cross_regime_mse": mse(yt, pred_b),
            "joint_cross_regime_mse": mse(yt, pred_j),
            "fitted_alpha": float(pjoint[2]),
            "fitted_tau": float(pjoint[3]),
        })

    # Prediction 3: no actuator disturbance (alpha=0) removes the mechanism;
    # the two models should have comparable cross-regime error.
    zero_rows = []
    for i in range(8):
        x, u, y = make_data(0.0, 2000, 1.5, SEED + 300 + i)
        pb = fit_baseline(x, u, y)
        pj = fit_joint(x, u, y)
        xt, ut, yt = make_data(0.0, 2000, 2.2, SEED + 400 + i)
        zero_rows.append({
            "baseline_mse": mse(yt, pb[0] * xt + pb[1] * ut + pb[2]),
            "joint_mse": mse(yt, pj[0] * xt + pj[1] * (ut + disturbance(ut, pj[2], pj[3])) + pj[4]),
            "fitted_alpha": float(pj[2]),
        })

    # Fit slopes through the origin for the small-alpha predicted/observed bias.
    small = rows[:4]
    alpha_small = np.array([r["alpha"] for r in small])
    observed_bias_small = np.array([r["observed_baseline_b_bias"] for r in small])
    predicted_bias_small = np.array([r["predicted_linear_bias"] for r in small])
    observed_slope = float(np.dot(alpha_small, observed_bias_small) / np.dot(alpha_small, alpha_small))
    predicted_slope = float(np.dot(alpha_small, predicted_bias_small) / np.dot(alpha_small, alpha_small))
    zero_base = float(np.mean([r["baseline_mse"] for r in zero_rows]))
    zero_joint = float(np.mean([r["joint_mse"] for r in zero_rows]))
    output = {
        "seed": SEED,
        "derivative_bound_check": derivative_check,
        "bias_scaling_prediction": {
            "predicted_slope": predicted_slope,
            "observed_slope": observed_slope,
            "slope_ratio_observed_over_predicted": observed_slope / predicted_slope,
            "rows": rows,
        },
        "zero_disturbance_prediction": {
            "mean_baseline_mse": zero_base,
            "mean_joint_mse": zero_joint,
            "joint_over_baseline": zero_joint / zero_base,
            "rows": zero_rows,
        },
        "interpretation": "Joint model should recover saturation and improve held-out command-regime error when alpha is nonzero; at alpha=0 no systematic benefit is expected.",
    }
    Path("results.json").write_text(json.dumps(output, indent=2))
    print(json.dumps(output, indent=2))


if __name__ == "__main__":
    main()