from dmor.model import Model

T = 10
demand = {1: 50, 2: 80, 3: 40, 4: 60, 5: 30, 6: 20, 7: 90, 8: 70, 9: 50, 10: 60}
daverage= sum(demand[t] for t in range(1, T+1)) / T
capacity = 100
h = 3
c = {1: 5, 2: 4, 3: 5, 4: 9,  5: 6, 6: 4, 7: 8, 8: 7, 9: 5, 10: 6}
cost_over = 2
cost_under = 3
I1 = 0

m = Model()

m.add_set("T", range(1, T+1))
m.add_set("T+1", range(2, T+2))

m.add_parameter("d", demand, index="T")
m.add_parameter("dav", daverage)
m.add_parameter("h", h)
m.add_parameter("cost_over", cost_over)
m.add_parameter("cost_under", cost_under)
m.add_parameter("c", c, index="T")
m.add_parameter("capacity", capacity)

m.add_variable("x", index="T")
m.add_variable("I", index="T+1")
m.add_variable("over", index="T", lb=0)
m.add_variable("under", index="T", lb=0)

def balance(m, t):
    if t == 2:
        return m.I[t] == I1 + m.x[1] - m.d[1]
    return m.I[t] == m.I[t-1] + m.x[t-1] - m.d[t-1]

m.add_constraint("balance", balance, index="T+1")

def capacity_constraint(m, t):
    return m.x[t] <= m.capacity

m.add_constraint("capacity_const", capacity_constraint, index="T")

def smoothing_constraint(m, t):
    return m.over[t] - m.under[t] == m.x[t] - m.dav

m.add_constraint("smoothing_const", smoothing_constraint, index="T")

def obj(foo):
    return sum(foo.c[t] * foo.x[t] + foo.h * foo.I[t+1] + foo.cost_over * foo.over[t] + foo.cost_under * foo.under[t]    for  t in foo.T)

m.set_objective(obj, sense="min")


m.solve(solver="appsi_highs")
m.display()
m.slacks()