import json
import numpy as np


def make_certified(r=8, d=6, seed=0):
    rng = np.random.default_rng(seed)
    # C=B^T makes the KYP off-diagonal blocks vanish. Positive A,D
    # then give a strict certificate while retaining a nontrivial rational term.
    A = 0.35 * np.eye(r)
    B = rng.normal(size=(r, d)) / np.sqrt(d)
    C = B.T
    D = 0.12 * np.eye(d)
    P1 = np.diag([1.0] * (r // 2) + [0.0] * (r - r // 2))
    P2 = np.eye(r) - P1
    return A, B, C, D, P1, P2


def kyp_matrix(A, B, C, D):
    return np.block([[A + A.T, C.T - B], [C - B.T, D + D.T]])


def H_of_z(z1, z2, A, B, C, D, P1, P2):
    M = A + z1 * P1 + z2 * P2
    return D + C @ np.linalg.solve(M, B)


def resolvent(x, eta, H):
    return np.linalg.solve(np.eye(H.shape[0]) + eta * H, x)


def verify(seed=0):
    A, B, C, D, P1, P2 = make_certified(seed=seed)
    eig_kyp = np.linalg.eigvalsh(kyp_matrix(A, B, C, D))
    rng = np.random.default_rng(seed + 11)
    min_herm = np.inf
    max_sigma = 0.0
    max_energy_ratio = 0.0
    samples = []
    for z1, z2 in np.exp(rng.uniform(np.log(1e-3), np.log(20.0), size=(500, 2))):
        H = H_of_z(z1, z2, A, B, C, D, P1, P2)
        herm_eig = np.linalg.eigvalsh((H + H.T) / 2).min()
        R = np.linalg.inv(np.eye(H.shape[0]) + 0.8 * H)
        min_herm = min(min_herm, herm_eig)
        max_sigma = max(max_sigma, np.linalg.svd(R, compute_uv=False)[0])
        x = rng.normal(size=H.shape[0])
        max_energy_ratio = max(max_energy_ratio, np.linalg.norm(R @ x) / np.linalg.norm(x))
    # The control uses the same positive H but an explicit additive residual.
    H0 = H_of_z(0.5, 0.5, A, B, C, D, P1, P2)
    x = rng.normal(size=H0.shape[0])
    implicit_norms, additive_norms = [np.linalg.norm(x)], [np.linalg.norm(x)]
    for _ in range(24):
        x = resolvent(x, 0.8, H0)
        implicit_norms.append(np.linalg.norm(x))
    x = rng.normal(size=H0.shape[0])
    for _ in range(24):
        x = x + 0.8 * (H0 @ x)
        additive_norms.append(np.linalg.norm(x))
    return {
        "kyp_min_eigenvalue": float(eig_kyp.min()),
        "min_hermitian_part_eigenvalue": float(min_herm),
        "max_resolvent_singular_value": float(max_sigma),
        "max_sampled_energy_ratio": float(max_energy_ratio),
        "implicit_norm_start_end": [float(implicit_norms[0]), float(implicit_norms[-1])],
        "additive_norm_start_end": [float(additive_norms[0]), float(additive_norms[-1])],
        "implicit_norm_max": float(max(implicit_norms)),
        "additive_norm_max": float(max(additive_norms)),
    }


def toy_training(seed=123, n=128, steps=40, depth=8, eta=0.8):
    # Same fixed operator and same initial batch. Target is the stable zero state;
    # this isolates whether the update itself suppresses activation growth.
    rng = np.random.default_rng(seed)
    A, B, C, D, P1, P2 = make_certified(seed=seed)
    H = H_of_z(0.5, 0.5, A, B, C, D, P1, P2)
    x0 = rng.normal(size=(n, H.shape[0]))
    baseline = x0.copy()
    idea = x0.copy()
    rows = []
    for step in range(steps + 1):
        rows.append({
            "step": step,
            "baseline_loss": float(np.mean(baseline ** 2)),
            "idea_loss": float(np.mean(idea ** 2)),
            "baseline_activation_rms": float(np.sqrt(np.mean(baseline ** 2))),
            "idea_activation_rms": float(np.sqrt(np.mean(idea ** 2))),
            "baseline_max_norm": float(np.max(np.linalg.norm(baseline, axis=1))),
            "idea_max_norm": float(np.max(np.linalg.norm(idea, axis=1))),
        })
        if step == steps:
            break
        # Standard explicit residual versus the proposed implicit resolvent.
        baseline = baseline + eta * (baseline @ H.T)
        idea = np.stack([resolvent(v, eta, H) for v in idea])
    # Explicit residual controls at smaller learning/update coefficients.
    sweep = {}
    for explicit_eta in (0.05, 0.1, 0.2, 0.4, 0.8):
        z = x0.copy()
        for _ in range(steps):
            z = z + explicit_eta * (z @ H.T)
        sweep[str(explicit_eta)] = {
            "final_rms": float(np.sqrt(np.mean(z ** 2))),
            "max_rms_during_run": float(np.sqrt(np.mean(x0 ** 2))) if explicit_eta == 0 else float(np.sqrt(np.mean(z ** 2))),
        }
    return {"operator_eigenvalues": np.linalg.eigvalsh(H).tolist(), "trace": rows, "explicit_step_sweep_final_rms": sweep}


def main():
    out = {"verification": verify(), "toy_training": toy_training()}
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()