using LinearAlgebra, Plots, Random, SparseArrays

"""
conls
=====

Solve a constrained least squares problem 

    conls(A,b,B,d) -> x 

solves 

   minimize     ||b - Ax||_2
   subject to   Bx = d

when B is p-by-n, p < n, and rank(B) = p, by using the QR factorization
of B' to find a basis for the nullspace of B.
"""
function conls(A,b,B,d)
    m,n = size(A)
    p = size(B,1)
    
    @assert size(B,2) == n "incompatible dimensions for A and B"
    
    # Step 1, QR of B' to get the nullspace
    qrf = qr(B')
    R = qrf.R
    @assert all(diag(R) != 0) "B is not full-rank"
    
    Q = qrf.Q*Matrix(I,n,n) # this forms the full matrix Q. 
    Q1 = view(Q,:,1:p)     # view extracts a submatrix
    Q2 = view(Q,:,(p+1):n) 
    
    # form a solution
    xc = Q1*(R'\d)
    # solve the new LS problem
    y = (A*Q2)\(b-A*(xc))
    # recombine the solution
    x =xc + Q2*y
    return x
end

conls

teams = ["duke","miami","unc","uva","vt"]
data = [ # team 1 team 2, team 1 pts, team 2 pts
    1 2  7 52 # duke played Miami and lost 7 to 52 
    1 3 21 24 # duke played unc and lost 21 to 24 
    1 4  7 38
    1 5  0 45
    2 3 34 16
    2 4 25 17
    2 5 27  7
    3 4  7  5
    3 5  3 30
    4 5 14 52
]
ngames = size(data,1)
nteams = length(teams)

G = zeros(ngames, nteams)
p = zeros(ngames, 1)

for g=1:ngames
    i = data[g,1]
    j = data[g,2]
    Pi = data[g,3]
    Pj = data[g,4]
  
    G[g,i] = 1
    G[g,j] = -1
    p[g] = Pi - Pj
end
s1 = conls(G,p,ones(1,nteams),[100.0])

5×1 Matrix{Float64}:
 -4.800000000000004
 38.2
 11.999999999999996
 16.599999999999998
 38.0

sum(s1)

1.0000000000000107

G

10×5 Matrix{Float64}:
 1.0  -1.0   0.0   0.0   0.0
 1.0   0.0  -1.0   0.0   0.0
 1.0   0.0   0.0  -1.0   0.0
 1.0   0.0   0.0   0.0  -1.0
 0.0   1.0  -1.0   0.0   0.0
 0.0   1.0   0.0  -1.0   0.0
 0.0   1.0   0.0   0.0  -1.0
 0.0   0.0   1.0  -1.0   0.0
 0.0   0.0   1.0   0.0  -1.0
 0.0   0.0   0.0   1.0  -1.0

for i=1:nteams
    println(teams[i], " ", s1[i])
end

duke -4.800000000000004
miami 38.2
unc 11.999999999999996
uva 16.599999999999998
vt 38.0

Look at the residual

G*s1 - p

10×1 Matrix{Float64}:
   1.999999999999993
 -13.800000000000004
   9.599999999999994
   2.1999999999999957
   8.200000000000003
  13.600000000000001
 -19.799999999999997
  -6.600000000000001
   1.0
  16.6

sum(s1)

1.0000000000000107

"""
conls_lagrange
=====

Solve a constrained least squares problem 

    conls(A,b,B,d) -> x, lam

solves 

   minimize     ||b - Ax||_2
   subject to   Bx = d

when B is p-by-n, p < n, and rank(B) = p, by using the augmented
system of normal equations

   [ A'*A  B' ] [  x  ] = [ A'*b]
   [ B     0  ] [ lam ] = [  d  ]

and returns both x and the lagrange multipliers
"""
function conls_lagrange(A,b,B,d)
    m,n = size(A)
    p = size(B,1)
    
    @assert size(B,2) == n "incompatible dimensions for A and B"
    
    # Step 1, form the block system
    M = [A'*A B'; B zeros(p,p) ]

    # Step 2, solve
    z = M\[A'*b; d]

    # Step 3, extract 
    x = z[1:n]
    lam = z[n+1:end]
    
    return x, lam
end

conls_lagrange

s2,lam = conls_lagrange(G,p,ones(1,nteams),[100.0])

([-4.800000000000001, 38.2, 11.999999999999998, 16.6, 38.0], [0.0])

s1 - s2

5×1 Matrix{Float64}:
 -3.552713678800501e-15
  0.0
 -1.7763568394002505e-15
 -3.552713678800501e-15
  0.0