#################################################################################################################
## Package

class Tank:
    def __init__(self, level, rate, upstream_tanks=None):
        self.level = level
        self.rate = rate
        self.upstream_tanks = upstream_tanks
        return
    
    @property # allows the method use without brackets 
    def Q(self):
        tank_area = 1
        return self.level * self.rate / tank_area
        

class LevelGauge:
    def __init__(self,tank):
        self.tank = tank
        self.data = []
        return
    
    def take_measurement(self):
        self.data.append(self.tank.level)
        return


class Network:
    def __span_network(self, i):
        self.tanks[i] = Tank(levels[i], rates[i], upstream_tanks=network_structure.get(i))
        if network_structure.get(i) is not None:
            for tank in network_structure.get(i):
                #recursion
                self.__span_network(tank) 
    
    def generate(self, network_structure, levels, rates):
        
        self.final_tank = max(network_structure.keys())
        
        if network_structure[self.final_tank] == []:
            print("The last tank has no upstream tanks. The network structure you provided is invalid.")
        else:
            self.tanks = {}
            self.__span_network(self.final_tank)
        
        return "network generated"
    
    def __fill_tank(self,tank_id):
        tank = self.tanks[tank_id]
        if tank.upstream_tanks is not None:
            for uptank_id in tank.upstream_tanks:
                # recursion
                self.__fill_tank(uptank_id)          
                
                uptank = self.tanks[uptank_id]
                uptank.level -= uptank.Q
                tank.level += uptank.Q
    
    def install_level_gauges(self):
        self.gauges = {}
        for tank_id in self.tanks:
            self.gauges[tank_id] = LevelGauge(self.tanks[tank_id])
    
    def simulate(self,iterations):
        
        for gauge in self.gauges.values():
            gauge.take_measurement()
        
        for dt in range(iterations):
            self.__fill_tank(self.final_tank)
            
            for gauge in self.gauges.values():
                gauge.take_measurement()
            
        return

#################################################################################################################
## Top level Api use

network_structure = {6:[4,5],4:[1,2],5:[3]}
levels = {1:10,2:20,3:15,4:8,5:44,6:2.5}
rates = {1:0.0,2:0.06,3:0.9,4:0.3,5:0.5,6:0}

network = Network()
network.generate(network_structure,levels,rates)
network.install_level_gauges()
network.simulate(110)


#################################################################################################################
## plotting results

import matplotlib.pyplot as plt
fig, axis = plt.subplots(6,1,figsize = (12,8))
for idx, ax in enumerate(axis):
    ax.plot(network.gauges[idx+1].data)
    ax.set_title(idx+1)
plt.tight_layout()