import sys
from docplex.cp.model import *

model = CpoModel()

####################
# Variables declaration
####################

NbLane  = 1
Lane = range(NbLane)

PowerUnit = ["nuclear_1","nuclear_2","nuclear_3","nuclear_4"]
Duration =  [5, 5, 5, 5]
ReleaseDate = [0]
Precedences = [("nuclear_1", "nuclear_2"), ("nuclear_2", "nuclear_3"), ("nuclear_3", "nuclear_4")]


itvs = {}
for h in Lane:
    for i,t in enumerate(PowerUnit):
        itvs[h,t] = model.interval_var(start = [ReleaseDate[h], INTERVAL_MAX], size=Duration[i])


energy_level = step_at(0, 0)
for h in Lane:
    for i,t in enumerate(PowerUnit):
        energy_level += model.step_at_start(itvs[h,t], Duration[i])
   
            
####################
# Constraints
####################

for h in Lane:
    for p in Precedences:
        model.add( model.end_before_start(itvs[h,p[0]], itvs[h,p[1]]) )
    

####################
# Objective function
####################

obj = model.max( model.end_of(itvs[h,"nuclear_4"]) for h in Lane)
model.add( model.minimize( obj ) )

print("\nSolving model....")
solution = model.solve()
print("Done")


####################
# Results display
####################

print("Cost will be " + str( solution.get_objective_values()[0] ))

import docplex.cp.utils_visu as visu
import matplotlib.pyplot as plt

from pylab import rcParams
rcParams['figure.figsize'] = 15, 3


energy_levelF = CpoStepFunction()

for h in Lane:
    for i,t in enumerate(PowerUnit):
        itv = solution.get_var_solution(itvs[h,t])
        energy_levelF.add_value(itv.start, INT_MAX, Duration[i])

visu.timeline('Solution SchedCumul')

visu.panel(name="Schedule")
for h in Lane:        
    for i,t in enumerate(PowerUnit):
        visu.interval(solution.get_var_solution(itvs[h,t]), h, t)
    
visu.panel(name="Energy level")
visu.function(segments=energy_levelF, style='area')
visu.show()