"""Minimal numerical test of a power-preserving formation GNN.

Run with:
  /home/maxwelhelp/main/bin/python3 power_formation_gnn.py
"""
import json
import numpy as np


def formation_matrix(n, d, edges):
    B = np.zeros((len(edges), n))
    for k, (i, j) in enumerate(edges):
        B[k, i] = 1.0
        B[k, j] = -1.0
    return np.kron(B, np.eye(d))


def port_hamiltonian_matrices(F, damping_node=0.15, damping_edge=0.15):
    nd, me = F.shape
    J = np.block([[np.zeros((nd, nd)), -F.T], [F, np.zeros((me, me))]])
    R = np.diag(np.r_[np.full(nd, damping_node), np.full(me, damping_edge)])
    return J, R


def energy(x, mass):
    return 0.5 * np.sum(x * x / mass)


def verify_math(rng):
    F = formation_matrix(4, 2, [(0, 1), (1, 2), (2, 3), (0, 3)])
    J, R = port_hamiltonian_matrices(F)
    mass = np.ones(J.shape[0])
    x = rng.normal(size=J.shape[0])
    grad = x / mass
    interconnection_power = float(grad @ J @ grad)
    dissipation_power = float(grad @ R @ grad)
    dx = (J - R) @ grad
    analytic_dh = float(grad @ dx)
    eps = 1e-7
    finite_difference = (energy(x + eps * dx, mass) - energy(x, mass)) / eps
    return {
        "skew_error": float(np.max(np.abs(J + J.T))),
        "power_cancellation_abs": abs(interconnection_power),
        "dissipation_power": dissipation_power,
        "analytic_dH": analytic_dh,
        "finite_difference_dH": float(finite_difference),
        "derivative_match_abs": abs(analytic_dh - finite_difference),
    }


def simulate_ph(J, R, x0, dt, steps):
    x = x0.copy()
    hist = [energy(x, np.ones_like(x))]
    for _ in range(steps):
        x = x + dt * ((J - R) @ x)
        hist.append(energy(x, np.ones_like(x)))
    return np.asarray(hist), x


def stability_sweep(rng):
    F = formation_matrix(4, 2, [(0, 1), (1, 2), (2, 3), (0, 3)])
    J, R = port_hamiltonian_matrices(F, 0.2, 0.2)
    x0 = rng.normal(size=J.shape[0])
    stated_bound = 2.0 / np.linalg.eigvalsh(R).max()
    # Actual linear explicit-Euler condition is rho(I + dt(J-R)) <= 1.
    A = J - R
    eig = np.linalg.eigvals(A)
    actual_limit = np.inf
    for lam in eig:
        if lam.real < 0:
            candidate = -2.0 * lam.real / (abs(lam) ** 2)
            actual_limit = min(actual_limit, candidate)
    dts = [0.001, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 10.0]
    records = []
    for dt in dts:
        h, _ = simulate_ph(J, R, x0, dt, 100)
        records.append({
            "dt": dt,
            "spectral_radius": float(max(abs(1.0 + dt * eig))),
            "final_energy": float(h[-1]),
            "max_energy": float(h.max()),
        })
    return {
        "stated_damping_only_bound": float(stated_bound),
        "actual_linear_euler_limit": float(actual_limit),
        "records": records,
    }


def run_depth_experiment(rng):
    n, d = 12, 2
    edges = [(i, i + 1) for i in range(n - 1)] + [(0, n - 1)]
    F = formation_matrix(n, d, edges)
    J, R = port_hamiltonian_matrices(F, damping_node=0.12, damping_edge=0.12)
    x0 = rng.normal(size=J.shape[0])
    node0 = x0[: n * d]
    L = F.T @ F
    A = np.eye(n * d) - L / max(np.linalg.eigvalsh(L).max(), 1e-12)
    unconstrained = node0.copy()
    ph = x0.copy()
    rows = []
    for k in range(51):
        rows.append({
            "step": k,
            "unconstrained_node_norm": float(np.linalg.norm(unconstrained)),
            "ph_total_energy": float(energy(ph, np.ones_like(ph))),
            "ph_node_norm": float(np.linalg.norm(ph[: n * d])),
        })
        unconstrained = 1.18 * (A @ unconstrained)
        ph = ph + 0.08 * ((J - R) @ ph)
    return rows


def main():
    rng = np.random.default_rng(3165)
    math = verify_math(rng)
    stability = stability_sweep(rng)
    depth = run_depth_experiment(rng)
    out = {
        "math": math,
        "stability": stability,
        "depth_summary": {
            "unconstrained_norm_step0": depth[0]["unconstrained_node_norm"],
            "unconstrained_norm_step50": depth[-1]["unconstrained_node_norm"],
            "ph_energy_step0": depth[0]["ph_total_energy"],
            "ph_energy_step50": depth[-1]["ph_total_energy"],
            "ph_max_energy": max(x["ph_total_energy"] for x in depth),
            "ph_node_norm_step0": depth[0]["ph_node_norm"],
            "ph_node_norm_step50": depth[-1]["ph_node_norm"],
        },
    }
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()