using LowRankModels
using LinearAlgebra
using SparseArrays

# testing Gram matrix recovery using the LowRankModels package

# initialize the partial gram matrix for an arrangement of seven spheres in
# which spheres 1 through 5 are mutually tangent, and spheres 3 through 7 are
# also mutually tangent
I = Int64[]
J = Int64[]
values = Float64[]
for i in 1:7
  for j in 1:7
    if (i <= 5 && j <= 5) || (i >= 3 && j >= 3)
      push!(I, i)
      push!(J, j)
      push!(values, i == j ? 1 : -1)
    end
  end
end
gram = sparse(I, J, values)

# in this initial guess, the mutual tangency condition is satisfied for spheres
# 1 through 5
X₀ = sqrt(0.5) * [
  1 0  1  1  1;
  1 0  1 -1 -1;
  1 0 -1  1 -1;
  1 0 -1 -1  1;
  2 -sqrt(6) 0 0 0;
  0.2  0.3 -0.1 -0.2  0.1;
  0.1 -0.2  0.3  0.4 -0.1
]'
Y₀ = diagm([-1, 1, 1, 1, 1]) * X₀

# search parameters
search_params = ProxGradParams(
  1.0;
  max_iter = 100,
  inner_iter = 1,
  abs_tol = 1e-16,
  rel_tol = 1e-9,
  min_stepsize = 0.01
)

# complete gram matrix
model = GLRM(gram, QuadLoss(), ZeroReg(), ZeroReg(), 5, X = X₀, Y = Y₀)
X, Y, history = fit!(model, search_params)