# Copyright (c) 2021 Satlab A/S <satlab@satlab.com>
# SPDX-License-Identifier: GPL-3.0

import numpy as np

from gnuradio import gr


class RowCollector(gr.sync_block):
    def __init__(self, row_size, max_rows=100000):
        gr.sync_block.__init__(
            self,
            name=type(self).__name__,
            in_sig=[(np.float32, row_size)],
            out_sig=[]
        )

        self.data = np.zeros((max_rows, row_size), dtype=np.float32)
        self.index = 0
        self.active = False

    def start_collect(self):
        self.index = 0
        self.active = True

    def stop_collect(self):
        self.active = False

    def get_collected(self):
        return self.data[:self.index, ]

    def work(self, input_items, output_items):
        bins = input_items[0][0]
        if self.active:
            if self.index < self.data.shape[0]:
                self.data[self.index, ] = bins
                self.index += 1
            else:
                self.active = False

        return len(input_items[0])
