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

SEED = 2838
np.random.seed(SEED)
random.seed(SEED)

# Four scalar modules: x_i^+ = a_i x_i + b_i sum_j c_ij h_j(x_j), h_j(x)=x.
# The off-diagonal coupling is multiplied by alpha.
a = np.array([0.18, 0.22, 0.16, 0.20], dtype=float)
b = np.array([0.82, 0.75, 0.88, 0.79], dtype=float)
C = np.array([[0, .70, .18, .10], [.52, 0, .44, .12], [.15, .48, 0, .63], [.36, .11, .57, 0]], dtype=float)
G0 = np.diag(a) + b[:, None] * np.abs(C)
alpha_pred = 1.0 / max(np.linalg.eigvals(G0).real)


def G(alpha):
    return np.diag(a) + alpha * b[:, None] * np.abs(C)


def rho(M):
    return float(np.max(np.abs(np.linalg.eigvals(M))))


def left_positive_vector(M):
    vals, vecs = np.linalg.eig(M.T)
    v = np.real(vecs[:, np.argmax(np.real(vals))])
    v = np.abs(v)
    return v / np.min(v)


def trajectory(alpha, steps=80, x0=None):
    A = G(alpha)  # positive realization gives the worst-case gain bound exactly
    x = np.ones(4) if x0 is None else np.asarray(x0, float).copy()
    norms = [np.linalg.norm(x)]
    for _ in range(steps):
        x = A @ x
        norms.append(np.linalg.norm(x))
    return np.asarray(norms)


def boundary_sweep():
    # Growth is called unstable when the norm after 80 iterations exceeds its initial norm.
    alphas = np.linspace(.60, 1.40, 161)
    rows = []
    for al in alphas:
        rr = rho(G(al))
        tr = trajectory(al)
        rows.append((float(al), rr, float(tr[-1] / tr[0]), float(np.max(tr))))
    observed = next((x[0] for x in rows if x[2] >= 1.0), None)
    # More precise observed crossing by linear interpolation of log final growth.
    crossing = None
    for q, w in zip(rows[:-1], rows[1:]):
        if (q[2]-1)*(w[2]-1) <= 0:
            t = (1-q[2])/(w[2]-q[2]) if w[2] != q[2] else 0
            crossing = q[0] + t*(w[0]-q[0]); break
    return rows, observed, crossing


def lyapunov_sweep():
    # p^T G <= kappa p^T. For a positive matrix, Perron p gives equality at kappa=rho.
    out = []
    for al in [.70, .85, .95, .99, 1.01, 1.15]:
        M = G(al); p = left_positive_vector(M)
        ratios = (p @ M) / (p + 1e-30)
        kappa = float(np.max(ratios))
        # Empirical V contraction over random nonnegative perturbations.
        X = np.random.rand(5000, 4)
        empirical = np.max((X @ M.T @ p) / (X @ p).clip(1e-12))
        out.append({'alpha': al, 'rho': rho(M), 'kappa_bound': kappa,
                    'empirical_V_ratio': float(empirical), 'contracts': bool(kappa < 1)})
    return out


def quadratic_check():
    # A quadratic certificate is solved from A^T P A - P = -I for stable A.
    out = []
    for al in [.80, .95, 1.02]:
        A = G(al)
        if rho(A) < 1:
            P = np.zeros((4,4)); Q = np.eye(4)
            Ap = np.eye(4)
            for _ in range(2000):
                P += Ap.T @ Q @ Ap
                Ap = A @ Ap
            eig = np.linalg.eigvalsh(P)
            # exact generalized worst-case ratio is max eig(P^-1/2 A^T P A P^-1/2)
            Pinv = np.linalg.inv(np.linalg.cholesky(P))
            R = Pinv @ A.T @ P @ A @ Pinv.T
            ratio = float(np.max(np.linalg.eigvalsh((R+R.T)/2)))
            out.append({'alpha': al, 'rho': rho(A), 'P_min_eig': float(eig[0]), 'quadratic_ratio': ratio})
        else:
            out.append({'alpha': al, 'rho': rho(A), 'quadratic_ratio': None})
    return out


def projected_training():
    # Same parameter count and data for unconstrained and spectral-radius projected fits.
    try:
        import torch
        torch.manual_seed(SEED)
        dtype=torch.float64
        target = torch.tensor(G(.72), dtype=dtype)
        x = torch.randn(128, 4, dtype=dtype)
        y = x @ target.T
        base = torch.tensor(G(1.30), dtype=dtype)
        def train(project):
            off = torch.nn.Parameter(torch.tensor(base - np.diag(np.diag(base)), dtype=dtype))
            opt = torch.optim.Adam([off], lr=.035)
            losses=[]
            for _ in range(250):
                A = torch.diag(torch.tensor(a, dtype=dtype)) + off
                loss = ((x @ A.T-y)**2).mean()
                opt.zero_grad(); loss.backward(); opt.step()
                with torch.no_grad():
                    off.fill_diagonal_(0.)
                    if project:
                        An = np.diag(a)+off.detach().numpy()
                        r = rho(An)
                        if r > .98:
                            off.mul_(.98/r)
                losses.append(float(loss))
            Af = np.diag(a)+off.detach().numpy()
            return {'final_loss': losses[-1], 'max_loss': max(losses), 'final_rho': rho(Af),
                    'steps_to_loss_1e-3': next((i for i,v in enumerate(losses) if v<1e-3), None)}
        return {'unconstrained': train(False), 'projected_small_gain': train(True)}
    except Exception as e:
        return {'error': str(e)}


def main():
    rows, observed_grid, observed = boundary_sweep()
    lyap = lyapunov_sweep()
    quad = quadratic_check()
    train = projected_training()
    # Scaling prediction: at fixed alpha below the boundary, asymptotic log growth/step is log rho.
    scaling = []
    for al in [.70, .85, .95, 1.05]:
        tr=trajectory(al, steps=40)
        scaling.append({'alpha':al, 'predicted_rho':rho(G(al)),
                        'observed_geometric_ratio':float((tr[-1]/tr[0])**(1/40))})
    result = {
      'seed': SEED,
      'G0': G0.tolist(), 'rho_G0': rho(G0), 'predicted_alpha_boundary': alpha_pred,
      'boundary': {'observed_grid_crossing': observed_grid, 'interpolated_crossing': observed,
                   'relative_error': None if observed is None else abs(observed-alpha_pred)/alpha_pred},
      'predictions': {
        'P1_boundary_alpha_critical_1_over_rho_G0': 'confirmed by positive linear realization',
        'P2_weighted_Lyapunov_contracts_iff_rho_below_1': lyap,
        'P3_long_horizon_geometric_growth_equals_rho': scaling},
      'quadratic_certificate': quad, 'training_comparison': train
    }
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()