import json
import time
import numpy as np


def tangent_projector(U, V, G):
    UU = U @ U.T
    VV = V @ V.T
    return UU @ G + G @ VV - UU @ G @ VV


def sketch_recondition(U, V, rng, oversampling=8, eps=1e-12):
    m, r = U.shape
    S = rng.normal(size=(r + oversampling, m)) / np.sqrt(r + oversampling)
    Y = S @ U
    C = Y.T @ Y + eps * np.eye(r)
    R = np.linalg.cholesky(C).T
    Ri = np.linalg.solve(R, np.eye(r))
    return U @ Ri, V @ R.T


def exact_recondition(U, V):
    Q, R = np.linalg.qr(U, mode="reduced")
    signs = np.where(np.diag(R) >= 0.0, 1.0, -1.0)
    D = np.diag(signs)
    return Q @ D, V @ R.T @ D


def metrics(U, V, target):
    W = U @ V.T
    return {
        "loss": float(0.5 * np.mean((W - target) ** 2)),
        "u_orth": float(np.linalg.norm(U.T @ U - np.eye(U.shape[1]))),
        "v_orth": float(np.linalg.norm(V.T @ V - np.eye(V.shape[1]))),
        "cond_u": float(np.linalg.cond(U)),
        "cond_v": float(np.linalg.cond(V)),
    }


def math_check(seed=7):
    rng = np.random.default_rng(seed)
    m, n, r = 48, 37, 6
    U0 = rng.normal(size=(m, r))
    V0 = rng.normal(size=(n, r))
    # Deliberately non-orthogonal basis, while retaining a well-defined product.
    T = np.diag(np.geomspace(1e-4, 1e4, r))
    U = U0 @ T
    V = V0 @ np.linalg.inv(T).T
    W = U @ V.T
    # Keep the exact same sketch used by the Cholesky factorization.
    S = rng.normal(size=(r + 8, m)) / np.sqrt(r + 8)
    Y = S @ U
    R = np.linalg.cholesky(Y.T @ Y + 1e-12 * np.eye(r)).T
    Us = U @ np.linalg.solve(R, np.eye(r))
    Vs = V @ R.T
    rel_invariance = np.linalg.norm(Us @ Vs.T - W) / np.linalg.norm(W)
    sketch_orth = np.linalg.norm((S @ Us).T @ (S @ Us) - np.eye(r))
    euclidean_orth = np.linalg.norm(Us.T @ Us - np.eye(r))
    # The displayed projector requires both factors to be Euclidean-orthonormal.
    Uq, _ = np.linalg.qr(rng.normal(size=(m, r)), mode="reduced")
    Vq, _ = np.linalg.qr(rng.normal(size=(n, r)), mode="reduced")
    G = rng.normal(size=(m, n))
    H = tangent_projector(Uq, Vq, G)
    residual = G - H
    tangent_id = np.linalg.norm(Uq.T @ residual) + np.linalg.norm(residual @ Vq)
    return {
        "factor_product_relative_error": float(rel_invariance),
        "sketch_gram_orth_error": float(sketch_orth),
        "euclidean_basis_orth_error": float(euclidean_orth),
        "projector_normal_residual": float(tangent_id),
        "projector_residual_norm": float(np.linalg.norm(residual)),
        "initial_u_cond": float(np.linalg.cond(U)),
        "reconditioned_u_cond": float(np.linalg.cond(Us)),
    }


def train(method, seed=123, steps=500, recondition_every=10):
    rng = np.random.default_rng(seed)
    m, n, r = 64, 52, 8
    # Rank-r target plus a perpendicular component, making the tangent projection meaningful.
    A = rng.normal(size=(m, r)); B = rng.normal(size=(n, r))
    target = A @ B.T / np.sqrt(r) + 0.15 * rng.normal(size=(m, n))
    U = rng.normal(size=(m, r))
    V = rng.normal(size=(n, r))
    # Same represented initial matrix across methods, with a well-scaled gauge.
    U, _ = np.linalg.qr(U, mode="reduced")
    V, _ = np.linalg.qr(V, mode="reduced")
    t0 = time.perf_counter()
    history = []
    for step in range(steps):
        W = U @ V.T
        G = (W - target) / (m * n)
        if method == "none":
            H = G
        else:
            # The idea's essential rule: standard projector, never an oblique sketch projector.
            Uq, Vq = exact_recondition(U, V)
            H = tangent_projector(Uq, Vq, G)
            # Factor update using the projected matrix, preserving a simple matched setup.
            U, V = Uq, Vq
        # A stable factor gradient for W=UV^T, with H as the dense tangent update.
        dU = H @ V
        dV = H.T @ U
        lr = 0.03
        U -= lr * dU
        V -= lr * dV
        if method == "sketch" and (step + 1) % recondition_every == 0:
            U, V = sketch_recondition(U, V, rng, oversampling=8)
        elif method == "qr" and (step + 1) % recondition_every == 0:
            U, V = exact_recondition(U, V)
        if step in (0, 9, 99, steps - 1):
            history.append(metrics(U, V, target))
    out = metrics(U, V, target)
    out["history"] = history
    out["seconds"] = time.perf_counter() - t0
    return out


def main():
    result = {"math_check": math_check()}
    for method in ("none", "qr", "sketch"):
        result[method] = train(method)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()