include("Engine.jl")

using SparseArrays

# --- construct the nucleus spheres ---

println("--- Nucleus spheres ---\n")

# initialize the partial gram matrix for the circumscribing and nucleus spheres
J = Int64[]
K = Int64[]
values = BigFloat[]
for n in 1:3
  push!(J, n)
  push!(K, n)
  push!(values, 1)
  if n > 1
    append!(J, [1, n])
    append!(K, [n, 1])
    append!(values, [1, 1])
  end
end
gram_nuc = sparse(J, K, values)

# make an initial guess
guess_nuc = hcat(
  Engine.sphere(BigFloat[0, 0, 0], BigFloat(15)),
  Engine.sphere(BigFloat[0, 0, -10], BigFloat(5)),
  Engine.sphere(BigFloat[0, 0, 11], BigFloat(3)),
)
frozen_nuc = [CartesianIndex(4, k) for k in 1:3]

# complete the gram matrix using Newton's method with backtracking
L_nuc, success_nuc, history_nuc = Engine.realize_gram(gram_nuc, guess_nuc, frozen_nuc)
completed_gram_nuc = L_nuc'*Engine.Q*L_nuc
println("Completed Gram matrix:\n")
display(completed_gram_nuc)
if success_nuc
  println("\nTarget accuracy achieved!")
else
  println("\nFailed to reach target accuracy")
end
println("Steps: ", size(history_nuc.scaled_loss, 1))
println("Loss: ", history_nuc.scaled_loss[end], "\n")

# --- construct the chain of spheres ---

# initialize the partial gram matrix for the chain of spheres
J = Int64[]
K = Int64[]
values = BigFloat[]
for a in 4:9
  push!(J, a)
  push!(K, a)
  push!(values, 1)
  
  # each chain sphere is internally tangent to the circumscribing sphere
  append!(J, [a, 1])
  append!(K, [1, a])
  append!(values, [1, 1])
  
  # each chain sphere is externally tangent to the nucleus spheres
  for n in 2:3
    append!(J, [a, n])
    append!(K, [n, a])
    append!(values, [-1, -1])
  end
  
  # each chain sphere is externally tangent to the next sphere in the chain
  #=
  a_next = 4 + mod(a-3, 6)
  append!(J, [a, a_next])
  append!(K, [a_next, a])
  append!(values, [-1, -1])
  =#
end
gram_chain = sparse(J, K, values)

if success_nuc
  println("--- Chain spheres ---\n")
  
  # make an initial guess, with the circumscribing and nucleus spheres included
  # as frozen elements
  guess_chain = hcat(
    L_nuc,
    (
      Engine.sphere(10*BigFloat[cos(k*π/3), sin(k*π/3), 0], BigFloat(2.5))
      for k in 1:6
    )...
  )
  frozen_chain = [CartesianIndex(j, k) for k in 1:3 for j in 1:5]
  
  # complete the gram matrix using Newton's method with backtracking
  L_chain, success_chain, history_chain = Engine.realize_gram(gram_chain, guess_chain, frozen_chain)
  completed_gram_chain = L_chain'*Engine.Q*L_chain
  println("Completed Gram matrix:\n")
  display(completed_gram_chain)
  if success_chain
    println("\nTarget accuracy achieved!")
  else
    println("\nFailed to reach target accuracy")
  end
  println("Steps: ", size(history_chain.scaled_loss, 1))
  println("Loss: ", history_chain.scaled_loss[end], "\n")
end