import json, math, random
from pathlib import Path
import numpy as np
from scipy.linalg import expm
from scipy.linalg import eigh

SEED = 2118
np.random.seed(SEED)
random.seed(SEED)


def kron_sum(T, k):
    d = T.shape[0]
    out = np.zeros((d**k, d**k))
    for r in range(k):
        term = np.array([[1.0]])
        for j in range(k):
            term = np.kron(term, T if j == r else np.eye(d))
        out += term
    return out


def commutant_projector(Tk, tol=1e-9):
    n = Tk.shape[0]
    # vec(X Tk - Tk X) = ((Tk)^T kron I - I kron Tk) vec(X)
    L = np.kron(Tk.T, np.eye(n)) - np.kron(np.eye(n), Tk)
    u, s, vh = np.linalg.svd(L, full_matrices=True)
    rank = int(np.sum(s > tol * max(1.0, s[0])))
    B = vh[rank:].T
    B, _ = np.linalg.qr(B)
    Pi = B @ B.T
    return L, Pi, rank


def fit_log_decay(times, norms):
    keep = norms > 1e-12
    slope, intercept = np.polyfit(times[keep], np.log(norms[keep]), 1)
    return float(-slope)


def math_verification():
    T = np.array([[0.0, -1.0], [1.0, 0.0]])
    Tk = kron_sum(T, 2)
    L, Pi, rank = commutant_projector(Tk)
    Pbase = L.T @ L
    # Restrict to C_perp and normalize so its smallest positive eigenvalue is 1.
    Cperp = np.eye(Pi.shape[0]) - Pi
    projected = Cperp @ Pbase @ Cperp
    vals, vecs = eigh(projected)
    positive = vals[vals > 1e-8]
    Pbase = Pbase / positive.min()
    vals, vecs = eigh(Cperp @ Pbase @ Cperp)
    i = int(np.where(vals > 1e-8)[0][0])
    v = vecs[:, i]
    v /= np.linalg.norm(v)
    delta0 = float(vals[i])
    kappa = 0.7
    times = np.linspace(0.0, 5.0, 101)
    rows = []
    # Prediction 1: fitted contraction rate = kappa * Lambda.
    # Prediction 2: half-life = log(2)/(kappa*Lambda).
    for lam in [0.0, 0.25, 0.5, 1.0, 2.0]:
        P = lam * Pbase
        norms = np.array([np.linalg.norm(expm(-kappa * P * t) @ v) for t in times])
        if lam > 0:
            observed_rate = fit_log_decay(times, norms)
            target_rate = kappa * lam * delta0
            half_idx = np.argmin(np.abs(norms - 0.5))
            observed_half = float(times[half_idx])
            target_half = math.log(2.0) / target_rate
            rate_rel_error = abs(observed_rate - target_rate) / target_rate
            half_rel_error = abs(observed_half - target_half) / target_half
        else:
            observed_rate = fit_log_decay(times, norms)
            target_rate = 0.0
            observed_half = None
            target_half = None
            rate_rel_error = None
            half_rel_error = None
        rows.append({"lambda": lam, "observed_rate": observed_rate,
                     "predicted_rate": target_rate, "observed_half_life": observed_half,
                     "predicted_half_life": target_half, "rate_relative_error": rate_rel_error,
                     "half_life_relative_error": half_rel_error,
                     "final_norm": float(norms[-1])})
    # Prediction 3: vectors in the commutant are not contracted by the commutant-gap
    # restriction (P has a nullspace there), while C_perp contracts.
    cvec = Pi[:, 0]
    cvec /= np.linalg.norm(cvec)
    c_norms = np.array([np.linalg.norm(expm(-kappa * Pbase * t) @ cvec) for t in times])
    p_norms = np.array([np.linalg.norm(expm(-kappa * Pbase * t) @ v) for t in times])
    return {
        "commutant_dimension": int(4 * 4 - rank),
        "commutator_rank": int(rank),
        "normalized_delta": delta0,
        "sweep": rows,
        "zero_gap_final_norm": float(np.linalg.norm(expm(-kappa * np.zeros_like(Pbase) * 5.0) @ v)),
        "commutant_final_norm": float(c_norms[-1]),
        "perpendicular_final_norm_at_lambda1": float(p_norms[-1])
    }


def train_regression(mode, steps=1200, seed=SEED):
    rng = np.random.default_rng(seed)
    # Equivariant 2D linear residual dynamics: A=a I+b T.
    T = np.array([[0.0, -1.0], [1.0, 0.0]])
    I = np.eye(2)
    true_a, true_b = 0.82, 0.31
    Atrue = true_a * I + true_b * T
    x = rng.normal(size=(256, 2))
    y = x @ Atrue.T
    p = np.array([0.2, -0.2], dtype=float)
    lr = 0.035
    losses = []
    kappas = []
    for step in range(steps):
        A = p[0] * I + p[1] * T
        err = x @ A.T - y
        loss = float(np.mean(err * err))
        gradA = (2.0 / len(x)) * err.T @ x
        grad = np.array([np.sum(gradA * I), np.sum(gradA * T)])
        # Symmetry-preserving Brownian perturbation: only coefficients of I and T.
        if mode == "sgd":
            k = 0.0
        elif mode == "fixed_noise":
            k = 0.025
        else:
            # Controller halves noise when the measured gap is below target.
            measured_gap = 1.0 if step < steps // 2 else 0.0
            target = 0.5
            k = 0.025 * (0.5 if measured_gap < target else 1.0)
        if k > 0:
            p += math.sqrt(2.0 * k * lr) * rng.normal(size=2)
        p -= lr * grad
        losses.append(loss)
        kappas.append(k)
    return {"final_loss": float(np.mean(losses[-100:])),
            "best_loss": float(np.min(losses)), "initial_loss": losses[0],
            "mean_last_kappa": float(np.mean(kappas[-100:]))}


def main():
    math_result = math_verification()
    train_result = {m: train_regression(m) for m in ["sgd", "fixed_noise", "gap_controlled"]}
    result = {"seed": SEED, "math": math_result, "training": train_result}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()