import json, time
import numpy as np


def rank_project(X, q):
    U, s, Vt = np.linalg.svd(X, full_matrices=False)
    return (U[:, :q] * s[:q]) @ Vt[:q, :], U[:, :q], Vt[:q, :].T, s


def tangent_project(H, U, V):
    UtH = U.T @ H
    HV = H @ V
    UtHV = U.T @ H @ V
    return U @ UtH + HV @ V.T - U @ UtHV @ V.T


def normal_project(H, U, V):
    return H - tangent_project(H, U, V)


def adaptive_step(X, B, b, q, c=1.0, rho=1.0, eps=1e-6):
    Y, U, V, s = rank_project(X, q)
    r = np.linalg.norm(X-Y) / (np.linalg.norm(X) + 1e-12)
    mu = c * (r + eps) ** rho
    d = X.size
    # Explicit Q is appropriate for this deliberately small numerical check.
    Q = np.zeros((d, d))
    for j in range(d):
        E = np.zeros((X.shape[0], X.shape[1])); E.flat[j] = 1.0
        Q[:, j] = (normal_project(E, U, V) + mu * E).ravel()
    K = np.block([[Q, B.T], [B, np.zeros((B.shape[0], B.shape[0]))]])
    rhs = np.concatenate([(mu * X + normal_project(Y, U, V)).ravel(), b])
    sol = np.linalg.solve(K, rhs)
    Z = sol[:d].reshape(X.shape)
    return Z, Y, r, mu


def affine_projection(X, B, b):
    # Euclidean projection onto B vec(Z)=b.
    BBt = B @ B.T
    return (X.ravel() + B.T @ np.linalg.solve(BBt, b - B @ X.ravel())).reshape(X.shape)


def verify_math(rng):
    m, n, q = 7, 6, 2
    X = rng.normal(size=(m, n))
    Y, U, V, _ = rank_project(X, q)
    H = rng.normal(size=(m, n))
    PT = tangent_project(H, U, V); PN = normal_project(H, U, V)
    projector_err = max(np.linalg.norm(tangent_project(PT, U, V)-PT),
                        np.linalg.norm(normal_project(PN, U, V)-PN))
    orth_err = abs(np.sum(PT * PN))
    decomp_err = np.linalg.norm(PT + PN - H)
    # The residual rule should have precisely the claimed power-law log slope.
    rs = np.logspace(-6, -1, 20)
    mus = np.array([(rs[i] + 1e-12) ** 0.7 for i in range(len(rs))])
    slope = np.polyfit(np.log(rs), np.log(mus), 1)[0]
    return {"projector_idempotence_error": float(projector_err),
            "projector_orthogonality_error": float(orth_err),
            "projector_decomposition_error": float(decomp_err),
            "mu_log_slope_rho_0.7": float(slope),
            "math_pass": bool(projector_err < 1e-10 and orth_err < 1e-10 and decomp_err < 1e-10 and abs(slope-.7)<1e-3)}


def run_method(name, W0, target, B, b, q, steps=100, c=1.0, rho=1.0):
    W = W0.copy(); hist = []; t0 = time.perf_counter()
    for k in range(steps):
        # Same gradient step for every method.
        X = W - 0.35 * (W - target)
        if name == "alternating":
            Y, _, _, _ = rank_project(X, q)
            W = affine_projection(Y, B, b)
            # Keep the standard alternating scheme's next iterate rank feasible.
            W, _, _, _ = rank_project(W, q)
            mu = None; r = np.linalg.norm(X-Y)/(np.linalg.norm(X)+1e-12)
        else:
            use_rho = 0.0 if name == "fixed_mu" else rho
            Z, Y, r, mu = adaptive_step(X, B, b, q, c=c, rho=use_rho)
            W, _, _, _ = rank_project(Z, q)
        rank_res = np.linalg.norm(W - rank_project(W, q)[0])/(np.linalg.norm(W)+1e-12)
        aff_res = np.linalg.norm(B @ W.ravel() - b)
        err = np.linalg.norm(W-target)/np.linalg.norm(target)
        hist.append((err, rank_res, aff_res, r, -1 if mu is None else mu))
    return {"final_relative_target_error": float(hist[-1][0]),
            "final_affine_residual": float(hist[-1][2]),
            "final_rank_residual": float(hist[-1][1]),
            "initial_affine_residual": float(hist[0][2]),
            "min_error_last_20": float(min(x[0] for x in hist[-20:])),
            "time_sec": time.perf_counter()-t0,
            "history": np.array(hist).tolist()}


def main():
    rng = np.random.default_rng(284)
    math = verify_math(rng)
    m, n, q = 8, 7, 2
    # A rank-q target defines a consistent affine calibration constraint.
    target = rng.normal(size=(m,q)) @ rng.normal(size=(q,n))
    d = m*n
    # Full-row-rank random linear calibration observations.
    B0 = rng.normal(size=(12,d)); Qb, _ = np.linalg.qr(B0.T); B = Qb[:, :12].T
    b = B @ target.ravel()
    W0 = target + 1.8*rng.normal(size=(m,n))
    configs = {"fixed_mu": (1.0, 0.0), "adaptive_rho05": (1.0, 0.5), "adaptive_rho1": (1.0, 1.0), "adaptive_c10_rho1": (10.0, 1.0), "alternating": (0.0, 0.0)}
    results = {name: run_method(name, W0, target, B, b, q, c=c, rho=rho) for name, (c, rho) in configs.items()}
    # Strip histories from summary but save them for reproducibility.
    out = {"math": math, "benchmark": results,
           "setup": {"shape": [m,n], "rank": q, "constraints": B.shape[0], "steps": 100,
                     "learning_rate": .35, "seed": 284}}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps({"math": math, "benchmark_summary":
        {k:{x:v for x,v in val.items() if x != "history"} for k,val in results.items()}}, indent=2))

if __name__ == "__main__": main()