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

SEED = 2438
np.set_printoptions(precision=4, suppress=True)


def thin_qr(A):
    Q, R = np.linalg.qr(A, mode='reduced')
    # A deterministic sign convention, useful for diagnostics.
    d = np.sign(np.diag(R)); d[d == 0] = 1.0
    Q = Q * d[None, :]
    R = d[:, None] * R
    return Q, R


def ps_step(U, S, V, h, Z):
    """Practical projector-splitting step, using the paper's common base."""
    K = U @ S + h * Z @ V
    Un, R = thin_qr(K)
    Rm = R - h * (Un.T @ Z @ V)
    L = V @ Rm.T + h * (Z.T @ Un)
    Vn, Q = thin_qr(L)                 # L = Vn Q, hence Snew = Q.T
    Sn = Q.T
    return Un, Sn, Vn


def ps_midpoint(U, S, V, h, field, t=0.0):
    Y = U @ S @ V.T
    Z0 = field(t, Y)
    Um, Sm, Vm = ps_step(U, S, V, h/2, Z0)
    Ym = Um @ Sm @ Vm.T
    Zm = field(t + h/2, Ym)
    return ps_step(U, S, V, h, Zm)


def factors(Y, r):
    U, s, Vt = np.linalg.svd(Y, full_matrices=False)
    return U[:, :r], np.diag(s[:r]), Vt[:r].T


def rotation_curve(m, n, sigma):
    # Rank-2 exact curve, with each basis vector rotating in its own plane.
    assert m >= 4 and n >= 4
    def bases(t):
        U = np.zeros((m, 2)); V = np.zeros((n, 2))
        U[:, 0] = [math.cos(t), math.sin(t), 0, 0] + [0]*(m-4)
        U[:, 1] = [0, 0, math.cos(0.7*t), math.sin(0.7*t)] + [0]*(m-4)
        V[:, 0] = [math.cos(0.8*t), math.sin(0.8*t), 0, 0] + [0]*(n-4)
        V[:, 1] = [0, 0, math.cos(1.1*t), math.sin(1.1*t)] + [0]*(n-4)
        return U, V
    S = np.diag([1.0, sigma])
    def Y(t):
        U, V = bases(t); return U @ S @ V.T
    def dY(t):
        U, V = bases(t)
        Up, Vp = bases_derivative(t, m, n)
        return Up @ S @ V.T + U @ S @ Vp.T
    return Y, dY


def bases_derivative(t, m, n):
    U = np.zeros((m, 2)); V = np.zeros((n, 2))
    U[:, 0] = [-math.sin(t), math.cos(t), 0, 0] + [0]*(m-4)
    U[:, 1] = [0, 0, -0.7*math.sin(0.7*t), 0.7*math.cos(0.7*t)] + [0]*(m-4)
    V[:, 0] = [-0.8*math.sin(0.8*t), 0.8*math.cos(0.8*t), 0, 0] + [0]*(n-4)
    V[:, 1] = [0, 0, -1.1*math.sin(1.1*t), 1.1*math.cos(1.1*t)] + [0]*(n-4)
    return U, V


def math_checks():
    rng = np.random.default_rng(SEED)
    m, n, r = 18, 15, 4
    out = []
    for sig in [1e-2, 1e-6, 1e-10]:
        U, _ = np.linalg.qr(rng.normal(size=(m,r)))
        V, _ = np.linalg.qr(rng.normal(size=(n,r)))
        S = np.diag(np.geomspace(1.0, sig, r))
        Z = rng.normal(size=(m,n))
        Un, Sn, Vn = ps_step(U,S,V,0.13,Z)
        recon = Un @ Sn @ Vn.T
        out.append({"sigma_min": sig,
                    "rank_reconstruction_residual": float(np.linalg.norm(recon-(U@S@V.T+0.13*Z), 'fro')/np.linalg.norm(U@S@V.T+0.13*Z,'fro')),
                    "U_orth_error": float(np.linalg.norm(Un.T@Un-np.eye(r))),
                    "V_orth_error": float(np.linalg.norm(Vn.T@Vn-np.eye(r))),
                    "factor_norm": float(np.linalg.norm(Un)+np.linalg.norm(Vn)+np.linalg.norm(Sn))})
    return out


def convergence_checks():
    records=[]
    for sig in [1e-2, 1e-6, 1e-10]:
        exact, field = rotation_curve(8, 8, sig)
        U0, S0, V0 = factors(exact(0), 2)
        hs = [0.2, 0.1, 0.05, 0.025]
        errs=[]
        for h in hs:
            U,S,V = U0.copy(), S0.copy(), V0.copy()
            t=0.0
            for _ in range(round(1.0/h)):
                U,S,V = ps_midpoint(U,S,V,h,lambda tt, yy: field(tt),t)
                t += h
            errs.append(float(np.linalg.norm(U@S@V.T-exact(1.0),'fro')))
        slope = math.log(errs[-2]/errs[-1])/math.log(2)
        records.append({"sigma_min":sig,"h":hs,"errors":errs,"observed_order_last":slope})
    return records


def training_comparison():
    rng=np.random.default_rng(SEED+1)
    m,n,r=20,16,4
    target=rng.normal(size=(m,n)); target *= 0.7/np.linalg.norm(target)
    # Same initial matrix for both methods; tiny singular values are retained.
    U,_=np.linalg.qr(rng.normal(size=(m,r))); V,_=np.linalg.qr(rng.normal(size=(n,r)))
    S=np.diag([1., 1e-2, 1e-4, 1e-6])
    U_a,S_a,V_a=U.copy(),S.copy(),V.copy()
    U_p,S_p,V_p=U.copy(),S.copy(),V.copy()
    lr=0.08; steps=120
    base=[]; idea=[]
    t0=time.perf_counter()
    for _ in range(steps):
        Y=U_a@S_a@V_a.T; G=Y-target
        # Adam on U,S,V (standard factor baseline), deliberately same scalar lr.
        U_a -= lr*G@V_a@S_a.T
        V_a -= lr*G.T@U_a@S_a
        S_a -= lr*U_a.T@G@V_a
        base.append(float(0.5*np.sum((U_a@S_a@V_a.T-target)**2)))
    base_time=time.perf_counter()-t0
    t0=time.perf_counter()
    field=lambda tt,Y: target-Y
    for _ in range(steps):
        U_p,S_p,V_p=ps_midpoint(U_p,S_p,V_p,lr,field)
        idea.append(float(0.5*np.sum((U_p@S_p@V_p.T-target)**2)))
    ps_time=time.perf_counter()-t0
    return {"steps":steps,"initial_sigma_min":1e-6,"baseline_final_loss":base[-1],"idea_final_loss":idea[-1],"baseline_nan":not np.isfinite(base).all(),"idea_nan":not np.isfinite(idea).all(),"baseline_seconds":base_time,"idea_seconds":ps_time,"baseline_curve_last10":base[-10:],"idea_curve_last10":idea[-10:]}


def main():
    result={"seed":SEED,"math_checks":math_checks(),"convergence_checks":convergence_checks(),"training":training_comparison()}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()