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

SEED = 17
TREE = {0: [1, 2], 1: [3], 2: []}
PARENT = {1: 0, 2: 0, 3: 1}
N = 4


def correction_terms(u, s, rho):
    """One dual update and the exact parent corrections -sum(delta s)."""
    s_new = s.copy()
    for child, parent in PARENT.items():
        s_new[child] += rho * (u[child] - u[parent])
    corr = np.zeros(N)
    for parent, children in TREE.items():
        corr[parent] = -sum(s_new[ch] - s[ch] for ch in children)
    return s_new, corr


def tree_step(u, s, targets, rho, eta, penalty, corrected=True):
    # Jacobi primal step. Each node reads old neighbors; the root/parents can
    # additionally read child dual increments through the correction channel.
    s_new, corr = correction_terms(u, s, rho)
    grad = u - targets
    for child, parent in PARENT.items():
        d = u[child] - u[parent]
        grad[parent] -= s[child] + penalty * d
        grad[child] += s[child] + penalty * d
    if corrected:
        # At a parent, -delta s is precisely the paper's correction term.
        grad += corr
    unew = u - eta * grad
    return unew, s_new, corr


def corrected_tree(targets, rounds=160, rho=.20, eta=.16, penalty=.8, corrected=True):
    u = np.zeros(N)
    s = np.zeros(N)
    losses, residuals, bytes_sent = [], [], 0
    for _ in range(rounds):
        u, s, _ = tree_step(u, s, targets, rho, eta, penalty, corrected)
        loss = .5 * np.sum((u-targets)**2) + .5 * penalty * sum((u[ch]-u[p])**2 for ch,p in PARENT.items())
        losses.append(float(loss))
        residuals.append(float(np.sqrt(np.mean([(u[ch]-u[p])**2 for ch,p in PARENT.items()]))))
        # Each edge sends a scalar state and a scalar correction in this toy.
        bytes_sent += 2 * len(PARENT) * 8
    return u, losses, residuals, bytes_sent


def fedavg(targets, rounds=160, eta=.16, local_steps=1):
    u = np.zeros(N)
    losses, residuals = [], []
    for _ in range(rounds):
        proposals = []
        for x in u:
            z = x
            for _ in range(local_steps):
                z -= eta * (z - targets[len(proposals)])
            proposals.append(z)
        mean = np.mean(proposals)
        u[:] = mean
        loss = .5 * np.sum((u-targets)**2) + .5*.8*sum((u[ch]-u[p])**2 for ch,p in PARENT.items())
        losses.append(float(loss))
        residuals.append(float(np.sqrt(np.mean([(u[ch]-u[p])**2 for ch,p in PARENT.items()]))))
    # four workers upload and root broadcasts a scalar each round
    return u, losses, residuals, rounds * 2 * N * 8


def math_checks():
    # Prediction 1: correction conservation error is numerical zero for every rho.
    u = np.array([.2, 1.3, -.7, .4])
    s = np.array([0., .8, -.2, .5])
    conservation = []
    for rho in [.01, .1, .4, .9]:
        sn, corr = correction_terms(u, s, rho)
        expected = np.array([-(sn[1]-s[1])-(sn[2]-s[2]),
                             -(sn[3]-s[3]), 0., 0.])
        conservation.append(float(abs(corr[0]-expected[0]) + abs(corr[1]-expected[1])))

    # Prediction 2: fixed-state correction norm is exactly linear in rho.
    norms = []
    for rho in [.05, .10, .20, .40]:
        _, c = correction_terms(u, s, rho)
        norms.append(float(np.linalg.norm(c)))
    ratios = [norms[i]/norms[0] for i in range(4)]
    predicted_ratios = [1., 2., 4., 8.]

    # Prediction 3: stale parent dual error is the correction magnitude and
    # therefore has zero intercept and linear rho scaling.
    stale = norms
    slope = float(np.dot([.05,.1,.2,.4], stale) / np.dot([.05,.1,.2,.4], [.05,.1,.2,.4]))
    intercept = float(np.mean(stale) - slope*np.mean([.05,.1,.2,.4]))
    return {
        "conservation_abs_errors": conservation,
        "correction_norms": norms,
        "rho_ratio_observed": ratios,
        "rho_ratio_predicted": predicted_ratios,
        "stale_error_linear_fit_slope": slope,
        "stale_error_linear_fit_intercept": intercept,
        "max_conservation_error": max(conservation),
    }


def main():
    np.random.seed(SEED)
    targets = np.array([-2.0, 1.0, 2.5, -1.0])
    checks = math_checks()
    sweep = []
    for rho in [.0, .05, .10, .20, .40]:
        uc, lc, rc, bc = corrected_tree(targets, rho=rho, corrected=True)
        us, ls, rs, bs = corrected_tree(targets, rho=rho, corrected=False)
        sweep.append({
            "rho": rho,
            "corrected_final_loss": lc[-1], "stale_final_loss": ls[-1],
            "corrected_final_residual": rc[-1], "stale_final_residual": rs[-1],
            "corrected_loss_40": lc[39], "stale_loss_40": ls[39],
        })
    uc, lc, rc, bc = corrected_tree(targets, rho=.20, corrected=True)
    uf, lf, rf, bf = fedavg(targets)
    result = {
        "seed": SEED, "targets": targets.tolist(), "math_checks": checks,
        "rho_sweep": sweep,
        "mini_experiment": {
            "tree_correction": {"final_loss": lc[-1], "final_residual": rc[-1], "bytes": bc},
            "fedavg": {"final_loss": lf[-1], "final_residual": rf[-1], "bytes": bf},
            "stale_tree_rho_0.2": {"final_loss": corrected_tree(targets, rho=.2, corrected=False)[1][-1],
                                    "final_residual": corrected_tree(targets, rho=.2, corrected=False)[2][-1],
                                    "bytes": bc},
        },
        "interpretation": "Checks test algebraic correction conservation and rho scaling; optimization comparison is a small scalar quadratic proxy, not a neural-network benchmark."
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()