import numpy as np
from gnuradio import gr
import pmt


class MTS_basic_py(gr.basic_block):
    """
    docstring for block MTS_basic_py
    """

    def __init__(self, pos, vel, rcs, fs, fc, prf, len_tag_key, pulse_len):
        gr.basic_block.__init__(self,
                                name="MTS_basic_py",
                                in_sig=[np.complex64],
                                out_sig=[np.complex64]),
        self.pos = pos
        self.vel = vel
        self.rcs = rcs
        self.fs = fs
        self.fc = fc
        self.PRF = prf
        self.PRI = 1 / self.PRF
        self.len_tag_key = len_tag_key
        self.curr_pulse = 0
        self.pulse_len = pulse_len

    def forecast(self, noutput_items, ninput_items_required):
        # setup size of input_items[i] for work call
        for i in range(len(ninput_items_required)):
            ninput_items_required[i] = noutput_items

    def general_work(self, input_items, output_items):
        in0 = input_items[0][:len(output_items[0])]
        out = output_items[0]
        # Start and end of current input
        start_n = self.nitems_read(0)
        end_n = start_n + len(in0)

        # Start of burst tags
        tags = self.get_tags_in_range(0, start_n, end_n,
                                      pmt.string_to_symbol(self.len_tag_key))
        # No tags found in the input. Consume all, return nothing
        if len(tags) == 0:
            out[:] = 0
            out = np.trim_zeros(out)
            self.consume(0, len(input_items[0]))
            return len(out)
        # Since we're moving the first tag in the input to offset 0, get info
        # for the first tag
        tag = tags[0]
        # Offset not aligned. Consume samples before the tag and output nothing
        if tag.offset != start_n:
            out[:] = 0
            out = np.trim_zeros(out)
            self.consume(0, tag.offset - start_n)
            return len(out)
        # Tags should be aligned, get the pulse length from the tag value,
        # and (for now) copy the input pulse directly to the output and
        # return only the part of the output with the pulse
        else:
            self.curr_pulse += 1
            out[:self.pulse_len] = in0[:self.pulse_len]
            out[self.pulse_len:] = 0
            out = np.trim_zeros(out)
            self.consume(0, self.pulse_len)
            return len(out)

    def simulate_target(self):
        return 1