"""  This example illustrates how the Toolkit could be used to develop a hydrant rating curve used in fire flow studies. 
 This curve shows the amount of flow available at a node in the system as a function of pressure. 
 The curve is generated by running a number of steady state hydraulic analyses with the node of interest subjected 
 to a different demand in each analysis. For this example we assume that the ID label of the node of interest is MyNode 
 and that N different demand levels stored in the array D need to be examined. The corresponding pressures will be stored in P.
 To keep the code more readable, no error checking is made on the results returned from the Toolkit function calls.
  
  https://github.com/OpenWaterAnalytics/EPANET/wiki/Example-2)
"""
from epyt import epanet

# Load a network
d = epanet('Net1.inp')

# Select node of interest
nodeID = '32'

# Select number of iterations
r = d.arange(1.0, 10.1, 0.5)

d.openHydraulicAnalysis()

nodeIndex = d.getNodeIndex(nodeID)

categoryIndex = 1
B = d.getNodeBaseDemands()[categoryIndex][nodeIndex - 1]
D = [i * B for i in r]

P = []
for i in range(len(D)):
    d.setNodeBaseDemands(nodeIndex, D[i])
    d.initializeHydraulicAnalysis()
    d.runHydraulicAnalysis()
    P.append(d.getNodePressure(nodeIndex))

d.closeHydraulicAnalysis()

d.plot_ts(X=D, Y=P, xlabel=f'Base Demand ({d.units.NodeDemandUnits})',
          ylabel=f'Pressure ({d.units.NodePressureUnits})',
          marker='x', figure_size=[3, 2.5], constrained_layout=True)
d.plot_show()

# Unload library
d.unload()
