using JLD
using JuMP
using SCS

using GroupRings
using PropertyT

using ArgParse

immutable ProblemData{T}
   name::String
   Us::Vector{SparseMatrixCSC{Float64,Int}}
   Ps::Vector{Array{JuMP.Variable,2}}
   cnstr::Vector{T}
   laplacian::SparseVector{Float64}
   laplacianSq::SparseVector{Float64}
   dims::Vector{Int}
end

function sparsify!{T}(U::AbstractArray{T}, eps=eps(T))
#     n = rank(U)
    U[abs.(U) .< eps] = zero(T)
#     @assert rank(U) == n
    return sparse(U)
end

sparsify{T}(U::AbstractArray{T}, eps=eps(T)) = sparsify!(deepcopy(U), eps)

small_to_zero!{T}(A::AbstractArray{T}, eps=eps(T)) = A[abs(A) .< eps] = zero(T)

function orbit_spvector(vect::AbstractVector, orbits)
    orb_vector = spzeros(length(orbits))

    for (i,o) in enumerate(orbits)
        k = vect[collect(o)]
        val = k[1]
        @assert all(k .== val)
        orb_vector[i] = val
    end

    return orb_vector
end

function orbit_constraint(cnstrs::Vector{Vector{Vector{Int64}}}, n)
    result = spzeros(n,n)
    for cnstr in cnstrs
        for p in cnstr
            result[p[1],p[2]] += 1.0
        end
    end
    return 1/length(cnstrs)*result
end

function init_model(Uπs)
    m = JuMP.Model();
    l = size(Uπs,1)
    P = Vector{Array{JuMP.Variable,2}}(l)

    for k in 1:l
        s = size(Uπs[k],2)
        P[k] = JuMP.@variable(m, P[k][i=1:s, j=1:s])
        JuMP.@SDconstraint(m, P[k] >= 0.0)
    end

    JuMP.@variable(m, λ >= 0.0)
    JuMP.@objective(m, Max, λ)
    return m, P
end

function init_ProblemData(name::String)
   splap = load(joinpath(name, "delta.jld"), "delta");
   pm = load(joinpath(name, "pm.jld"), "pm");
   cnstr = PropertyT.constraints_from_pm(pm);
   splap² = GroupRings.mul(splap, splap, pm);

   Uπs = load(joinpath(name, "U_pis.jld"), "Uπs");
   #dimensions of the corresponding πs:
   dims = [1,1,2,3,3,4,4,8,6,6,6,6,4,4,8,1,1,2,3,3]
   # dims load(joinpath(name, "U_pis.jld"), "dims")
   Uπs = sparsify.(Uπs);

   m, P = init_model(Uπs);

   orbits = load(joinpath(name, "orbits.jld"), "orbits");
   n = size(Uπs[1],1)
   orb_spcnstrm = [orbit_constraint(cnstr[collect(orb)], n) for orb in orbits]
   orb_splap = orbit_spvector(splap, orbits)
   orb_splap² = orbit_spvector(splap², orbits)

   orb_SOutF4data = ProblemData(name, Uπs, P, orb_spcnstrm, orb_splap, orb_splap², dims);

   return m, orb_SOutF4data
end

function transform{T}(U::AbstractArray{T,2}, V::AbstractArray{T,2}, eps=eps(T))
    w = U'*V*U
    sparsify!(w, eps)
    dropzeros!(w)
    return w
end

A(data::ProblemData, π, t) = data.dims[π]*transform(data.Us[π], data.cnstr[t])

function constrLHS(m::JuMP.Model, data::ProblemData, t)
    l = endof(data.Us)
    lhs = @expression(m, sum(vecdot(A(data, π, t), data.Ps[π]) for π in 1:l))
    return lhs
end

function addconstraints!(m::JuMP.Model, data::ProblemData, l::Int=length(data.cnstr); var::Symbol = :λ)
    λ = getvariable(m, var)
    for t in 1:l
        d, d² = data.laplacian[t], data.laplacianSq[t]
        lhs = constrLHS(m, data, t)
        if lhs == zero(lhs)
            if d == 0 && d² == 0
                info("Detected empty constraint")
                continue
            else
                warn("Adding unsatisfiable constraint!")
            end
        end
        JuMP.@constraint(m, lhs == d² - λ*d)
    end
end

function reconstructP(m::JuMP.Model, data::ProblemData)
    computedPs = [getvalue(P) for P in data.Ps]
    return sum(data.dims[π]*data.Us[π]*computedPs[π]*data.Us[π]' for π in 1:endof(data.Ps))
end

function cpuinfo_physicalcores()
    maxcore = -1
    for line in eachline("/proc/cpuinfo")
        if startswith(line, "core id")
            maxcore = max(maxcore, parse(Int, split(line, ':')[2]))
        end
    end
    maxcore < 0 && error("failure to read core ids from /proc/cpuinfo")
    return maxcore + 1
end

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "--tol"
            help = "set numerical tolerance for the SDP solver (default: 1e-5)"
            arg_type = Float64
            default = 1e-5
        "--iterations"
            help = "set maximal number of iterations for the SDP solver (default: 20000)"
            arg_type = Int
            default = 20000
        "--upper-bound"
            help = "Set an upper bound for the spectral gap (default: Inf)"
            arg_type = Float64
            default = Inf
        "--cpus"
            help = "Set number of cpus used by solver (default: auto)"
            arg_type = Int
            required = false
      #   "-N"
      #       help = "Consider automorphisms of free group on N generators (default: N=3)"
      #       arg_type = Int
      #       default = 2
    end

    return parse_args(s)
end

function create_SDP_problem(name::String; upper_bound=Inf)
   info(PropertyT.logger, "Loading data....")
   t = @timed SDP_problem, orb_data = init_ProblemData(name);
   info(PropertyT.logger, PropertyT.timed_msg(t))

   if upper_bound < Inf
      λ = JuMP.getvariable(SDP_problem, :λ)
      JuMP.@constraint(SDP_problem, λ <= upper_bound)
   end

   info(PropertyT.logger, "Adding constraints... ")
   t = @timed addconstraints!(SDP_problem, orb_data)
   info(PropertyT.logger, PropertyT.timed_msg(t))

   return SDP_problem, orb_data
end

function λandP(m::JuMP.Model, data::ProblemData)
   info(PropertyT.logger, "Solving SDP problem...")
   varλ = JuMP.getvariable(m, :λ)
   varP = data.Ps
   λ, P = PropertyT.λandP(data.name, m, varλ, varP)

   recP = reconstructP(m, data)
   fname = PropertyT.λSDPfilenames(data.name)[2]
   save(fname, "origP", P, "P", recP)
   return λ, recP
end

function main()
   name = "SOutF4_E4"

   if !isdir(name)
      throw("Create dir with all the required orbit data first!")
   end

   logger = PropertyT.setup_logging(name)

   parsed_args = parse_commandline()

   if parsed_args["cpus"] != nothing
      if parsed_args["cpus"] > cpuinfo_physicalcores()
           warn("Number of specified cores exceeds the physical core cound. Performance will suffer.")
      end
      Blas.set_num_threads(parsed_args["cpus"])
   end

   tol = parsed_args["tol"]
   iterations = parsed_args["iterations"]
   upper_bound = parsed_args["upper-bound"]

   solver = SCS.SCSSolver(eps=tol, max_iters=iterations, verbose=true, linearsolver=SCS.Indirect)

   fnames = PropertyT.λSDPfilenames(name)
   if all(isfile.(fnames)) && false
      λ, P = PropertyT.λandP(name)
   else
      info(logger, "Creating SDP problem...")
      SDP_problem, orb_data = create_SDP_problem(name, upper_bound=upper_bound)
      JuMP.setsolver(SDP_problem, solver)
      λ, P = λandP(SDP_problem, orb_data)
   end

   info(PropertyT.logger, "λ = $λ")
   info(PropertyT.logger, "sum(P) = $(sum(P))")
   info(PropertyT.logger, "maximum(P) = $(maximum(P))")
   info(PropertyT.logger, "minimum(P) = $(minimum(P))")

end

main()