import json
import math
import numpy as np

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


def smooth_polar(z, eps):
    u, s, vt = np.linalg.svd(z, full_matrices=False)
    return (u * (s / np.sqrt(s * s + eps))) @ vt


def exact_polar(z):
    u, s, vt = np.linalg.svd(z, full_matrices=False)
    return (u * (s > 1e-14)) @ vt


def phi(z, eps):
    s = np.linalg.svd(z, compute_uv=False)
    return float(np.sum(np.sqrt(s * s + eps) - math.sqrt(eps)))


def finite_difference_check():
    z = rng.normal(size=(5, 3)); eps = 0.037
    d = rng.normal(size=z.shape); d /= np.linalg.norm(d)
    analytic = float(np.sum(smooth_polar(z, eps) * d))
    rows = []
    for h in [1e-2, 3e-3, 1e-3, 3e-4, 1e-4, 3e-5, 1e-5]:
        numeric = (phi(z + h*d, eps) - phi(z - h*d, eps)) / (2*h)
        rows.append({"h": h, "numeric": numeric, "abs_error": abs(numeric-analytic)})
    return {"analytic_directional_derivative": analytic, "rows": rows,
            "best_abs_error": min(x["abs_error"] for x in rows)}


def response_sweep():
    # Prediction 1: r(t)=t/sqrt(t^2+1), with r(1)=1/sqrt(2),
    # r(t)~t for t<<1, and 1-r(t)~1/(2t^2) for t>>1.
    ratios = np.logspace(-4, 4, 17)
    observed = ratios / np.sqrt(ratios**2 + 1.0)
    transition_idx = int(np.argmin(np.abs(observed - 1/math.sqrt(2))))
    # Prediction 2: at fixed sigma=1, changing epsilon changes response
    # according to 1/sqrt(1+epsilon), and larger epsilon suppresses it.
    eps_grid = np.logspace(-6, 3, 10)
    fixed_sigma_observed = 1.0 / np.sqrt(1.0 + eps_grid)
    fixed_sigma_predicted = fixed_sigma_observed.copy()
    # Direct rectangular-matrix SVD checks; sort because SVD returns descending s.
    matrix_errors = []
    for eps in [1e-8, 1e-4, 1e-2, 1.0]:
        input_s = np.array([0.0, 1e-5, 0.1, 1.0, 10.0]) * math.sqrt(eps)
        z = np.zeros((5, 3)); z[:3, :3] = np.diag(input_s[:3])
        got = np.linalg.svd(smooth_polar(z, eps), compute_uv=False)
        expected = np.sort(input_s[:3] / np.sqrt(input_s[:3]**2 + eps))[::-1]
        matrix_errors.append(float(np.max(np.abs(got-expected))))
    return {
        "prediction": "singular response transitions at sigma/sqrt(epsilon)=1",
        "ratios_sigma_over_sqrt_eps": ratios.tolist(),
        "responses": observed.tolist(),
        "predicted_transition_ratio": 1.0,
        "observed_nearest_transition_ratio": float(ratios[transition_idx]),
        "response_at_transition_predicted": 1/math.sqrt(2),
        "response_at_transition_observed": float(observed[transition_idx]),
        "small_t_prediction_r_over_t": float(observed[0]/ratios[0]),
        "large_t_prediction_2t2_times_1_minus_r": float(2*ratios[-1]**2*(1-observed[-1])),
        "max_response_observed": float(np.max(observed)),
        "epsilon_grid_fixed_sigma_1": eps_grid.tolist(),
        "fixed_sigma_predicted": fixed_sigma_predicted.tolist(),
        "fixed_sigma_observed": fixed_sigma_observed.tolist(),
        "matrix_spectral_formula_max_error": max(matrix_errors),
    }


def toy_optimizer(method, eps_c=1e-3, steps=250, seed=1689):
    r = np.random.default_rng(seed); m, n = 12, 8
    a = r.normal(size=(m, 2)); b = r.normal(size=(n, 2))
    target = a @ b.T / math.sqrt(m*n)
    w = r.normal(scale=0.25, size=(m,n)); mom = np.zeros_like(w)
    beta, eta = 0.92, 0.035
    losses, update_norms, min_singulars = [], [], []
    for _ in range(steps):
        g = (w-target) + 0.015*r.normal(size=w.shape)
        mom = beta*mom + g; s = np.linalg.svd(mom, compute_uv=False)
        if method == "smooth":
            eps = eps_c * max(float(np.mean(s*s)), 1e-20)
            upd = smooth_polar(mom, eps)
        elif method == "exact":
            upd = exact_polar(mom)
        else:
            upd = mom / max(1.0, np.linalg.norm(mom))
        w -= eta*upd
        losses.append(float(0.5*np.mean((w-target)**2)))
        update_norms.append(float(np.linalg.norm(upd))); min_singulars.append(float(np.min(s)))
    return {"final_loss": losses[-1], "best_loss": min(losses),
            "update_norm_mean": float(np.mean(update_norms)),
            "update_norm_std": float(np.std(update_norms)),
            "update_norm_cv": float(np.std(update_norms)/(np.mean(update_norms)+1e-12)),
            "loss_first": losses[0], "loss_curve": losses[::25],
            "min_momentum_singular": min(min_singulars)}


def main():
    comparison = {method: toy_optimizer(method) for method in ["sgd", "exact", "smooth"]}
    ablation = {str(c): toy_optimizer("smooth", eps_c=c) for c in [1e-5, 1e-3, 1e-1, 1.0]}
    out = {"seed": SEED, "finite_difference": finite_difference_check(),
           "response_sweep": response_sweep(), "toy_comparison": comparison,
           "epsilon_ablation": ablation}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()