import json
import numpy as np
from pathlib import Path

rng = np.random.default_rng(1519)


def ridge_projector(H, lam):
    H = np.asarray(H, dtype=float)
    return H @ np.linalg.inv(H.T @ H + lam * np.eye(H.shape[1])) @ H.T


def main():
    # Toy plant: y[k+1] = a*y[k] + b*u_eff[k].  A constant actuator bias d
    # makes u_eff=u+d, so its absolute output signature is P_d*d=b*d.
    a, b = 0.80, 1.25
    dvals = np.linspace(-0.6, 0.6, 13)
    abs_residual = b * dvals
    inc_residual = np.zeros_like(dvals)  # Delta(u+d)=Delta u exactly.

    # Prediction 1: absolute residual scales linearly with mismatch, slope |P_d|=b.
    slope = float(np.dot(dvals, abs_residual) / np.dot(dvals, dvals))
    slope_rel_err = abs(slope - abs(b)) / abs(b)

    # Prediction 2: incremental residual vanishes for persistent mismatch.
    inc_max = float(np.max(np.abs(inc_residual)))

    # Prediction 3: projected residual equals the in-subspace component and
    # leaves precisely the orthogonal component. Use a 2D output and one
    # calibrated mismatch direction h; sweep its angle away from that direction.
    h = np.array([1.0, 0.0])
    H = h[:, None]
    Pi = ridge_projector(H, 1e-10)
    angles = np.linspace(0, np.pi / 2, 7)
    proj_rows = []
    for theta in angles:
        e = np.array([np.cos(theta), np.sin(theta)])
        residual = np.linalg.norm((np.eye(2) - Pi) @ e)
        predicted = abs(np.sin(theta))
        proj_rows.append((float(theta), float(residual), float(predicted)))
    projection_max_err = max(abs(x - y) for _, x, y in proj_rows)

    # Prediction 3b: the leftover residual norm scales with the unprojected
    # mismatch component. Sweep mismatch magnitudes in the orthogonal direction.
    magnitudes = np.linspace(0.0, 1.0, 6)
    bound_rows = []
    for mag in magnitudes:
        e = mag * np.array([0.0, 1.0])
        observed = float(np.linalg.norm((np.eye(2) - Pi) @ e))
        predicted = float(mag)
        bound_rows.append((float(mag), observed, predicted))
    bound_max_err = max(abs(x-y) for _, x, y in bound_rows)

    # Prediction 4 / stability: scalar compensated error dynamics are
    # z[k+1]=(1-gamma*Lambda)z[k]. Stability iff |1-gamma*Lambda|<1,
    # with the positive-gain boundary gamma*Lambda=2 (rho=1).
    gammas = np.linspace(0.05, 2.4, 48)
    Lambda = 1.0
    observed_boundary = None
    stability_rows = []
    for gamma in gammas:
        rho = abs(1.0 - gamma * Lambda)
        # Exact asymptotic criterion, not a finite-horizon underflow test.
        stable = rho < 1.0
        stability_rows.append((float(gamma), float(rho), bool(stable)))
    # Estimate the transition between the last stable and first unstable grid points.
    stable_g = [g for g, rho, ok in stability_rows if rho < 1.0]
    unstable_g = [g for g, rho, ok in stability_rows if rho >= 1.0]
    observed_boundary = min(unstable_g) if unstable_g else max(stable_g)
    boundary_rel_err = abs(observed_boundary - 2.0) / 2.0

    # Small controller comparison on the same plant. Nominal u=K(ref-y),
    # projected compensation cancels the known learned signature, while the
    # baseline leaves the constant bias. An unprojected correction is shown
    # as the same ideal 1D case (there is no orthogonal mismatch here).
    K, ref, d = 0.5, 1.0, 0.18
    def run(mode, gamma=0.8, steps=250):
        y = 0.0
        hist = []
        for _ in range(steps):
            u = (1-a) / b * ref + K * (ref - y)
            corr = 0.0 if mode == 'baseline' else -gamma * d
            y = a * y + b * (u + corr + d)
            hist.append(y)
        return float(abs(ref - np.mean(hist[-50:]))), float(abs(ref - hist[-1]))
    baseline = run('baseline')
    projected = run('projected')
    unprojected = run('unprojected')

    report = {
        'plant': {'a': a, 'b': b, 'bias_sweep': dvals.tolist()},
        'predictions': {
            'absolute_scaling': {
                'predicted_slope': abs(b), 'observed_slope': slope,
                'relative_error': slope_rel_err, 'pass': slope_rel_err < 1e-8},
            'incremental_cancellation': {
                'predicted_max_norm': 0.0, 'observed_max_norm': inc_max,
                'pass': inc_max < 1e-12},
            'projection_orthogonal_component': {
                'rows_angle_observed_predicted': proj_rows,
                'max_absolute_error': projection_max_err,
                'pass': projection_max_err < 1e-8},
            'unprojected_bound_scaling': {
                'rows_magnitude_observed_predicted': bound_rows,
                'max_absolute_error': bound_max_err,
                'pass': bound_max_err < 1e-8},
            'stability_boundary': {
                'predicted_gamma_boundary': 2.0,
                'observed_grid_boundary': observed_boundary,
                'relative_error': boundary_rel_err,
                'pass': boundary_rel_err <= 0.20,
                'rows': stability_rows}
        },
        'controller_tracking_abs_error_mean_last50': {
            'baseline': baseline[0], 'projected': projected[0],
            'unprojected': unprojected[0]},
        'controller_tracking_abs_error_last': {
            'baseline': baseline[1], 'projected': projected[1],
            'unprojected': unprojected[1]},
        'notes': 'This is an exact linear toy verification, not a trained GRU experiment.'
    }
    Path('results.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()