import numpy as np
import scipy
import scipy.linalg

def ensure_close(P, A):
    eps = np.max(np.abs(P)) * 10e-10
    assert(np.max(np.abs(P - A)) <= eps)

def qr_decomposition_np(A):
    Q, R = np.linalg.qr(A)

    ensure_close(A, Q @ R)

    return Q, R

def qr_decomposition_my(A):
    n, m = A.shape
    
    # Gram–Schmidt process
    # https://en.wikipedia.org/wiki/QR_decomposition#Using_the_Gram%E2%80%93Schmidt_process
    
    Q, R = np.zeros((n, n), np.float64), np.zeros((n, m), np.float64)
    Q[:, 0] = A[:, 0] / np.linalg.norm(A[:, 0])
    Q[:, 1] = A[:, 1] - np.dot(Q[:, 0], A[:, 1]) * Q[:, 0]
    Q[:, 1] = Q[:, 1] / np.linalg.norm(Q[:, 1])

    R = Q.transpose() @ A
    R[1, 0] = 0

    ensure_close(A, Q @ R)

    return Q, R

def rq_decomposition_scipy(A):
    R, Q = scipy.linalg.rq(A, mode='economic')

    ensure_close(A, R @ Q)

    return R, Q

def rq_decomposition_my(A):
    A2 = A[::-1].transpose()
    Q2, R2 = qr_decomposition_my(A2)
    Q = Q2.transpose()[::-1]
    R = (R2.transpose())[::-1][:, ::-1]

    ensure_close(A, R @ Q)

    # Mimicking mode='economic' https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.rq.html
    m, n = A.shape
    assert(Q.shape == (n, n))
    assert(R.shape == (m, n))
    k = min(m, n)
    R = R[:, n-k:]
    Q = Q[n-k:, :]
    assert(Q.shape == (k, n))
    assert(R.shape == (m, k))

    ensure_close(A, R @ Q)

    return R, Q

P = np.array([
    [1.95397, -0.0048056, -0.323478, 18145.2],
    [0.0196696, -1.99207, 0.165801, 31412.1],
    [0, 0, 0, 1]
])[:2, :3]

print("P={}".format(P))
assert(P.shape == (2, 3))
# P=[[ 1.95397   -0.0048056 -0.323478 ]
#  [ 0.0196696 -1.99207    0.165801 ]]

# QR decomposition

Q, R = qr_decomposition_np(P)
print("numpy P=QR:")
print("Q={}".format(Q))
print("R={}".format(R))
print("max error in (P-Q@R): {}".format(np.max(np.abs(P - Q @ R))))

print()

Q, R = qr_decomposition_my(P)
print("my P=QR:")
print("Q={}".format(Q))
print("R={}".format(R))
print("max error in (P-Q@R): {}".format(np.max(np.abs(P - Q @ R))))

# numpy P=QR:
# Q=[[-0.99994934 -0.01006597]
#  [-0.01006597  0.99994934]]
# R=[[-1.954069    0.02485747  0.32179266]
#  [ 0.         -1.9919207   0.16904872]]
# max error in (P-Q@R): 1.734723475976807e-18

# my P=QR:
# Q=[[ 0.99994934  0.01006597]
#  [ 0.01006597 -0.99994934]]
# R=[[ 1.954069   -0.02485747 -0.32179266]
#  [ 0.          1.9919207  -0.16904872]]
# max error in (P-Q@R): 1.734723475976807e-18

# RQ decomposition

R, Q = rq_decomposition_scipy(P)
print("scipy P=RQ:")
print("R={}".format(R))
print("Q={}".format(Q))
print("max error in (P-R@Q): {}".format(np.max(np.abs(P - R @ Q))))

print()

R, Q = rq_decomposition_my(P)
print("my P=RQ (via P=QR):")
print("R={}".format(R))
print("Q={}".format(Q))
print("max error in (P-R@Q): {}".format(np.max(np.abs(P - R @ Q))))

# scipy P=RQ:
# R=[[ 1.98056859  0.00281437]
#  [ 0.         -1.99905471]]
# Q=[[ 0.98658421 -0.0038424  -0.16320797]
#  [-0.00983945  0.99650599 -0.0829397 ]]
# max error in (P-R@Q): 2.220446049250313e-16

# my P=RQ (via P=QR):
# R=[[ 1.98056859 -0.00281437]
#  [ 0.          1.99905471]]
# Q=[[ 0.98658421 -0.0038424  -0.16320797]
#  [ 0.00983945 -0.99650599  0.0829397 ]]
# max error in (P-R@Q): 2.220446049250313e-16