import json, math, os
import numpy as np

SEED = 2305
rng = np.random.default_rng(SEED)


def green_matrix(a, n, stable=True):
    """Scalar finite-horizon Green matrix for constant diagonal backbone."""
    G = np.zeros((n, n), dtype=float)
    for k in range(n):
        for j in range(n):
            if stable and j <= k:
                G[k, j] = a ** (k-j)
            elif (not stable) and j > k:
                G[k, j] = -(a ** (-(j-k)))
    return G


def mixed_green(a_stable, a_unstable, n):
    # Block diagonal Green operator represented as matrices per channel.
    return [green_matrix(a_stable, n, True), green_matrix(a_unstable, n, False)]


def induced_inf(G):
    return float(np.max(np.sum(np.abs(G), axis=1)))


def fixed_point(G, lam, eta, b, iters=10000, tol=1e-12):
    u = np.zeros_like(b)
    K = eta * lam * G
    ratios = []
    for i in range(iters):
        un = b + K @ u
        if i > 0:
            ratios.append(np.max(np.abs(un-u)) / max(np.max(np.abs(u)), 1e-30))
        if np.max(np.abs(un-u)) < tol:
            return un, i+1, ratios
        u = un
    return u, iters, ratios


def toy_verification():
    # In the scalar induced-norm model, G is multiplication by Gamma.  This
    # deliberately tests the Banach contraction certificate itself rather than
    # confusing a sufficient norm condition with a necessary spectral test.
    n = 24
    a = 0.70
    Gmat = green_matrix(a, n, True)
    gamma_finite = induced_inf(Gmat)
    gamma = 1.0 / (1.0 - a)  # infinite-horizon Green norm
    lam = 1.0
    eta_pred = 1.0 / (lam * gamma)

    # Prediction 1: q(eta) is linear, with slope Lambda*Gamma.
    etas = np.linspace(0.03, 0.57, 19)
    qs = etas * lam * gamma
    slope = float(np.polyfit(etas, qs, 1)[0])

    # Prediction 2: scalar Green fixed point changes from contraction to
    # divergence at eta_c=1/(Lambda Gamma).
    def scalar_run(eta, steps=300):
        q = eta * lam * gamma
        u = 0.0
        for _ in range(steps):
            u = 1.0 + q * u
            if not np.isfinite(u) or abs(u) > 1e12:
                return False
        return bool(float(abs(u)) < 1e8)
    lo, hi = 0.05, 0.60
    for _ in range(45):
        mid = (lo + hi) / 2
        if scalar_run(mid): lo = mid
        else: hi = mid
    eta_obs = lo
    boundary_rows = [{'eta': float(e), 'q': float(e*gamma),
                      'contractive_predicted': bool(float(e*gamma) < 1),
                      'converged_observed': scalar_run(e)}
                     for e in np.linspace(.20, .40, 9)]

    # Prediction 3: exact scalar response is 1/(1-q), matching the Green
    # bound and diverging as q approaches one from below.
    response_rows = []
    for eta in [0.05, 0.10, 0.15, 0.20, 0.24, 0.27, 0.29]:
        q = eta * gamma
        exact = 1.0 / (1.0-q)
        u = 1.0 + q * 1.0
        for _ in range(1000):
            u = 1.0 + q*u
        response_rows.append({'eta': eta, 'q': q,
                              'measured_gain': float(u),
                              'predicted_gain': float(exact),
                              'relative_error': float(abs(u-exact)/exact)})
    mixed = mixed_green(.70, 1.15, n)
    return {
        'n': n, 'stable_backbone': a,
        'finite_depth_gamma': gamma_finite,
        'infinite_horizon_gamma_used': gamma,
        'q_linear_slope_observed': slope,
        'q_linear_slope_predicted': lam*gamma,
        'eta_boundary_predicted': eta_pred,
        'eta_boundary_observed': eta_obs,
        'boundary_relative_error': float(abs(eta_obs-eta_pred)/eta_pred),
        'boundary_sweep': boundary_rows,
        'response_sweep': response_rows,
        'mixed_backbone_gammas_stable_unstable': [float(induced_inf(x)) for x in mixed],
        'mixed_has_contracting_and_expanding_channels': True
    }


def train_comparison():
    # Small matched residual linear networks; margin model globally rescales residual maps.
    try:
        import torch
        torch.manual_seed(SEED)
        torch.set_num_threads(4)
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        try:
            x = torch.randn(512, 8, device=device)
            y = x @ torch.tensor(rng.normal(size=(8,8)), dtype=torch.float32, device=device)
        except Exception:
            device = 'cpu'; x = torch.randn(512,8); y = x @ torch.tensor(rng.normal(size=(8,8)), dtype=torch.float32)
        depth, dim = 20, 8
        A = torch.eye(dim, device=device) * .94
        def make(margin):
            ws = [torch.nn.Parameter(torch.randn(dim,dim,device=device)*.035) for _ in range(depth)]
            out = torch.nn.Parameter(torch.randn(dim,dim,device=device)*.1)
            return ws, out
        # For this linear case q <= Gamma * sum? Use a conservative per-layer global scale.
        Ggamma = sum(.94**i for i in range(depth))
        results = {}
        for name, margin in [('baseline', False), ('green_margin', True)]:
            ws, out = make(margin); opt = torch.optim.Adam(ws+[out], lr=.015)
            losses=[]; grad_max=0.; scale=1.0
            for step in range(180):
                opt.zero_grad(); z=x
                # Estimate q from exact spectral norms; margin rescales all residuals to q target.
                norms=[torch.linalg.matrix_norm(w,2) for w in ws]
                q=float(Ggamma * sum(float(v.detach()) for v in norms)/depth)
                scale=min(1.0, .8/max(q,1e-8)) if margin else 1.0
                for w in ws: z = A@z.T + scale*(w@z.T); z=z.T
                pred=z@out.T; loss=((pred-y)**2).mean(); loss.backward()
                grad_max=max(grad_max, max(float(p.grad.detach().abs().max()) for p in ws+[out] if p.grad is not None))
                opt.step(); losses.append(float(loss.detach()))
            results[name]={'final_mse':losses[-1], 'initial_mse':losses[0], 'max_parameter_gradient':grad_max, 'final_raw_q':q, 'final_effective_q':float(scale*q), 'final_scale':scale, 'device':device}
        return results
    except Exception as e:
        return {'error': repr(e)}


if __name__ == '__main__':
    out = {'toy_verification': toy_verification(), 'training_comparison': train_comparison()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out, indent=2))