using Plots
using JuMP
using Juniper, SCS
using LinearAlgebra
using Test

m = 20     # number of families
n = 3      # number of schools
ndims = 2  # number of Euclidean dimensions in which points are embedded

# Generate points approximately on the unit circle
ps = map(1:m) do _
    p = randn(ndims)
    normalize!(p)
    p += 0.1*randn(ndims)
    return p
end

pl = plot(size=(600,600), xlim=(-1.5, 1.5), ylim=(-1.5, 1.5))
scatter!(pl, [tuple(p...) for p in ps], m=:square, c=:black, msc=:auto, label="families")

# Set SCS as the solver for the relaxation problems. 
nl_solver = optimizer_with_attributes(SCS.Optimizer, "verbose"=>0)

# Wrap SCS in Juniper to enable integrality constraints.
mdl = Model(
    optimizer_with_attributes(
        Juniper.Optimizer,
        "nl_solver"=>nl_solver,
        "atol"=>1e-3,
    )
)

set_silent(mdl)

# Set the value of big M, which needs to be an upper bound on d[i, j]. 
M = maximum(norm(ps[i] - ps[k], 1) for i in 1:m for k in i+1:m)

@variable(mdl, t)
@variable(mdl, d[1:n, 1:m])
@variable(mdl, s[1:n, 1:m], Bin)
@variable(mdl, x[1:ndims, 1:n])

for j in 1:m
    for i in 1:n  
        @constraint(mdl, [d[i, j]; x[:, i] .- ps[j]] in MOI.SecondOrderCone(ndims+1))  
        
        # To use one norm (taxicab distance) instead
        # @constraint(mdl, [d[i, j]; x[:, i] .- ps[j]] in MOI.NormOneCone(ndims+1))   

        # s[i, j] = 1 ⟹ d[i, j] ≤ t
        @constraint(mdl, t ≥ d[i, j] - M * (1 - s[i, j]))
    end
    @constraint(mdl, sum(s[i, j] for i in 1:n) ≥ 1)
end

@objective(mdl, Min, t)

mdl

optimize!(mdl)
solution_summary(mdl)

xs = [value.(x[:, i]) for i in 1:n]

attend_school = map(ps) do p
    argmin(i -> norm(p - xs[i]), 1:n)
end

@test attend_school == map(argmin, eachcol(value.(d))) == map(argmax, eachcol(value.(s)))

pl = plot(size=(600,600), xlim=(-1.5, 1.5), ylim=(-1.5, 1.5))

for (j, p) in enumerate(ps)
    plot!(
        pl,
        [tuple(p...), tuple(xs[attend_school[j]]...)],
        label= j==1 ? "local school" : nothing,
        c = "navy",
        ls = :dash
    )
end

scatter!(pl, [tuple(p...) for p in ps], m=:square, c=:black, msc=:auto, label="families")
scatter!(pl, [tuple(x...) for x in xs], c=:crimson, msc=:auto, ms=7, label="schools")

pl

using Clustering

res = kmeans(hcat(ps...), n)
xs_kmeans = map(Vector{Float64}, eachcol(res.centers))

attend_school_kmeans = attend_school = map(ps) do p
    argmin(i -> norm(p - xs_kmeans[i]), 1:n)
end

pl = plot(size=(600,600), xlim=(-1.5, 1.5), ylim=(-1.5, 1.5))

for (j, p) in enumerate(ps)
    plot!(
        pl,
        [tuple(p...), tuple(xs_kmeans[attend_school_kmeans[j]]...)],
        label= j==1 ? "local school" : nothing,
        c = "navy",
        ls = :dash
    )
end

scatter!(pl, [tuple(p...) for p in ps], m=:square, c=:black, msc=:auto, label="families")
scatter!(pl, [tuple(x...) for x in xs_kmeans], c=:crimson, msc=:auto, ms=7, label="schools (k means)")

pl