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}
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}
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("h", h)
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")

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 obj(foo):
    return sum(foo.c[t] * foo.x[t] + foo.h * foo.I[t+1] for t in foo.T)

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


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