from time import perf_counter
class FlowCollector(object):
    
    def __init__(self, model_node):
        self.node = model_node
        self.collects = [[] for _ in range(len(self.node.tracks))]
        self.timer = [[] for _ in range(len(self.node.tracks))]
        
    def clean(self):
        self.collects = [[] for _ in range(len(self.node.tracks))]

    def reset_timer(self):
        self.timer = [[] for _ in range(len(self.node.tracks))]

    def time_table(self):
        return [sum(i) for i in self.timer]
        
    def collect(self):
        for idx, track in enumerate(self.node.tracks):
            _t_s = perf_counter()
            self.collects[idx].append(track.flow().detach().numpy())

            self.timer[idx].append(perf_counter() - _t_s)
        