#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 
# Copyright 2014 <+YOU OR YOUR COMPANY+>.
# 
# This is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3, or (at your option)
# any later version.
# 
# This software is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this software; see the file COPYING.  If not, write to
# the Free Software Foundation, Inc., 51 Franklin Street,
# Boston, MA 02110-1301, USA.
# 

import numpy
from gnuradio import gr
import os, sys, inspect
sys.path.insert(0,'/home/gnuradio/gnuradio_trlcode/gr-trl/examples/ofdm_bsld/Pfiles/LDPC_lib')
from rew import *
from LDPC_lib import *
import datetime as dt
import LDPC_dec_lib
import pmt
import time

class blsd_enc_b(gr.sync_block):
    """
    docstring for block blsd_enc_b
    """
    def __init__(self, kA1=4,kB1=4,k2=4,NA1=8,NB1=8,N2=8,M=8,num_packets=1):
        gr.sync_block.__init__(self,
            name="blsd_enc_b",
            in_sig=None,
            out_sig=[numpy.uint8,numpy.uint8,numpy.uint8,numpy.uint8,numpy.uint8,numpy.uint8,numpy.uint8])
        global d_kA1
        global d_kB1
        global d_k2
        global d_NA1
        global d_NB1
        global d_N2
        global d_M
        global d_packet_num
        global d_num_packets
        
        d_kA1 = kA1
        d_kB1 = kB1
        d_k2 = k2
        d_NA1 = NA1
        d_NB1 = NB1
        d_N2 = N2
        d_M = M    
        d_packet_num = 0
        d_num_packets = num_packets

	#self.set_relative_rate(pow(2,10))
	#self.set_min_output_buffer(1536)
	#self.set_min_output_items(1536)

    def work(self, input_items, output_items):
        out_cA = output_items[0]
        out_cB = output_items[1]
        out_bA1 = output_items[2]
        out_bB1 = output_items[3]
        out_b2 = output_items[4]
        out_bA = output_items[5]
        out_bB = output_items[6]

        global d_kA1
        global d_kB1
        global d_k2
        global d_NA1
        global d_NB1
        global d_N2
        global d_M
        global d_packet_num
        kA1 = d_kA1
        kB1 = d_kB1
        k2 = d_k2
        NA1 = d_NA1
        NB1 = d_NB1
        N2 = d_N2
        M = d_M

        bA1 = numpy.random.randint(2, size=(M, kA1)) # Data in Source A
        bA2 = numpy.random.randint(2, size=(M, k2))

        bB1 = numpy.random.randint(2, size=(M, kB1)) # Data in Source B
        bB2 = numpy.random.randint(2, size=(M, k2))

        if M == 1:
          bA1 = bA1.flatten()
          bA2 = bA2.flatten()
          bB1 = bB1.flatten()
          bB2 = bB2.flatten()

        b2 = bB2^bA2 # Hierarchical data
        bAB = numpy.hstack([bA1, bB1, b2]) # Concatenate, bA1, bB1, b2, M*(kA1+kB1+k2)

        # Loading LDPC code, only a certain family of the LDPC codes is available
        naA1 = '/home/gnuradio/gnuradio_trlcode/gr-trl/examples/ofdm_bsld/Pfiles/LDPC_lib/LDPC_library/LDPC_{0}_{1}'.format(str(NA1),str(NA1-kA1)) 
        H = LDPC(NA1, kA1)
        H.load_o(naA1) 
        cA1 = H.encode(bA1) # M*NA1
        cB1 = H.encode(bB1) # M*NB1
        cA2 = H.encode(bA2) # M*N2
        cB2 = H.encode(bB2) #M*N2

        H.prepare_decoder_C()

        if M == 1:
          cA1 = cA1.flatten()
          cA2 = cA2.flatten()
          cB1 = cB1.flatten()
          cB2 = cB2.flatten()

        # Definition of the hierarchical mapping
        eXmap_file = '/home/gnuradio/gnuradio_trlcode/gr-trl/examples/ofdm_bsld/eXmap_111'
        X = eXmap()
        X.load_extend(eXmap_file)

        nbitsA1 = X.hier_str_sizes[0] # Number of bits per symbol in stream A1
        nbitsB1 = X.hier_str_sizes[1] # Number of bits per symbol in stream B1
        nbits2  = X.hier_str_sizes[2] # Number of bits per symbol in both streams A2 and B2
        nbitsA  = nbitsA1 + nbits2
        nbitsB  = nbitsB1 + nbits2
        nbits = nbitsA + nbitsB
        if (M == 1):
          num_symbA = len(cA1) / nbitsA1 # Number of channel symbols from source A
          num_symbB = len(cB1) / nbitsB1 # Number of channel symbols from source B
        else:
          num_symbA = cA1.shape[1] / nbitsA1
          num_symbB = cB1.shape[1] / nbitsB1

        #print 'bits/symb: A1 = ', nbitsA1, ', B1 = ', nbitsB1, ', 2 = ', nbits2, ', A = ', nbitsA, 'B = ', nbitsB, ', total = ', nbits
        #print 'num_symbA = ', num_symbA, ', num_symbB = ', num_symbB

        # Reshaping source streams for the signal space mapping (one symbol from each stream per one channel
        # symbol)
        cAr = numpy.hstack([cA1.reshape(M * num_symbA, nbitsA1), cA2.reshape(M * num_symbA, nbits2)])
        cA = cAr.reshape(M, num_symbA * nbitsA) 
        cBr = numpy.hstack([cB1.reshape(M * num_symbB, nbitsB1), cB2.reshape(M * num_symbB, nbits2)])
        cB = cBr.reshape(M, num_symbB * nbitsB)

        bAr = np.hstack([bA1.reshape(M * kA1, nbitsA1), bA2.reshape(M * k2, nbits2)])
        bA = bAr.reshape(M, kA1 * nbitsA)
        bBr = np.hstack([bB1.reshape(M * kB1, nbitsB1), bB2.reshape(M * k2, nbits2)])
        bB = bBr.reshape(M, kB1 * nbitsB)

        # Reshape to continuous stream
        cAm = cA.reshape(M * num_symbA * nbitsA)
        cBm = cB.reshape(M * num_symbB * nbitsB)

        bAm = bA.reshape(M * kA1 * nbitsA) # Only used for o/p
        bBm = bB.reshape(M * kB1 * nbitsB) # Only used for o/p

	#tmpbA1 = bA1.flatten()
	#tmpbA2 = bA2.flatten()
	#tmpbB1 = bB1.flatten()
	#tmpbB2 = bB2.flatten()
	
	#f_out = open('bA1_bA2.txt','w',0)
	#f_out.write('bA1,bA2.\n')
	#for idx in range(0,len(tmpbA1)):
	#	f_out.write( str(tmpbA1[idx]) + ', ' + str(tmpbA2[idx]) + '\n' );

	#f_out = open('bB1_bB2.txt','w',0)
	#f_out.write('bB1,bB2.\n')
	#for idx in range(0,len(tmpbB1)):
	#	f_out.write( str(tmpbB1[idx]) + ', ' + str(tmpbB2[idx]) + '\n' );

        # Add packet_len tags
        self.add_item_tag(0, d_packet_num*len(cAm), pmt.string_to_symbol("packet_len"), pmt.from_long(len(cAm)/8) )
        self.add_item_tag(1, d_packet_num*len(cAm), pmt.string_to_symbol("packet_len"), pmt.from_long(len(cAm)/8) )

        #print 'bA1 = ', bA1
        #print 'bA2 = ', bA1
        #print 'cA1 = ', cA1
        #print 'cA2 = ', cA2
        #print 'cA = ', cA
        #print 'cB = ', cB

        #print 'out_bA.shape', out_bA.shape
        #print 'bAm.shape', bAm.shape

        bA1m = bA1.flatten()
        bB1m = bB1.flatten()
        b2m = b2.flatten()

        print 'len(cAm) = ', len(cAm), ', len(out_cA) = ', len(out_cA)
	
        out_cA[0:len(cAm)] = cAm
        out_cB[0:len(cBm)] = cBm
        out_bA1[0:len(bA1m)] = bA1m
        out_bB1[0:len(bB1m)] = bB1m
        out_b2[0:len(b2m)] = b2m
        out_bA[0:len(bAm)] = bAm
        out_bB[0:len(bBm)] = bBm
        self.produce(0,len(cAm))   
        self.produce(1,len(cBm))   
        self.produce(2,len(bA1m))   
        self.produce(3,len(bB1m))   
        self.produce(4,len(b2m))   
        self.produce(5,len(bAm))   
        self.produce(6,len(bBm))   

        #print 'out_cA = ', cAm, ', len(out_cA) = ', len(cAm)
        #print 'out_cB = ', cBm, ', len(out_cB) = ', len(cBm)

        print 'packet_num = ', d_packet_num
	#time.sleep(0.01)
        d_packet_num = d_packet_num + 1
        if (d_packet_num >= d_num_packets and d_num_packets > 0):
            return -1
        else:
            return -2