from gnuradio import gr,window
from local_code import local_code

class acquisition(gr.hier_block2):
    def __init__(self, fs, fd):

        self.fft_size = int( 1e-3*fs)
        self.window = window.blackmanharris(self.fft_size)

        gr.hier_block2.__init__(self,
            "acquisition",
            gr.io_signature(1,1,gr.sizeof_gr_complex),
            gr.io_signature(1,1, int(self.fft_size * gr.sizeof_float)))

        # aliases:
        c = lambda i, o, i_p=0, o_p=0: self.connect(i,i_p,o,o_p)
        d = lambda n, f: self.define_component( n, f)

        # Input signal; get Fourier transform of signal.
        d( "s2v_input",
            gr.stream_to_vector(gr.sizeof_gr_complex, self.fft_size))
        d( "fft_input",
            gr.fft_vcc(self.fft_size, True, self.window))

        # Local code.
        d( "local_code", local_code(svn=1, fs=fs, fd=fd))

        # Multiply local code with recv. signal.
        d( "mult", gr.multiply_vcc(self.fft_size))

        # Invers transform result.
        d( "ifft", gr.fft_vcc(self.fft_size, False, self.window))

        # Get signal magnitude.
        d( "mag", gr.complex_to_mag(self.fft_size) )
        d( "v2s", gr.vector_to_stream(gr.sizeof_float, self.fft_size))

        # Split vector into N branches for filtering.
        d( "one2N", gr.deinterleave(gr.sizeof_float))

        # Integrate signal: y(n) = x(n) + x(n-1) + x(n-2)
        for i in range(self.fft_size):
           d( ( "fir_%d" % i ), gr.fir_filter_fff( 1, [1.0, 1.0, 1.0] ))

        # Combine signal again.
        d( "N2one", gr.interleave(gr.sizeof_float))
        d( "s2v", gr.stream_to_vector(gr.sizeof_float, self.fft_size))


        # Connect everything.
        c( "self", "s2v_input")
        c( "s2v_input", "fft_input")
        c( "local_code", "mult")
        c( "fft_input", "mult", o_p=1)
        c( "mult", "ifft")
        c( "ifft", "mag" )
        c( "mag", "v2s" )
        c( "v2s", "one2N" )

        for i in range(self.fft_size):
            fir = ("fir_%d" % i )
            c( "one2N", fir, i_p=i )
            c( fir, "N2one", o_p=i)
        
        c( "N2one", "s2v" )
        c( "s2v", "self" )