# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="zSf0zxukK6bz"
# # ISE-533 Project 4: Allocation Problem for COVID-19 Ventilators
# ## Forecast with Uncertainty - Simple Recourse Model
#
# By: <NAME>, <NAME>, <NAME>, <NAME>

# + id="nU5ULYslK6b3"
import os
import numpy as np
import pandas as pd
import math
from pyomo.environ import *
from collections import defaultdict
import pyomo.environ as pyo
from pyomo.opt import SolverFactory
from datetime import timedelta
import matplotlib.pyplot as plt

# + id="ZD8ZAxSeK6b3"
raw = pd.read_csv('project 4 dataset/init_data.csv', header=0, index_col=0)
raw

# + id="ZtD0Zj9VK6b4"
first_date = '2020/03/25'
t = 7     # for a week
T = 14
P = 10
lbd = 0.1 # penalty

S = 51    # 50 states and SNS
states_name = raw['State']
SNS_stock = 12000

init_ratio = 0.8        # 0.6 of the stock is used for covid-19
flow_bound_ratio = 0.3  #0.2, 0.25, 0.3
stock_bound_ratio = 0.5 #0.6, 0.5

# + id="EMwY3tD_K6b5"
#first iteration values
data = raw.rename(columns={'Available capacity': "stock_self"})
data = data[:-1]
data['stock_self'] = init_ratio * data['stock_self']
data['stock_nb'] = None
for i in range(S):
    data.loc[i, 'stock_nb'] = str([0 for _ in eval(data.loc[i, 'nb'])])
data['stock_sns'] = [0 for _ in range(S)]

# + id="clWSz2HcK6b5"
idx = 0
flow_mapping = {}
flow_mapping_rev = {}
stocks = {}
out_flow = defaultdict(list)
in_flow = defaultdict(list)

for i in range(S):
    nbs = eval(data.iloc[i, 3])
    out_flow[i] = nbs
    stocks[i ,i] = data['stock_self'][i]
    for n in nbs:
        stocks[i, n] = 0
        flow_mapping[(i, n)] = idx
        flow_mapping_rev[idx] = (i, n)
        in_flow[n].append(i)
        idx += 1

in_flow[51] = []
out_flow[51] = list(range(51))
for n in range(S):
    stocks[51, n] = 0
    flow_mapping[(51, n)] = idx
    flow_mapping_rev[idx] = (51, n)
    idx += 1

# + id="Pc7c61bCK6b6"
stocks[51, 51] = SNS_stock 


# + id="PFL-lXarK6b6"
def first_stage(stocks:dict):
    model = ConcreteModel()

    
    #Sets
    F = idx
    model.time_set = list(range(t))
    model.state_set = list(range(S))
    model.state_p_set = list(range(S + 1))
    model.flow_set = list(range(F))
    
    # Parameters
    model.demand = Param(model.state_set, model.time_set, within=NonNegativeReals, mutable=True)

    model.U = [round(flow_bound_ratio * stocks[i, i]) for i in model.state_p_set]
    model.G = [round(stock_bound_ratio * stocks[i, i]) for i in model.state_p_set]
    
    
    # Variables
    model.s = Var(model.state_p_set, within=NonNegativeReals) # s_jj
    model.s_nb = Var(model.flow_set, within=NonNegativeReals) # s_ij
    model.x = Var(model.flow_set, within=NonNegativeReals)    # x_ij
    
    
    # Objective
    model.first = Var(within = NonNegativeReals)
    model.second = Var(within = NonNegativeReals)
    model.first_penalty = Constraint(expr = model.first == lbd * summation(model.s_nb))
    
    model.cost = Objective(expr = model.first + model.second, sense = minimize)
    
    
    # Constraints
    model.stock_self = Constraint(model.state_set, rule=lambda model, j: \
                                  model.s[j] + sum(model.x[flow_mapping[j, i]] for i in out_flow[j]) - model.x[flow_mapping[51, j]]== \
                                  stocks[j, j])
    model.stock_SNS_self = Constraint(expr = \
                                      model.s[51] + sum(model.x[flow_mapping[(51, j)]] for j in out_flow[51]) == \
                                      stocks[51, 51])

    model.stock_nb = ConstraintList()
    for j in model.state_set:
        for i in in_flow[j]:
            model.stock_nb.add(model.s_nb[flow_mapping[i,j]] - model.x[flow_mapping[i,j]] == stocks[i,j])
            
    model.s_bound = Constraint(model.state_p_set, rule=lambda model, j: model.s[j] >= model.G[j])
    
    model.x_bound = ConstraintList()
    for i in model.state_p_set:
        for j in out_flow[i]:
            model.x_bound.add(model.x[flow_mapping[i,j]] <= model.U[i])
            model.x_bound.add(-stocks[i,j] <= model.x[flow_mapping[i,j]])
            
    model.benders_cuts = ConstraintList()
    
    
    return model


# + id="ts82ID9JK6b8"
def second_stage(demands:list, u:float):
    model = ConcreteModel()
    
    # Parameters
    model.time_set = list(range(t))
    
    # Variables
    model.delta = Var(model.time_set, within=NonNegativeReals)
    model.dual = Suffix(direction=Suffix.IMPORT)
    
    # Objective
    model.obj = Objective(expr = summation(model.delta), sense = minimize)
    
    # Constraint
    model.delta_bound = ConstraintList()
    for day in model.time_set:
        model.delta_bound.add(expr = model.delta[day] >= demands[day] - u)
        model.delta_bound.add(expr = model.delta[day] >= 0)
        
    return model


# + id="Jckgyi_dK6b9"
def benders(date, stocks, max_iteration = 200):
    opt = SolverFactory('glpk')
    master = first_stage(stocks)
    
    iter = 0
    
    lower = -float("inf")
    upper = float("inf")
    
    while iter < max_iteration:
        iter += 1
        # print("iteration", iter)
        
        opt.solve(master)
        lower = max(lower, master.cost.expr())
        # print("Lower Bound:", lower)
        
        if upper - lower <= 0.001:
            print("Converged")
            break
            
        x = master.x.get_values()   # x_ij
        s = master.s.get_values()   # s_jj
        s_nb = master.s_nb.get_values()   # s_ij
        
        sub_bound = 0
        sub_duals = defaultdict(int)
        benders_cut_constant = 0
        
        for j in range(S):
            state = states_name[j]
            state_demand = pd.read_csv("project 4 dataset/prediction/{}/{}.csv".format(date.strftime('%m%d'), state))
            
            path = [state_demand['InvVen_mean'], state_demand['InvVen_upper'], state_demand['InvVen_lower']]
            
            u = s[j] + sum(s_nb[flow_mapping[i, j]] for i in in_flow[j])

            for d in path:
                demands = d.iloc[:7]  
                
                sub = second_stage(demands, u)
                opt.solve(sub)
                
                duals = []
                
                for day in range(1, 1 + t):
                    duals.append(sub.dual[sub.delta_bound[day]])
                    
            sub_duals[j] += sum(duals)/3
            benders_cut_constant += sum([demand * dual for demand, dual in zip(demands, duals)])/3
            sub_bound += sub.obj.expr()/3

                
        sub_bound += lbd * sum(s_nb[j] for j in range(idx-S))
        upper = min(upper, sub_bound)
        # print("Upper Bound:", upper)
        
        # Add benders cuts
        master.benders_cuts.add(expr = master.cost >= lbd * sum(master.s_nb[flow] for flow in range(idx-S)) - \
                               sum(sub_duals[j] * (sum(master.s_nb[flow_mapping[i,j]] for i in in_flow[j]) + \
                                                  master.s[j]) for j in range(S)) + \
                                benders_cut_constant)
        
    return master

# + [markdown] id="LstF7h0jK6b_"
# ### Allocation For 10 Weeks

# + id="Y-hEDhg-K6cA"
date = pd.to_datetime("2020/03/25")
all_models = []
data_list = []
final = pd.DataFrame()

for week in range(10):
    solved_master = benders(date, stocks, max_iteration = 20)
    all_models.append(solved_master)
    s_nb = solved_master.s_nb.get_values()
    s = solved_master.s.get_values()
    
    date += timedelta(7)
    
    
    stocks = {}
    for flow, index in flow_mapping_rev.items():
        stocks[index] = s_nb[flow]
    for i in range(S+1):
        stocks[i,i] = s[i]
        
    
    temp = raw.copy()
    nb = 0
    
    for i in range(S+1):
        for index, flow in flow_mapping.items():
            if index[0] == i:
               nb += stocks[index]
        
        temp.loc[i, 'Stock'] = stocks[i,i] + nb
    
    temp = temp['Stock']
    data_list.append(temp)
    
final = pd.concat(data_list, axis = 1).set_index(states_name)

# + id="CkOSMu8bK6cD"
final

# + [markdown] id="AMXHmVaNK6cE"
# ### Unmet Demand with Actual Demand

# + id="lqg5WMK6K6cF"
states_name = raw['State']
state_demand = pd.DataFrame(columns=['date']) 

for j in range(S):
    state = states_name[j]
    temp = pd.read_csv("project 4 dataset/eval/{}.csv".format(state), header=None)
    temp.rename(columns = {0: 'date', 1: state}, inplace = True)
    state_demand = state_demand.merge(temp, on='date', how='outer') 

# + id="4jnA5baYK6cG"
state_demand = state_demand.set_index('date')

# + id="avV7Le71K6cG"
eval = state_demand.groupby(np.arange(len(state_demand))//7).mean()
eval

# + id="_KrMVzmCK6cG"
final1 = final.copy()
final1 = final1.transpose()
final1 = final1.iloc[: , :-1]
final1.index = eval.index
final1

# + [markdown] id="tRsm8xInK6cH"
# ### Unmet Demand Every Week and Sum

# + id="JxwdbcE8K6cH"
final1.to_csv('project 4 dataset/SR_supply_08_03_05.csv')

# + id="gJ0tWA2yK6cH"
final = pd.read_csv('project 4 dataset/SR_supply_08_03_05.csv')
final = final.drop(['Unnamed: 0'], axis = 1)
final

# + id="Y38Nj25cK6cI"
states_name = raw['State']
state_demand = pd.DataFrame(columns=['date']) 

for j in range(S):
    state = states_name[j]
    temp = pd.read_csv("project 4 dataset/eval/{}.csv".format(state), header=None)
    temp.rename(columns = {0: 'date', 1: state}, inplace = True)
    state_demand = state_demand.merge(temp, on='date', how='outer') 

state_demand = state_demand.set_index('date')
state_demand = state_demand.reset_index()
state_demand.head()

# + id="lZ9FYUVZK6cI"
final_repeat = final.loc[final.index.repeat(7)].reset_index(drop=True)
final_repeat

# + id="RlEHq_kUK6cI"
unmet_everyday = state_demand.drop(['date'], axis = 1).subtract(final_repeat)
unmet_everyday[unmet_everyday <= 0] = 0
unmet_everyweek = unmet_everyday.groupby(np.arange(len(unmet_everyday))//7).sum()

unmet_everyweek

# + id="tWvJLwunK6cJ"
result_08_03_05 = pd.DataFrame(unmet_everyweek.sum(axis = 1), columns = ['Unmet Demand'])
result_08_03_05.index = [1,2,3,4,5,6,7,8,9,10]
result_08_03_05 = result_08_03_05.rename_axis('Week').reset_index()
result_08_03_05.to_csv('project 4 dataset/result_08_03_05.csv')
result_08_03_05

# + id="C9BziD2QK6cJ"
pd.DataFrame(result_08_03_05.sum(axis = 1), columns = ['Unmet Demand']).sum(axis = 0)

StarcoderdataJupyterScriptsDedupFiltered