using Test, LinearAlgebra

function test_naive()
    # X represents a square with edges parallel to coordinate axes
    # and center at (-1, -1),
    # Y is a square rotated by 45 degrees centered at (1, 1).
    s = 1/sqrt(2)
    X = [-1-s -1-s; -1+s -1-s; -1+s -1+s; -1-s -1+s]'
    Y = [1 0; 2 1; 1 2; 0 1]'
    # The actual translation vector and rotation matrix.
    b = [1; 1+2*s]
    Q = [s -s; s s]
    @test isapprox([Q, b], collect(naive(X, Y)); atol = sqrt(eps()))
end

# the actual 'naive' function ...
function naive(X, Y)
    # solve the corresponding linear system for the naive approach
    Q = Y/[X; ones(1, size(X, 2))]
    # extract b
    b = Q[:, end]
    # find the QR decomposition of Q'
    F = qr(Q[:, 1:(end-1)])
    # not every Q is good enough ...
    Q = F.Q * sign.(Diagonal(F.R))
    return (Q, b)
end

function kabsch(X, Y)
    # averages of X and Y datasets
    xbar = sum(X; dims=2)/size(X, 2)
    ybar = sum(Y; dims=2)/size(Y, 2)
    # centered datasets
    Xprime = X .- xbar
    Yprime = Y .- ybar
    # covariance matrix ...
    C = Yprime*Xprime'
    # ... and its SVD
    U, _, V = svd(C)
    # to replace S we can equivalently change the sign 
    # of the last column of U (or V)
    # U[:, end] = sign(det(C))*U[:, end]
    # evaluate Q ...
    Q = U*V'
    # ... and b
    b = ybar - Q*xbar
    
    return (Q, b)
end

function test_kabsch()
    # X represents a square with edges parallel to coordinate axes
    # and center at (-1, -1),
    # Y is a square rotated by 45 degrees centered at (1, 1).
    s = 1/sqrt(2)
    X = [-1-s -1-s; -1+s -1-s; -1+s -1+s; -1-s -1+s]'
    Y = [1 0; 2 1; 1 2; 0 1]'
    # The actual translation vector and rotation matrix.
    b = [1; 1+2*s]
    Q = [s -s; s s]
    @test isapprox([Q, b], collect(kabsch(X, Y)); atol = sqrt(eps()))
end