import json, math, os
import numpy as np

SEED = 1402
RNG = np.random.default_rng(SEED)


def geom_sensitivity(delta0, rho, K):
    return delta0 * sum(rho ** k for k in range(K))


def gaussian_proxy(delta, sigma, delta_fail=1e-5):
    return delta / sigma * math.sqrt(2.0 * math.log(1.25 / delta_fail))


def stability_sweep():
    # For x_{k+1}=(1-beta*lambda)x_k, contraction requires |1-beta*lambda|<1.
    lam = 1.0
    betas = [0.5, 1.5, 2.0, 2.1, 2.5]
    rows = []
    for beta in betas:
        rho = abs(1.0-beta*lam)
        x = 1.0
        for _ in range(40):
            x = (1.0-beta*lam)*x
        rows.append({"beta": beta, "predicted_rho": rho,
                     "observed_final_abs": abs(x),
                     "predicted_contractive": rho < 1.0,
                     "observed_decay": abs(x) < 1.0})
    return rows

def toy_predictions():
    # Prediction 1: a contractive recurrence has geometric sensitivity and a finite limit.
    delta0 = 1.0
    K = 100
    rhos = [0.2, 0.5, 0.8, 0.95, 1.05]
    rows = []
    for rho in rhos:
        vals = delta0 * rho ** np.arange(K)
        observed = float(vals.sum())
        predicted = float(delta0 * (1-rho**K)/(1-rho)) if rho != 1 else float(K)
        infinite = delta0/(1-rho) if rho < 1 else None
        rows.append({"rho":rho, "observed_sum":observed, "finite_geom_prediction":predicted,
                     "infinite_bound":infinite, "ratio_to_bound":(observed/infinite if infinite else None)})

    # Prediction 2: 95% of infinite sensitivity mass arrives by log(.05)/log(rho).
    k95_rows = []
    for rho in [0.5, 0.8, 0.9, 0.95]:
        pred = math.log(0.05)/math.log(rho)
        vals = rho ** np.arange(1000)
        cumulative = np.cumsum(vals)
        target = 0.95/(1-rho)
        observed = int(np.flatnonzero(cumulative >= target)[0] + 1)
        k95_rows.append({"rho":rho, "predicted_k95_steps":pred, "observed_k95_steps":observed,
                         "relative_error":abs(observed-pred)/pred})

    # Prediction 3: for fixed d, channel privacy proxy scales as alpha^-1/2.
    d = 2.0; sigma0 = 0.2; eps = 1e-6; p = 2.0; delta = 0.1
    alphas = np.array([0.01, 0.04, 0.16, 0.64])
    proxy = np.array([gaussian_proxy(delta, math.sqrt(sigma0**2+a*(d+eps)**p)) for a in alphas])
    scaling = proxy[0] / proxy
    expected = np.sqrt(alphas/alphas[0])
    scaling_rows = [{"alpha":float(a), "observed_proxy":float(q),
                     "observed_reduction_vs_first":float(r), "predicted_reduction":float(e)}
                    for a,q,r,e in zip(alphas,proxy,scaling,expected)]
    return {"contraction_sum":rows, "k95":k95_rows, "alpha_scaling":scaling_rows, "stability_boundary":stability_sweep()}


def clip(x, C):
    n = np.linalg.norm(x)
    return x if n <= C else x * (C/n)


def run_federated(mode, beta=0.25, sigma0=0.08, alpha=0.08, p=2.0, rounds=80,
                  clients=10, dim=8, seed=1402):
    rng = np.random.default_rng(seed)
    # Each client has a quadratic target; adjacent data changes client 0's target.
    targets = rng.normal(0, 1.0, size=(clients, dim))
    targets -= targets.mean(axis=0, keepdims=True)
    state = np.zeros(dim)
    previous_updates = np.zeros((clients, dim))
    losses=[]; proxy_sum=0.; noise_sum=0.; disagreement=[]
    C=1.5; eta=0.45; delta_fail=1e-5
    for k in range(rounds):
        updates=[]; variances=[]; ds=[]
        for i in range(clients):
            grad = state - targets[i]
            z = clip(-eta*grad, C)
            d = float(np.linalg.norm(z-previous_updates[i]))
            if mode == 'constant':
                var = sigma0**2 + alpha*(1e-6)**p
            elif mode == 'channel':
                var = sigma0**2 + alpha*(d+1e-6)**p
            else:
                var = 0.0
            updates.append(z + rng.normal(0, math.sqrt(var), dim))
            variances.append(var); ds.append(d)
            # The bound uses clipped per-client sensitivity; use a conservative proxy.
            proxy_sum += (float('inf') if var == 0.0 else gaussian_proxy(min(2*C, 0.25), math.sqrt(var), delta_fail))
        aggregate=np.mean(updates,axis=0)
        state=(1-beta)*state + beta*(state+aggregate)
        previous_updates=np.array([clip(-eta*(state-targets[i]), C) for i in range(clients)])
        losses.append(float(np.mean((state-targets)**2)))
        noise_sum += float(np.mean(variances)); disagreement.append(float(np.mean(ds)))
    return {"final_loss":losses[-1], "best_loss":min(losses), "privacy_proxy_sum":proxy_sum,
            "mean_noise_variance":noise_sum/rounds, "mean_disagreement":float(np.mean(disagreement)),
            "loss_curve":losses}


def main():
    math_checks = toy_predictions()
    results={"math_checks":math_checks, "federated":{}}
    for mode in ['none','constant','channel']:
        results['federated'][mode]=run_federated(mode)
    # A separate paired scalar recurrence numerically estimates rho from neighboring states.
    rho_hat=[]; a=1.0; b=1.001; rho=0.8
    for _ in range(20):
        rho_hat.append(abs(rho*b-rho*a)/abs(b-a)); a,b=rho*a,rho*b
    results['estimated_contraction']={"true_rho":rho,"median_rho_hat":float(np.median(rho_hat))}
    with open('results.json','w') as f: json.dump(results,f,indent=2)
    print(json.dumps(results, indent=2))

if __name__ == '__main__': main()