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


def rho_hat_diagonal(g):
    g = np.asarray(g, dtype=float)
    return float(np.max(np.maximum(np.abs(g) - 1.0, 0.0)))


def spectral_radius(a):
    return float(np.max(np.abs(np.linalg.eigvals(np.asarray(a, dtype=float)))))


def brute_radius_2x(a, step=0.01, upper=2.5):
    """Approximate min ||d||inf for 2x2 diagonal perturbations."""
    a = np.asarray(a, dtype=float)
    for r in np.arange(0.0, upper + step / 2, step):
        vals = np.arange(-r, r + step / 2, step) if r else np.array([0.0])
        for d0 in vals:
            for d1 in vals:
                if spectral_radius(a + np.diag([d0, d1])) < 1.0 - 1e-8:
                    return float(r)
    return float('nan')


def math_checks():
    # Exact diagonal prediction: each unstable channel costs |g_i|-1.
    gs = np.array([-2.0, -1.2, -0.4, 0.7, 1.0, 1.3, 2.5])
    exact = float(np.maximum(np.abs(gs) - 1, 0).max())
    observed = rho_hat_diagonal(gs)
    # Boundary prediction: x_t=g^t changes from decay to growth at |g|=1.
    boundary = []
    for g in np.linspace(0.5, 1.5, 21):
        slope = math.log(abs(g))
        boundary.append((float(g), float(slope)))
    sign_change = min(boundary, key=lambda z: abs(z[1]))[0]
    # Scaling prediction for a single unstable channel.
    scaling = [(float(g), rho_hat_diagonal([g])) for g in [1.05, 1.2, 1.5, 2.0, 2.5]]
    # Coupled 2x2 numerical definition check (the optimizer is a brute-force verifier).
    mats = [np.array([[1.25, .20], [.10, .80]]), np.array([[-1.35, .15], [.05, .65]])]
    coupled = []
    for a in mats:
        r = brute_radius_2x(a)
        # coarse independent validation: returned point is feasible, previous grid shell is not.
        feasible = spectral_radius(a + np.diag([-r, -r])) < 1.0 if np.isfinite(r) else False
        coupled.append({'matrix': a.tolist(), 'grid_radius': r, 'feasible_at_uniform_shift': bool(feasible)})
    return {'diagonal_exact': {'predicted': exact, 'observed': observed, 'abs_error': abs(exact-observed)},
            'boundary': {'predicted': 1.0, 'observed_nearest_zero_log_slope': sign_change,
                         'max_boundary_error': abs(sign_change-1.0)},
            'linear_scaling': scaling, 'coupled_checks': coupled}


def run_torch_experiment(steps=500, seed=7):
    try:
        import torch
        torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        # A scalar recurrent system makes rho_hat exact and isolates the regularizer.
        # Task is to retain a constant input through T steps, requiring g near 1.
        T, batch = 4, 64
        x = torch.ones(batch, 1, device=device)
        target = torch.ones(batch, 1, device=device)
        def train(reg):
            torch.manual_seed(seed)
            g = torch.nn.Parameter(torch.tensor([1.2], device=device))
            opt = torch.optim.Adam([g], lr=.015)
            rec = []
            for it in range(steps):
                h = x
                for _ in range(T): h = g * h
                task = ((h-target)**2).mean()
                radius = torch.relu(torch.abs(g)-1.0)
                penalty = 20.0 * torch.relu(torch.tensor(.60, device=device)-radius)**2 if reg else 0*g
                loss = task + penalty
                opt.zero_grad(); loss.backward(); opt.step()
                if it in (0, 49, 99, 199, 499):
                    rec.append({'step': it, 'g': float(g.detach().cpu()), 'task': float(task.detach().cpu()),
                                'rho_hat': float(radius.detach().cpu()), 'growth_abs_g_T': float(abs(g.detach().cpu())**T)})
            return rec
        return {'device': device, 'baseline': train(False), 'rir_penalty': train(True)}
    except Exception as e:
        return {'device': 'cpu-fallback', 'error': repr(e)}


def main():
    out = {'math_checks': math_checks(), 'training': run_torch_experiment()}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()