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

SEED = 2688
rng = np.random.default_rng(SEED)


def mgs(X, eps=1e-10):
    """Columns of X -> orthonormal columns, rejecting dependent columns."""
    qs = []
    residuals = []
    for i in range(X.shape[1]):
        v = X[:, i].copy()
        for q in qs:
            v -= q * (q @ v)
        nv = np.linalg.norm(v)
        residuals.append(float(nv))
        if nv > eps:
            qs.append(v / nv)
    Q = np.stack(qs, axis=1) if qs else np.zeros((X.shape[0], 0))
    return Q, np.asarray(residuals)


def orthogonality_check():
    d = 8
    # Deliberately badly conditioned but full-rank columns.
    U, _ = np.linalg.qr(rng.normal(size=(d, d)))
    V, _ = np.linalg.qr(rng.normal(size=(d, d)))
    X = U @ np.diag(np.geomspace(1.0, 1e-7, d)) @ V.T
    Q, res = mgs(X)
    return {
        "dimension": d,
        "sigma_min_feature_matrix": float(np.linalg.svd(X, compute_uv=False)[-1]),
        "mgs_rank": int(Q.shape[1]),
        "max_abs_QtQ_minus_I": float(np.max(np.abs(Q.T @ Q - np.eye(d)))),
        "min_mgs_residual": float(res.min()),
    }


def stability_sweep():
    # With a complete orthonormal basis, e <- (I-alpha I)e and norm ratio=|1-alpha|.
    d = 8
    Q = np.eye(d)
    e0 = rng.normal(size=d)
    e0 /= np.linalg.norm(e0)
    alphas = np.array([0.25, 0.75, 1.0, 1.25, 1.75, 1.99, 2.01, 2.5])
    rows = []
    for a in alphas:
        e = e0.copy()
        for _ in range(30):
            e = e - a * Q @ (Q.T @ e)
        ratio = float(np.linalg.norm(e)) # initial norm is 1
        predicted = abs(1.0 - a) ** 30
        rows.append({"alpha": float(a), "predicted_norm_after_30": predicted,
                     "observed_norm_after_30": ratio,
                     "observed_stable": bool(ratio < 1.0)})
    # empirical boundary is first alpha where amplification per step exceeds 1
    boundary = 2.0
    return {"predicted_boundary_alpha": boundary, "rows": rows}


def excitation_transition():
    # Memory of k orthogonal directions. Only those coordinates contract.
    d = 8
    alpha = 0.8
    e0 = np.ones(d) / math.sqrt(d)
    rows = []
    for k in range(d + 1):
        Q = np.eye(d)[:, :k]
        e = e0.copy()
        norms = []
        for t in range(16):
            norms.append(float(np.linalg.norm(e)))
            e = e - alpha * Q @ (Q.T @ e)
        # excited coordinates have exact predicted slope log|1-alpha|; unexcited remain.
        expected = math.sqrt((k / d) * abs(1-alpha) ** (2*15) + (d-k)/d)
        rows.append({"independent_directions": k,
                     "norm_at_start": norms[0], "norm_at_step_15": norms[-1],
                     "predicted_norm_at_step_15": expected,
                     "all_direction_contraction": bool(k == d)})
    return {"predicted_transition_k": d, "rows": rows}


def replay_comparison():
    # Same realizable linear task, with feature matrices having equal trace but
    # different conditioning. Raw replay uses X X^T; MGS uses Q Q^T.
    d = 8
    wstar = rng.normal(size=d)
    wstar /= np.linalg.norm(wstar)
    U, _ = np.linalg.qr(rng.normal(size=(d, d)))
    cases = [("well_conditioned", np.ones(d)),
             ("ill_conditioned", np.geomspace(1.0, 1e-3, d))]
    alpha = 0.8
    out = []
    for name, sing in cases:
        X = U @ np.diag(sing)
        y = X.T @ wstar
        Q, _ = mgs(X)
        # Normalize raw replay's mean Gramian to make it a fair gradient step;
        # conditioning, rather than overall scale, determines the rate.
        G = X @ X.T
        G = G / np.trace(G) * d
        eb = rng.normal(size=d); eb /= np.linalg.norm(eb)
        eo = eb.copy()
        raw_norms = []; ortho_norms = []
        for t in range(60):
            raw_norms.append(float(np.linalg.norm(eb)))
            ortho_norms.append(float(np.linalg.norm(eo)))
            eb = eb - alpha * G @ eb
            eo = eo - alpha * Q @ (Q.T @ eo)
        raw_rate = float(np.polyfit(np.arange(20, 60), np.log(np.maximum(raw_norms[20:], 1e-300)), 1)[0])
        ortho_rate = float(np.polyfit(np.arange(20, 60), np.log(np.maximum(ortho_norms[20:], 1e-300)), 1)[0])
        out.append({"case": name, "sigma_min": float(sing.min()),
                    "raw_norm_step_10": raw_norms[10], "mgs_norm_step_10": ortho_norms[10],
                    "raw_log_slope_late": raw_rate, "mgs_log_slope_late": ortho_rate})
    return {"predicted_mgs_slope": math.log(abs(1-alpha)), "rows": out}


def main():
    result = {"seed": SEED,
              "orthogonality": orthogonality_check(),
              "stability_sweep": stability_sweep(),
              "excitation_transition": excitation_transition(),
              "replay_comparison": replay_comparison()}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()