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

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


def roots(lam, q, alpha):
    # (r-lambda)^2 + q alpha (r-1) = 0
    return np.roots([1.0, q * alpha - 2.0 * lam, lam * lam - q * alpha])


def pole_radius(lams, qs, alpha):
    return float(max(abs(r) for lam in lams for q in qs for r in roots(lam, q, alpha)))


def predicted_boundary(lam, q):
    # Jury boundary for 0 <= lambda < 1: q alpha = (1+lambda)^2/2.
    return (1.0 + lam) ** 2 / (2.0 * q)


def make_symmetric_W(eigs):
    n = len(eigs) + 1
    # Deterministic orthonormal basis with first vector equal to consensus.
    A = rng.normal(size=(n, n))
    A[:, 0] = 1.0
    Q, _ = np.linalg.qr(A)
    if np.sum(Q[:, 0]) < 0:
        Q[:, 0] *= -1
    vals = np.array([1.0] + list(eigs), dtype=float)
    W = Q @ np.diag(vals) @ Q.T
    # Numerical cleanup preserves symmetry and stochasticity closely.
    W = (W + W.T) / 2.0
    return W


def diging(W, qs_worker, centers, alpha, rounds=250):
    n, d = centers.shape
    x = np.zeros((n, d))
    g = qs_worker[:, None] * (x - centers)
    y = g.copy()
    losses = []
    tracker = []
    for _ in range(rounds):
        x = W @ x - alpha * y
        g_new = qs_worker[:, None] * (x - centers)
        y = W @ y + g_new - g
        g = g_new
        local_loss = 0.5 * qs_worker[:, None] * (x - centers) ** 2
        losses.append(float(np.mean(np.sum(local_loss, axis=1))))
        tracker.append(float(np.linalg.norm(y - np.mean(y, axis=0, keepdims=True))))
    return np.asarray(losses), np.asarray(tracker), x


def main():
    # Core numerical sanity check: formula roots match polynomial evaluation.
    checks = []
    for lam, q, a in [(0.2, 1.0, .1), (.75, 4.0, .03), (-.3, 2.0, .2)]:
        rr = roots(lam, q, a)
        residual = max(abs((r-lam)**2 + q*a*(r-1)) for r in rr)
        checks.append(float(residual))

    # Prediction 1: stability transition at alpha_c=(1+lambda)^2/(2q).
    lam0, q0 = .6, 3.0
    ac = predicted_boundary(lam0, q0)
    ratios = np.array([.98, .995, 1.005, 1.02])
    boundary_rows = []
    for ratio in ratios:
        rho = pole_radius([lam0], [q0], ac * ratio)
        boundary_rows.append({"ratio": float(ratio), "alpha": float(ac*ratio),
                              "rho": rho, "stable_observed": bool(rho < 1.0)})

    # Prediction 2: boundary scales inversely with curvature q.
    lam1 = .6
    qvals = np.array([1., 2., 4., 8.])
    measured = []
    for q in qvals:
        grid = np.linspace(.0001, 1.2 * predicted_boundary(lam1, q), 3000)
        rhos = np.array([pole_radius([lam1], [q], a) for a in grid])
        # first grid point at/above one is a numerical estimate of the transition
        idx = np.flatnonzero(rhos >= 1.0)
        observed = float(grid[idx[0]]) if len(idx) else float("nan")
        measured.append({"q": float(q), "predicted_alpha_c": predicted_boundary(lam1, q),
                         "observed_alpha_c": observed,
                         "relative_error": abs(observed-predicted_boundary(lam1,q))/predicted_boundary(lam1,q)})

    # Prediction 3: minimax tuning picks a compromise over the lambda/q rectangle.
    lams = np.array([.2, .6, .85])
    qs = np.linspace(1., 5., 17)
    candidates = np.linspace(.0005, .30, 1200)
    candidate_rho = np.array([pole_radius(lams, qs, a) for a in candidates])
    best_i = int(np.argmin(candidate_rho))
    best_a = float(candidates[best_i])
    best_rho = float(candidate_rho[best_i])
    # A deliberately generic conservative choice, and its predicted contraction.
    fixed_a = .05
    fixed_rho = pole_radius(lams, qs, fixed_a)

    # Small decentralized quadratic experiment, same setup and rounds.
    W = make_symmetric_W(lams)
    worker_q = np.array([1., 2., 3., 5.])
    centers = rng.normal(scale=1.0, size=(4, 2))
    tuned_loss, tuned_tracker, tuned_x = diging(W, worker_q, centers, best_a)
    fixed_loss, fixed_tracker, fixed_x = diging(W, worker_q, centers, fixed_a)
    # Centralized optimum for reporting a meaningful parameter error.
    x_star = np.sum(worker_q[:, None] * centers, axis=0) / np.sum(worker_q)
    tuned_err = float(np.linalg.norm(np.mean(tuned_x, axis=0) - x_star))
    fixed_err = float(np.linalg.norm(np.mean(fixed_x, axis=0) - x_star))

    report = {
        "seed": SEED,
        "root_residual_max": max(checks),
        "prediction_1_boundary": {"lambda": lam0, "q": q0, "predicted_alpha_c": ac,
                                   "rows": boundary_rows},
        "prediction_2_inverse_curvature": measured,
        "prediction_3_minimax": {"lambdas": lams.tolist(), "q_interval": [1., 5.],
                                  "tuned_alpha": best_a, "tuned_predicted_rho": best_rho,
                                  "fixed_alpha": fixed_a, "fixed_predicted_rho": fixed_rho,
                                  "predicted_improvement": fixed_rho-best_rho},
        "toy_diging": {"rounds": 250, "tuned_alpha": best_a, "fixed_alpha": fixed_a,
                        "initial_loss": float(fixed_loss[0]),
                        "final_loss_fixed": float(fixed_loss[-1]),
                        "final_loss_tuned": float(tuned_loss[-1]),
                        "final_tracker_fixed": float(fixed_tracker[-1]),
                        "final_tracker_tuned": float(tuned_tracker[-1]),
                        "final_mean_error_fixed": fixed_err,
                        "final_mean_error_tuned": tuned_err},
    }
    Path("results.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()