# Models for 2.A
using JuMP

function maxBonus(end_time, start_time, bonus, conflicts)
    # end_time: the end time of each task
    # start_time: the start time of each task 
    # bonus: the bonus of each task
    # conflicts: conflicts[i] is the set of tasks whose time conflicts with task i
    #
    # returned: (x,obj), where 
    # 1. x[i] indicates whether task i is selected in the optimal solution;
    # 2. obj is the maximum total bonus 
    
    num = length(bonus) # number of tasks
    
    m = Model()
    
    @variable(m, x[1:num], Bin) # whether task i is selected to do
    for i=1:num
        for k in conflicts[i]
            # if two tasks conflicts, at most one of them is selected
            @constraint(m, x[i] + x[k] <= 1)
        end
    end
    @objective(m, Max, sum{bonus[i]*x[i], i=1:num}) # maximize the total bonus
    status = solve(m)
    
    return (getvalue(x), getobjectivevalue(m))
end

function minNumProcessors(end_time, start_time, conflicts)
    # end_time: the end time of each task
    # start_time: the start time of each task 
    # conflicts: conflicts[i] is the set of tasks whose time conflicts with task i
    #
    # returned: (x,y,obj), where 
    # 1. x[i,j] indicates whether task i is scheduled on processor j
    # 2. y[j] indicates whether processor j is used;
    # 3. obj is the maximum total bonus 
    
    num = length(bonus) # num of tasks
    
    m = Model()
    
    @variable(m, x[1:num, 1:num], Bin) # whether task i is scheduled on processor j
    @variable(m, y[1:num], Bin) # whether processor j is used
    
    for i=1:num # the set partition constraint
        @constraint(m, sum{x[i,j], j=1:num} == 1)
    end
    
    for i=1:num
        for j=1:num
            for k in conflicts[i] 
                # if two tasks conflicts, they cannot be scheduled on the same task
                @constraint(m, x[i, j] + x[k, j] <= 1)
            end
        end
    end
        
    for j=1:num # if processor j is assigned at least one task, then it is used
        @constraint(m, sum{x[i,j], i=1:num} <= num*y[j])
    end
    
    @objective(m, Min, sum{y[j], j=1:num})
    status = solve(m)
    
    return (getvalue(x), getvalue(y), getobjectivevalue(m))
end
;

# Models for 2.B
using JuMP

function minTime(workload, num_proc)
    # workload: workload of each task
    # num_proc: the number of processors available
    #
    # returned: (x,z) where
    # 1. x[i,j] indicates whether task i is scheduled on processor j
    # 2. z is the earliest completion time 
    
    num = length(workload) # number of tasks
    
    m = Model()
    
    @variable(m, x[1:num, 1:num_proc], Bin) # whether task i is scheduled on processor j
    @variable(m, z) # the earliest completion time 
    
    for i=1:num # set partition constraint
        @constraint(m, sum{x[i,j], j=1:num_proc} == 1)
    end
    for j=1:num_proc # z >= total workloads of tasks scheduled processor j
        @constraint(m, sum{workload[i]*x[i,j], i=1:num} <= z)
    end
    @objective(m, Min, z)
    status = solve(m)
    
    return (getvalue(x), getvalue(z))
end

function minResource(workload, T)
    # workload: workload of each task
    # T: the time limit (each task needs to be completed by T)
    #
    # returned: (x,y,obj) where
    # 1. x[i,j] indicates whether task i is scheduled on processor j
    # 2. y[i] indicates whether processor j is used
    # 3. obj: the number of used processors
    
    num = length(workload)
    
    m = Model()
    
    @variable(m, x[1:num, 1:num], Bin)
    @variable(m, y[1:num], Bin)
    for i=1:num # set partition constraint
        @constraint(m, sum{x[i,j], j=1:num} == 1)
    end
    for j=1:num 
        # for each processor, total workloads of tasks assigned to it cannot exceed T
        # and it is used if at least one task is assigned to it
        @constraint(m, sum{workload[i]*x[i,j], i=1:num} <= T*y[j])
    end
    @objective(m, Min, sum{y[j], j=1:num}) # minimize the number of used processors
    status = solve(m)
    
    return (getvalue(x), getvalue(y), getobjectivevalue(m))
end
;

using JuMP

# Models for 2.C
function minTimeWithTaskDependence(workload, num_proc, precedence, M)
    # workload: workload of each task
    # num_proc: number of processors available
    # precedence: precedence[i] is the set of tasks precede task i
    # M: a common upper bound used to ensure tasks on the same processor cannot overlap
    #    Note that M can be set to sum of worload of all tasks. 
    #
    # returned: (x,y,z) where
    # 1. x[i,j] is whether task i is scheduled to processor j
    # 2. y[i] is the start time of task i
    # 3. z is the earliest completion time
    
    num = length(workload)
    
    m = Model()
    @variable(m, x[1:num, 1:num_proc], Bin) # whether task i is scheduled to processor j
    @variable(m, y[1:num] >= 0) # the start time of task i
    @variable(m, z) # earlist completion time
    @variable(m, u[1:num, 1:num, 1:num_proc], Bin) # whether both task i,k are scheduled to processor j
    @variable(m, v[1:num, 1:num], Bin) # whether task i completes before task k
    @variable(m, r[1:num, 1:num], Bin) # whether task k completes before task i
    
    for i=1:num # set partition constraint
        @constraint(m, sum{x[i,j], j=1:num_proc} == 1)
    end
    
    for i=1:num
        for k in precedence[i] # precendence enforcement
            @constraint(m, y[i] >= y[k] + workload[k])
        end
    end
    
    for i=1:num-1
        for k=i+1:num
            for j=1:num_proc # whether both task i,k are scheduled to processor j
                @constraint(m, x[i,j] >= u[i,k,j])
                @constraint(m, x[k,j] >= u[i,k,j])
                @constraint(m, x[i,j] + x[k,j] <= u[i,k,j] + 1)
            end
        end
    end
    
    for i=1:num
        for k=1:num 
            # whether task i completes before task k 
            @constraint(m, y[i]+workload[i]-y[k] <= M*(1-v[i,k]))
            # whether task k completes before task i
            @constraint(m, y[k]+workload[k]-y[i] <= M*(1-r[i,k])) 
        end
    end
    
    for i=1:num
        for k=1:num 
            # if task i,k on same processor, task i completes before task k,
            # OR task k completes before task i
            @constraint(m, sum{u[i,k,j], j=1:num_proc} <= v[i,k] + r[i,k])
        end
    end
    
    for i=1:num # z >= last finished task
        @constraint(m, y[i] + workload[i] <= z)
    end
    
    @objective(m, Min, z) # minimize the completion time
    
    status = solve(m)
    
    return getvalue(x), getvalue(y), getvalue(z)
end
;

using JuMP
# Model for 2.D
function TSPwithTimeConstraint(A, times, v, upper, ϵ, lower)
    # A: the n*(n+1) distance matrix
    # times: times[i] is the time constraint of location i
    # v: travelling speed
    # upper: the common upper bound, can be set to n
    # ϵ: the positive number that is small enough, can be set to 1
    # lower: the lower bound, can be set to -1*(sum of length of all edges)
    #
    # returned: (u,y) where
    # 1. u: u[i] is the position of city i in the tour
    # 2. y: y[i] is the distance from the start location to city i in the tour
    # Note that city (n+1) also refers to the start city
    
    (row, col) = size(A) # n*(n+1)
    if col != row+1
        error("We need a nx(n+1) matrix!")
    end
    
    dists = v*times # transform time constrant to distance constrant
    
    m = Model()
    @variable(m, y[1:col]>=0) # distance from the starting node
    @variable(m, z[1:col, 1:col], Bin) # whether edge from i to j is in the tour
    @variable(m, u[1:col], Int) # position of each node in the tour
    @variable(m, x[1:col, 1:col], Bin) # x[i,k] is whether node i is visited before node k 
    
    for i=1:row # leave node 1 to n exactly once
        @constraint(m, sum{ z[i,k], k=1:col} == 1)
    end
    
    for k=2:col # enter node 2 to n+1 exactly once
        @constraint(m, sum{ z[i,k], i=1:row} == 1)
    end
    
    @constraint(m, z[1,col] == 0) # no edge from node 1 to node (n+1)
    @constraint(m, z[col,1] == 0) # no edge from node n+1 to node 1
    
    for i=1:col # no self loop
        @constraint(m, z[i,i] == 0)
    end
    
    for i=1:row # positions
        @constraint(m, 1<= u[i] <= row)
    end
    
    @constraint(m, u[1] == 1)
    @constraint(m, u[col] == col)
    
    for i=1:row
        for k=2:row # position constraint between node i and k
            @constraint(m, u[i]-u[k]+row*z[i,k] <= row-1)
        end
    end
    
    for i=1:row
        for k=1:col # if node i is visited before node k, then y[i]+A_ik<=y[k]
            @constraint(m, u[k]-u[i]-1 <= upper*x[i,k] - ϵ*(1-x[i,k]))
            @constraint(m, y[k]-y[i]-A[i,k] >= lower*(1-x[i,k]))
        end
    end
    
    @constraint(m, y[1] == 0) # start at first node 
    
    for i=2:row # we need to reach each location within its time constraint
        @constraint(m, y[i] <= dists[i])
    end
    
    status = solve(m)
    
    return getvalue(u), getvalue(y)
end
;

# Data generator
n = 20 # number of tasks
min_end = 5 # the min end time of any task
max_end = 15 # the max end time of any task
min_workload = 1 # the min workload of any task
max_workload = 4 # the max workload of any task, here we restrict max_workload <= min_end
min_bonus = 5 # the min bonus
max_bonus = 25 # the max bonus

function gen_data(is_float)
    if max_workload > min_end
        error("Min end time needs to be greater that max workload")
    end
    
    if is_float == false
        end_time = rand(min_end:max_end, n) # the end time of each task 
        workload = zeros(n)
        for i=1:n
            workload[i] = rand(min_workload:max_workload) # workload
        end
        start_time = end_time - workload # the start time of each task
        bonus = rand(min_bonus:max_bonus, n) # the bonus of each task
    else
        end_time = rand(n)*(max_end - min_end) + min_end # the end time of each task 
        workload = zeros(n)
        for i=1:n
            workload[i] = rand()*(max_workload - min_workload) + min_workload # workload
        end
        start_time = end_time - workload # the start time of each task
        bonus = rand(n)*(max_bonus - min_bonus) + min_bonus # the bonus of each task
    end
    
    data = cat(2, end_time, workload, start_time, bonus)
    
    # return it
    return data
end

# generate data and save it to file
#data = gen_data(false)
#writedlm("ZhangCui_0506.txt", data)
#data = gen_data( true)
#writedlm("ZhangCui_0506_float.txt",data)
;

# Read the saved data
data = readdlm("ZhangCui_0506.txt")
end_time = data[:,1]
workload = data[:,2]
start_time = data[:,3]
bonus = data[:,4]

# function to find conflict tasks for each task
function findConflictTasks(end_time, start_time)
    num = length(end_time)
    tmp = zeros((num, num))
    for i=1:num-1
        for k=i+1:num
            # does task i and k conflict? 
            if end_time[i] >= start_time[k] && start_time[i] <= end_time[k]
                tmp[i,k] = 1
                tmp[k,i] = 1
            end
        end
    end
    
    return [ [ k for k in filter(k -> tmp[i,k]==1, 1:num) ] for i=1:num ]
end

conflicts = findConflictTasks(end_time, start_time)
;

# 2.A.part(1): Select the subset of tasks with maximum total bonus, when only one processor is available
(A1_x, A1_obj) = maxBonus(end_time, start_time, bonus, conflicts)
println("Selected tasks are:")
for i=1:length(bonus)
    if A1_x[i]==1
        print(i, "    ")
    end
end
println()
println("Maximum bonus is: ", A1_obj)

# 2.A.part(2): Decide the minimum number of processor required to complete all tasks
(A2_x, A2_y, A2_obj) = minNumProcessors(end_time, start_time, conflicts)
println("Minimum number of processors needed: ", A2_obj)

# plot them
using PyPlot

fig = figure(figsize=(10,4))
ax = fig[:add_subplot](111)

ylabel("Process ID")
xlabel("Timeline")

num_tasks = length(A2_y)
procIDs = zeros(num_tasks)
for i=1:num_tasks
    for k=1:num_tasks
        if A2_x[i,k]==1
            procIDs[i] = k
        end
    end
end

for i=1:num_tasks # annotate each Task with task ID
    procID = procIDs[i]
    ax[:annotate](string("task", string(i)), (start_time[i]+0.1, procID+0.1), textcoords="data")
end

for i=1:num_tasks
     # plot line segment between start and end time of each task
    procID = procIDs[i]
    plot([start_time[i], end_time[i]], [procID, procID], marker="o", linestyle="-")
end
;

# 2.B.part(1): Minimize the time to complete all tasks when no constraint on the start and end time of tasks
num_proc = 5 # the number of processors available
(B1_x, B1_z) = minTime(workload, num_proc)
println("All tasks are completed by ", B1_z)

# 2.B.part(2): Bin Packing -- minimum number of processors needed to complete all tasks before given deadline T
T = 10 # the deadline (bin size)
(B2_x, B2_y, B2_obj) = minResource(workload, T)
println("Minimum number of processors needed: ", B2_obj)
for j=1:length(end_time)
    if B2_y[j]==0
        continue
    end
    print("Tasks scheduled on processor ", j, ": ")
    for i=1:length(end_time)
        if B2_x[i,j]==1
            print(i, "  ")
        end
    end
    println()
end

# Data for 2.C
# This is a manuall created small dataset
# four levels: 1-(2,3,4,5)-(6,7,8)-(9,10)
# The execution time increases very fast since the number of variables and constraints is O(n^2 * m)
workload = [2 3 4 5 3 2 2 2 3 4]
precedence = Dict(1=>[], 2=>[1], 3=>[1], 4=>[1], 5=>[1], 6=>[2 3], 7=>[2 4], 8=>[4 5], 9=>[6 8], 10=>[7 8])
println("Total workload:", sum(workload))
;

# Given m processors, find the earliest completion time
num_proc = 3
# lower bound and epsilon, based on analysis of the data
M = sum(workload) # this is an upper bound suitable for all we need 
(C_x, C_y, C_z) = minTimeWithTaskDependence(workload, num_proc, precedence, M)
println("All tasks are completed by: ", C_z)
;

# Let's show the task dependence and scheduling
xdata = C_y # xdata is the scheduled time for each task
ydata = zeros(length(xdata))
for j=1:num_proc
    tids = Int[]
    for i=1:length(workload)
        if C_x[i,j] == 1
            ydata[i] = j # ydata is the assigned processor
        end
    end
end

using PyPlot

fig = figure(figsize=(10,4))
ax = fig[:add_subplot](111)

ylim(0, num_proc+1)
xlim(0, C_z+1)
ylabel("Process ID")
xlabel("Timeline")
yticks(1:num_proc)

for i=1:length(xdata) # annotate each point with task ID, workload
    ax[:annotate](string("t", string(i), ", w=", workload[i]), (xdata[i], ydata[i]+0.1), textcoords="data")
end

for i=1:length(xdata)
    for k in precedence[i] # plot lines for the precedence
        plot([xdata[k], xdata[i]], [ydata[k], ydata[i]], marker="o", linestyle="-")
    end
end
;

# Let's vary the number of processors
num_tasks = length(workload)
num_procs = [i for i=1:num_tasks]
min_time = zeros(num_tasks)

M = sum(workload)
for i in num_procs
    (C_x, C_y, C_z) = minTimeWithTaskDependence(workload, i, precedence, M)
    min_time[i] = C_z
end
;

using PyPlot

fig = figure(figsize=(10,4))
ylabel("Earliest Completion Time")
xlabel("Number of Processors")

plot(num_procs, min_time, "ro-")
;

# the n*(n+1) distance matrix, borrowed from the distance matrix betweeen 10 major airports in the US
# we add one more column, which is same as first column
A = [     0   587  1212   701  1936   604   748  2139  2182   543   0
        587     0   920   940  1745  1188   713  1858  1737   597   587
        1212   920     0   879   831  1726  1631   949  1021  1494  1212
        701   940   879     0  1379   968  1420  1645  1891  1220   701
        1936  1745   831  1379     0  2339  2451   347   959  2300  1936 
        604  1188  1726   968  2339     0  1092  2594  2734   923   604
        748   713  1631  1420  2451  1092     0  2571  2408   205   748
        2139  1858   949  1645   347  2594  2571     0   678  2442  2139
        2182  1737  1021  1891   959  2734  2408   678     0  2329  2182
        543   597  1494  1220  2300   923   205  2442  2329     0   543 ]

upper = 10 # the upper bound for some constraint in the model
ϵ = 1 # the positive number that is small enough
lower = -sum(A) # the lower bound needed for some constraint in the model

v = 1 # the travelling speed
# original tour is cities 1-6-4-5-8-9-3-2-7-10-1

# set time constraints
# For this dataset, 200000 is equivalent to no time constraint on that node
# below one ensures node 7 and 9 needs to be visited first
#times = [0 200000 200000 200000 200000 200000 888 200000 5555 200000]

# below one ensures the tour is cities 1-10-7-2-4-6-3-8-9-5-1
times = [0 1500 7000 2500 8000 4000 800 7000 8000 600]

(D_u,D_y) = TSPwithTimeConstraint(A, times, v, upper, ϵ, lower)
println("Positions of each city in the tour:", [Int(x) for x in round(D_u)])
println("Distance to each city:", round(D_y))

# let's plot and compare original tour and new tour due to time constraint
function getTour(positions)
    num_cities = length(positions)-1
    tour = zeros(Int, num_cities+1)
    for i=1:num_cities
        for k=1:num_cities
            if positions[k] == i
                tour[i] = k
            end
        end
    end
    
    tour[num_cities+1] = 1 # start and end at node 1
    
    return tour
end

# coordinates of each city
data = [ 33.636700  -84.427863 
         41.977320  -87.908005
         39.861667 -104.673166
         29.645417  -95.278888
         33.942437 -118.408121
         25.795361  -80.290110
         40.639926  -73.778694
         37.618806 -122.375416
         47.449889 -122.311777
         38.851444  -77.037721 ]

# plot original tour and the tour with time constraint
original_tour = [1 6 4 5 8 9 3 2 7 10 1]
tour = getTour([Int(round(k)) for k in D_u])

using PyPlot

fig = figure(figsize=(12,3))

axis("equal")
ox = [ data[i,2] for i in original_tour ]
oy = [ data[i,1] for i in original_tour ]

nx = [ data[i,2] for i in tour ]
ny = [ data[i,1] for i in tour ]

ax1 = subplot(1,2,1)
title("Original Tour")
for i=1:length(original_tour)-1 # positions in the tour
    ax1[:annotate](string("#", i), (ox[i], oy[i]+0.1), textcoords="data")
end
plot(ox, oy, "ro-")

ax2 = subplot(1,2,2)
title("Tour with Time Constraint")
for i=1:length(tour)-1 # positions in the tour
    ax2[:annotate](string("#", i), (nx[i], ny[i]+0.1), textcoords="data")
end
plot(nx, ny, "bo-")
;