from dmor.model import Model

T = 4
demand = {1: 5, 2: 8, 3: 4, 4: 6}
h = 1
c = {1: 2, 2: 1, 3: 4, 4: 1}
I1 = 3

m = Model()

m.add_set("T", range(1, T+1))
m.add_set("T2", 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_variable("x", index="T")
m.add_variable("I", index="T2")

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

m.add_constraint("balance", balance, index="T2")


def obj(genmodel):
    return sum(genmodel.c[t] * genmodel.x[t] + genmodel.h * genmodel.I[t+1] for t in genmodel.T)

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


m.solve(solver="appsi_highs")
m.display()
m.slacks()
print("Constraint List")
print(m.constraints_list())
print("Objective Expression")
print(m.objective_expression())