using JuMP
using Gurobi
using CSV, DataFrames
using NLPModelsJuMP, NLPModels
using Random, Distributions
using LinearAlgebra
Random.seed!(3);

H = 10
J = 10
s = Diagonal( ones(H) )
T = 25
Ω = 25
R = 200
q = 3

zᵖ = [4 5 6 7 4 5 6 7 4 5 4 5 6 7 4 5 6 7 4 5   ]
zᶜ = [9 10 11 12 9 10 11 12 9 10 9 10 11 12 9 10 11 12 9 10  ]
zᵐ = ones(H,H)-s

D = fill(0.0,J,T,Ω)
for ω in 1:Ω
    for t in 1:T
        D[:,t,ω] = rand(1:3,J)
    end
end
p = [1/Ω for ω in 1:Ω ];
flexibility = 2



m = Model(Gurobi.Optimizer)
#set_optimizer_attribute(m, "MIPGap", 3)
@variable(m, ψ[1:H], Bin )
@variable(m, χ[1:H , 1:J], Bin  )
@variable(m, 0 ≤ γ[1:J , 1:T , 1:Ω])

@objective(m , Min , 
    sum(T * zᵖ[h] * ψ[h] for h in 1:H ) + 
    sum(zᵐ[h,j] *χ[h,j] for h in 1:H for j in 1:J ) +
    sum(p[ω]*zᶜ[j]*γ[j,t,ω] for j in 1:J for t in 1:T for ω in 1:Ω )
)

c1 = @constraint(m, [h in 1:H , j in 1:J], χ[h,j] ≤ ψ[h]  )
cc = @constraint(m, [h in 1:H , j in 1:J], s[h,j] + ψ[h] ≤ χ[h,j] + 1  )
# extra constraints
c5 = @constraint(m, [h in 1:H], sum(χ[h,j] for j in 1:J) ≤ flexibility )

#######################################################################################
# index generator
#######################################################################################

function index_generator(χᵏ)
    Hᵏ = [ [] for j in 1:J ] # for 1
    Hᵖ = [ [] for j in 1:J ] # for 0
    Hᵀ = [h for h in 1:H]

        for j in 1:J
        
            for h in 1:H
                if χᵏ[h,j] == 1
                    push!(Hᵏ[j] , h)
                end
            end
        end

        for j in 1:J
            Hᵖ[j] = setdiff( Hᵀ , Hᵏ[j] )
        end
   Jᵀ = [j for j in 1:J]
    Jᵏ = [ [] for h in 1:H]
    Jᵛ = [ [] for h in 1:H]
    for h in 1:H
        for j in 1:J
            if χᵏ[h,j] == 1
                push!(Jᵏ[h] , j)
            end
        end
    end
    for h in 1:H    
        Jᵛ[h] = setdiff(Jᵀ , Jᵏ[h])
    end
    return (Hᵏ , Hᵖ , Jᵏ , Jᵛ)
end

function f_zʳ(γᵏ) 
    s = [ [] for t in 1:T, ω in 1:Ω]
    for t in 1:T
        for ω in 1:Ω
            append!(s[t,ω] , [ones(Int(γᵏ[j,t,ω]) )  * zᶜ[j] for j in 1:J] )
            s[t,ω] = reduce(vcat , s[t,ω])
            s[t,ω] = sort(s[t,ω] , rev=true)
            if length(s[t,ω]) < H
                append!( s[t,ω] , [0.0 for i in 1:(H - length(s[t,ω]) ) ] )
            end
        end
    end
    return  s
end

#######################################################################################
# f_dʳ
#######################################################################################
function f_dʳ(zʳ)    
    
    d=[fill(0.0 , H) for t in 1:T, ω in 1:Ω]
    for t in 1:T
        for ω in 1:Ω
            for k in 1:H
                d[t,ω][k]= sum( zʳ[t,ω][l] for l in 1:k)
            end
        end
    end
    return d
end

function solve_sub_prob(χʳ)
    sp =[ Model(Gurobi.Optimizer) for t in 1:T, ω in 1:Ω]
    θ_1 = fill(0.0,T,Ω)
    γ_1 = fill(0.0,J,T,Ω)
    st_1 = fill(0.0,T,Ω)
    for v in 1:T
        for u in 1:Ω
            @variable(sp[v,u], 0 ≤ γ[1:J , [v] , [u] ], Int )
            @variable(sp[v,u], α[1:H , 1:J, [v] , [u] ], Bin )
            @objective(sp[v,u], Min, sum(zᶜ[j] * γ[j,v,u] for j in 1:J) )
            @constraint(sp[v,u], con_23b[h in 1:H , j in 1:J], α[h,j,v,u] ≤ χʳ[h,j] )
            @constraint(sp[v,u], con_23c[h in 1:H], sum( α[h,j,v,u] for j in 1:J) ≤ 1 )
            @constraint(sp[v,u], con_23d[j in 1:J], sum( α[h,j,v,u] for h in 1:H ) + γ[j,v,u] ≥ D[j,v,u] )
            optimize!(sp[v,u])
            θ_1[v,u] = objective_value(sp[v,u])
            γ_1[:,v,u] = round.(value.(γ[:,v,u]) , digits=1)
            st_1[v,u] = solve_time( sp[v,u] )
        end
    end
    
    return θ_1, γ_1 , sum(st_1)
end



#######################################################################################
# initiate
#######################################################################################

function initiate()
    set_r=[]
    set_cost_per=[]
    set_cost_train=[]
    set_βʳ=[]
    set_cost_cas=[]
    set_time_main=[]
    set_time_sub=[]
    set_LB=[]
    set_UB=[]
    set_Θ_min=[]
    set_gap=[]
    set_gap_min=[]
    set_time_tot=[] 
    for r in 1:R
        println("          main   ##############################")
        optimize!(m)
        ψʳ = value.(ψ)
        χʳ = value.(χ)
        βʳ = value.(γ)
        cost_per = sum( T*zᵖ[h]*ψʳ[h] for h in 1:H ) 
        cost_train = sum(χʳ[h,j] * zᵐ[h,j] for h in 1:H for j in 1:J) 
        βʳ = sum( p[ω] * βʳ[j,t,ω] for j in 1:J for t in 1:T for ω in 1:Ω )
        println("          index   ##############################")
        ig = index_generator(χʳ)
        Hʳ = ig[1]
        Hᵖ = ig[2]
        Jʳ = ig[3]
        Jᵖ = ig[4]
        println("          sub   ##############################")
        ss = solve_sub_prob(χʳ)
        θʳ=ss[1]
        γʳ=ss[2]
        solveTime_sub = ss[3]
        solveTime_sub = round(solveTime_sub, digits=2)
        cost_cas = sum( p[ω] * θʳ[t,ω] for t in 1:T for ω in 1:Ω)
        println("          z   ##############################")
        zʳ=f_zʳ(γʳ)
        println("          d   ##############################")
        @show value.(χ)
        l = [length(zʳ[t,ω]) for t in 1:T, ω in 1:Ω]
        dʳ=f_dʳ(zʳ)
        #################################
        println("")
        println("")
        println("          χ   ##############################")
        dt_χʳ = DataFrame(χʳ , :auto)
        @show dt_χʳ
    
        println("")
        println("")
        dt_θʳ = DataFrame(θʳ , :auto)
        #@show dt_θʳ
        push!(set_r , r)
        push!(set_cost_per, cost_per)
        push!(set_cost_train, cost_train)
        push!(set_βʳ, βʳ)
        push!(set_cost_cas, cost_cas)
        println("")
        println("")
        dt1 = DataFrame( r = set_r , cost_per = set_cost_per , cost_train = set_cost_train , βʳ = set_βʳ , cost_cas = set_cost_cas)
        @show dt1
        println("")
        println("")
        push!(set_time_main , round( solve_time(m),digits=2 ) )
        push!(set_time_sub , solveTime_sub)
        time_tot = solve_time(m) + solveTime_sub 
        time_tot = round(time_tot , digits = 2)
        push!(set_time_tot ,  time_tot )
        LB = objective_value(m)
        push!(set_LB , LB)
        UB = cost_per + cost_train + cost_cas
        push!(set_UB , UB)
        Θ_min = minimum(set_UB)
        push!(set_Θ_min , Θ_min)
        gap = (UB - LB) / UB
        push!(set_gap , gap)
        gap_min = minimum(set_gap)
        push!(set_gap_min , gap_min)
        dt2 = DataFrame( r = set_r , 
            time_main = set_time_main , time_sub = set_time_sub , time_tot = set_time_tot, 
            time_grand_tot = sum(set_time_tot),
            LB = set_LB , UB = set_UB , Θ_min = set_Θ_min , gap = set_gap , gap_min = set_gap_min )
        @show dt2
        
        if r ≥ q
            @variable(m , α[1:H , 1:J , 1:T , 1:Ω], Bin);
            c2 = @constraint(m, [h in 1:H , j in 1:J , t in 1:T , ω in 1:Ω], α[h,j,t,ω] ≤ χ[h,j] )
            c3 = @constraint(m, [h in 1:H , t in 1:T , ω in 1:Ω], sum( α[h,j,t,ω] for j in 1:J) ≤ 1 )
            c4 = @constraint(m, [j in 1:J , t in 1:T , ω in 1:Ω], sum( α[h,j,t,ω] for h in 1:H ) + γ[j,t,ω] ≥ D[j,t,ω] )
            optimize!(m)
            @show solve_time(m)
            @show solution_summary(m)
             ψʳ = value.(ψ)
             χʳ = value.(χ)
             γʳ = value.(γ)
            cost_per = sum(T * zᵖ[h] * ψʳ[h] for h in 1:H ) 
            cost_train = sum(zᵐ[h,j] *χʳ[h,j] for h in 1:H for j in 1:J ) 
            cost_cas = sum( p[ω] * zᶜ[j] * γʳ[j,t,ω]  for j in 1:J for t in 1:T for ω in 1:Ω )
            @show objective_value(m)
            @show cost_per
            @show cost_train
            @show cost_cas
            @show main_time = sum(set_time_main)
            @show sub_time = sum(set_time_sub)
            @show Org_time = solve_time(m)
            @show time_tot + solve_time(m)
            @show num_variables(m)
            noConLess = num_constraints(m, AffExpr, MOI.LessThan{Float64})
            noConMore = num_constraints(m, AffExpr, MOI.GreaterThan{Float64})
            noConEqual = num_constraints(m, AffExpr, MOI.EqualTo{Float64}) 
            noConTot = noConLess + noConMore + noConEqual
            @show noConTot
            df_result = DataFrame( r= r, gap_min = gap_min , no_Var = num_variables(m) , no_Con = noConTot , 
                main_time = main_time , sub_time = sub_time , Org_time = solve_time(m) , tot_time =  main_time + sub_time + Org_time  )
            @show df_result
            break
        end
        
        
        if gap_min ≤ 0.001
            println("##################################")
            println("########## optimality      #######")
            @show objective_value(m)
            break
        end
         b = @variable(m, [1:H ], Bin)
        c = @variable(m, [1:T , 1:Ω])
        @constraint(m, [t in 1:T , ω in 1:Ω], c[t,ω] ≥ 0 )
        a = @variable(m, [1:H], Bin)
        
       
        c_29a = @constraint(m, [t in 1:T , ω in 1:Ω], sum(zᶜ[j]*γ[j,t,ω] for j in 1:J ) ≥ θʳ[t,ω] - c[t,ω] - 
            θʳ[t,ω] * sum( (1 - χ[h,j] ) for j in 1:J for h in Hʳ[j] )   )
        c_29b = @constraint(m, [h in 1:H], J * b[h] ≥ sum(χ[h,j] for j in Jᵖ[h]) )  
        c_29c = @constraint(m, [h in 1:H], b[h] ≤ sum( χ[h,j] for j in Jᵖ[h] ) ) 
        
        c_29d = @constraint(m, [t in 1:T , ω in 1:Ω], c[t,ω] ≤ sum( a[k]*dʳ[t,ω][k] for k in 1:H ) )
        c_29e = @constraint(m, [t in 1:T , ω in 1:Ω], c[t,ω] ≤ sum( zʳ[t,ω][i] for i in 1:H ) )
        c_29f = @constraint(m, [k in 1:H], sum( k * a[k] for k in 1:H) == sum(b[h] for h in 1:H) )
        c_29g = @constraint(m, sum( a[k] for k in 1:H ) == 1 )
    end 
end

initiate()


1