from pyomo.environ import *
from numpy import *
import matplotlib.pyplot as plt

# model
model = ConcreteModel()

# Set
model.t = Set(initialize=[1, 2, 3, 4], doc='time')

# Parameter
model.pb = Param(model.t, initialize={1: 3, 2: 1, 3: 2, 4: 4}, doc='Base load')
model.c = Param(model.t, initialize={1: 0.5, 2: 0.1, 3: 1, 4: 0.3}, doc='Cost ')
model.dt = Param(initialize=1, doc='time duration ')
# Variable

model.pl = Var(model.t,bounds=(0, None)) #todo example: 2
# model.pl = Var(RangeSet(1, 4),bounds=(0, 3)) #todo example: 3
# model.pl = Var(RangeSet(1, 4),bounds=(0, 3)) #todo example: 4
# Constraints
model.c1 = Constraint(expr=sum([model.pl[i] for i in model.t]) * model.dt == 8)

# Objective

model.obj = Objective(expr=sum([(model.pb[i] + model.pl[i]) * model.c[i] for i in model.t]) * model.dt)

# Solve

opt = SolverFactory('cplex')
result = opt.solve(model)
print(result)
# model.display()
# model.pprint()

# print variables


# print('------------Countable variable Solution------------------------')
for i in model.t:
    print('x[%i] = %i' % (i, value(model.pl[i])))
print('objective function = ', value(model.obj))

# Plotting
time = []
pbase = []
pload = []
cost = []

for ll in range(len([model.t, model.pb, model.pl, model.c])):
    if ll == 0:
        for i in model.t:
            time.append(value(model.t[i]))
    elif ll == 1:
        for i in model.pb:
            pbase.append(value(model.pb[i]))
    elif ll == 2:
        for i in model.pl:
            pload.append(value(model.pl[i]))
    else:
        for i in model.c:
            cost.append(value(model.c[i]))

fig, (ax0, ax1) = plt.subplots(nrows=2,ncols= 1)

ax0.plot(time, pbase, 'b-', label='Base')
ax0.plot(time, pload, 'g--', label='Controable Load')
ptotal = [pbase[i] + pload[i] for i in range(len(pbase))]
ax0.plot(time, ptotal, 'r', label='Total Load')
ax0.set_ylabel('Electric power (W)')
ax0.legend(loc='upper left')
ax1.plot(time, cost, 'y-+', label='Cost')
ax1.legend(loc='upper left')
ax1.set_xlabel('Time (hour)')
ax1.set_ylabel('Electricity price ($/W/h)')
plt.suptitle("Energy management system-Example (4)")
plt.show()