"""Finite-horizon walk reciprocity control: verification and tiny training study."""
import json, math, random
from pathlib import Path
import numpy as np


def phi_star(g):
    return 1.0 - math.sqrt(max(0.0, 1.0 - g * g))


def walk_energy(A, L=4, g=0.8, normalize=True, eps=1e-8):
    A = np.asarray(A, dtype=np.float64)
    if normalize:
        s = np.linalg.svd(A, compute_uv=False)[0]
        A = A / (s + eps)
    n = A.shape[0]
    V = A.copy()
    terms = []
    total = 0.0
    for k in range(1, L + 1):
        D = V - V.T
        term = (g ** (2 * k - 2)) * np.sum(D * D) / n
        terms.append(float(term))
        total += term
        V = A @ V
    return float(total), terms


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


def math_sanity():
    n = 8
    # Strictly lower-triangular directed chain: nilpotent but transient walks exist.
    A = np.zeros((n, n))
    for i in range(n - 1):
        A[i + 1, i] = 1.0
    S = A + A.T                         # reciprocal control
    rn, tn = walk_energy(A, L=6, g=.8)
    rs, ts = walk_energy(S, L=6, g=.8)
    # Verify the iterative powers used by the implementation against A**k.
    An = A / (np.linalg.svd(A, compute_uv=False)[0] + 1e-8)
    V = An.copy()
    recursion_errors = []
    for k in range(1, 7):
        recursion_errors.append(float(np.max(np.abs(V - np.linalg.matrix_power(An, k)))))
        V = An @ V
    return {
        "nilpotent_spectral_radius": spectral_radius(A),
        "nilpotent_A_power_n_frobenius": float(np.linalg.norm(np.linalg.matrix_power(A, n))),
        "nilpotent_walk_energy": rn,
        "nilpotent_terms": tn,
        "symmetric_control_walk_energy": rs,
        "symmetric_control_terms": ts,
        "phi_star_0.8": phi_star(.8),
        "max_iterative_power_error": max(recursion_errors),
        "energy_positive_before_nilpotency": rn > 0,
    }


def train_compare(seed=7, steps=500, lam=0.0, target=0.0):
    import torch
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        x = torch.randn(96, 12, device=device)
        true_w = torch.randn(12, 12, device=device) / math.sqrt(12)
        y = torch.tanh(x @ true_w.T)
        model = torch.nn.Sequential(torch.nn.Linear(12, 12), torch.nn.Tanh(), torch.nn.Linear(12, 12)).to(device)
        opt = torch.optim.Adam(model.parameters(), lr=3e-3)
        losses=[]; energies=[]; gradnorms=[]
        g=.8; L=4
        for step in range(steps):
            opt.zero_grad(set_to_none=True)
            pred=model(x)
            loss=torch.mean((pred-y)**2)
            W=model[0].weight
            scale=torch.linalg.matrix_norm(W, ord=2) + 1e-6
            An=W/scale
            V=An
            R=torch.zeros((), device=device)
            for k in range(1,L+1):
                D=V-V.T
                R=R+(g**(2*k-2))*torch.sum(D*D)/W.shape[0]
                V=An@V
            total=loss+lam*torch.relu(R-target)**2
            total.backward()
            gn=0.0
            for p in model.parameters():
                if p.grad is not None: gn += float(torch.sum(p.grad.detach()**2).cpu())
            opt.step()
            if step >= steps-50:
                losses.append(float(loss.detach().cpu())); energies.append(float(R.detach().cpu())); gradnorms.append(math.sqrt(gn))
        return {"device":device, "final_loss_mean_last50":float(np.mean(losses)),
                "walk_energy_mean_last50":float(np.mean(energies)),
                "gradient_norm_mean_last50":float(np.mean(gradnorms)), "lambda":lam, "target":target}
    except (RuntimeError, torch.cuda.OutOfMemoryError):
        if device == "cuda":
            torch.cuda.empty_cache()
            return train_compare_cpu(seed, steps, lam, target)
        raise


def train_compare_cpu(seed, steps, lam, target):
    import torch
    old=torch.cuda.is_available
    torch.cuda.is_available=lambda: False
    try: return train_compare(seed, steps, lam, target)
    finally: torch.cuda.is_available=old


def main():
    sanity=math_sanity()
    # Per-seed baseline calibration follows the proposed first-epoch-style calibration.
    rows=[]
    for seed in (7, 19, 31):
        base=train_compare(seed=seed, lam=0.0)
        controlled=train_compare(seed=seed, lam=2.0, target=base["walk_energy_mean_last50"])
        rows.append({"seed":seed, "baseline":base, "walk_regularized":controlled})
    def avg(key, method):
        return float(np.mean([r[method][key] for r in rows]))
    out={"sanity":sanity, "per_seed":rows,
         "mean_last50": {
             "baseline_loss":avg("final_loss_mean_last50", "baseline"),
             "idea_loss":avg("final_loss_mean_last50", "walk_regularized"),
             "baseline_walk_energy":avg("walk_energy_mean_last50", "baseline"),
             "idea_walk_energy":avg("walk_energy_mean_last50", "walk_regularized"),
             "baseline_gradient_norm":avg("gradient_norm_mean_last50", "baseline"),
             "idea_gradient_norm":avg("gradient_norm_mean_last50", "walk_regularized")},
         "all_seeds_lower_energy": all(r["walk_regularized"]["walk_energy_mean_last50"] < r["baseline"]["walk_energy_mean_last50"] for r in rows),
         "all_seeds_lower_loss": all(r["walk_regularized"]["final_loss_mean_last50"] < r["baseline"]["final_loss_mean_last50"] for r in rows)}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()