import json
import numpy as np

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


def make_problem(gamma):
    # H = I + gamma*(ones-I): SPD for gamma < 1, while Jacobi rho=2*gamma.
    E = gamma * (np.ones((3, 3)) - np.eye(3))
    H = np.eye(3) + E
    r = np.array([1.0, -0.7, 0.45])
    return H, E, r


def spectral_radius(A):
    return float(np.max(np.abs(np.linalg.eigvals(A))))


def run_jacobi(H, E, r, steps, alpha=1.0):
    z = np.zeros(3)
    residuals = []
    Dinv = np.diag(1.0 / np.diag(H))
    M = (1-alpha)*np.eye(3) - alpha*Dinv@E
    for _ in range(steps):
        z = (1-alpha)*z + alpha*Dinv@(r-E@z)
        residuals.append(float(np.linalg.norm(H@z-r)))
    return z, np.array(residuals), spectral_radius(M)


def run_gs(H, r, steps):
    z = np.zeros(3)
    residuals = []
    for _ in range(steps):
        for i in range(3):
            z[i] = (r[i] - (H[i] @ z - H[i, i]*z[i])) / H[i, i]
        residuals.append(float(np.linalg.norm(H@z-r)))
    return z, np.array(residuals)


def run_gd(H, r, steps, lr=0.25):
    z = np.zeros(3)
    residuals = []
    for _ in range(steps):
        z -= lr*(H@z-r)
        residuals.append(float(np.linalg.norm(H@z-r)))
    return z, np.array(residuals)


def run_adam(H, r, steps, lr=0.15):
    z = np.zeros(3); m = np.zeros(3); v = np.zeros(3)
    residuals = []
    for t in range(1, steps+1):
        g = H@z-r
        m = .9*m + .1*g; v = .999*v + .001*g*g
        z -= lr*(m/(1-.9**t))/(np.sqrt(v/(1-.999**t))+1e-8)
        residuals.append(float(np.linalg.norm(H@z-r)))
    return z, np.array(residuals)


def main():
    gammas = [0.10, 0.30, 0.49, 0.51, 0.70, 0.90]
    rows = []
    for g in gammas:
        H,E,r = make_problem(g)
        rho = spectral_radius(E)
        _, res, rho_measured = run_jacobi(H,E,r,40)
        # empirical asymptotic ratio, avoiding initial transient
        ratio = float(np.median(res[-8:-1]/res[-9:-2]))
        _, damp_res, damp_rho = run_jacobi(H,E,r,200,alpha=.5)
        _, gs_res = run_gs(H,r,40)
        _, gd_res = run_gd(H,r,40,lr=.25)
        _, adam_res = run_adam(H,r,40)
        if rho < 0.8:
            policy, _, policy_res, policy_rho = 'parallel', *run_jacobi(H,E,r,40,alpha=1.0)
        elif rho < 1.0:
            policy, _, policy_res, policy_rho = 'damped', *run_jacobi(H,E,r,40,alpha=.5)
        else:
            policy, _, policy_res = 'sequential', *run_gs(H,r,40)
            policy_rho = None
        rows.append({
            'gamma': g, 'predicted_rho_2gamma': 2*g, 'policy': policy, 'policy_residual_40': float(policy_res[-1]), 'measured_rho': rho,
            'jacobi_matrix_rho': rho_measured, 'jacobi_ratio_last': ratio,
            'jacobi_residual_40': float(res[-1]), 'damped_alpha_.5_rho': damp_rho,
            'damped_residual_200': float(damp_res[-1]), 'gs_residual_40': float(gs_res[-1]),
            'gd_residual_40': float(gd_res[-1]), 'adam_residual_40': float(adam_res[-1]),
            'spd_min_eigenvalue': float(np.min(np.linalg.eigvalsh(H)))
        })

    boundary = []
    for g in np.arange(0.45, 0.551, 0.01):
        H,E,r = make_problem(float(g))
        _, rr, mr = run_jacobi(H,E,r,100)
        boundary.append({'gamma': float(g), 'predicted_rho': float(2*g),
                         'matrix_rho': mr, 'residual_growth_100': float(rr[-1]/max(rr[0],1e-300))})

    # Direct numerical checks of the three claims.
    stable_below = [x for x in rows if x['gamma'] < .5 and x['jacobi_residual_40'] < 1e-6]
    unstable_above = [x for x in rows if x['gamma'] > .5 and x['jacobi_residual_40'] > 1e2]
    ratio_check = [x for x in rows if .2 <= x['gamma'] <= .49 and abs(x['jacobi_ratio_last']-x['predicted_rho_2gamma']) < .03]
    damp_check = [x for x in rows if x['gamma'] == .9 and x['damped_alpha_.5_rho'] < 1 and x['damped_residual_200'] < 0.2]
    report = {
        'seed': SEED,
        'predictions': {
            'boundary': 'rho=2*gamma, transition at gamma=0.5',
            'contraction': 'stable Jacobi residual ratio approaches rho',
            'damping': 'alpha=0.5 stabilizes gamma=0.9 despite undamped rho=1.8'
        },
        'rows': rows, 'boundary_sweep': boundary,
        'checks': {
            'stable_below_count': len(stable_below), 'unstable_above_count': len(unstable_above),
            'ratio_match_count': len(ratio_check), 'damping_check': bool(damp_check)
        }
    }
    with open('results.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()