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

from gnuradio import gr, blocks, filter, digital, fec
from gnuradio.filter import firdes
from gnuradio import satlab


class srs4_tx(gr.hier_block2):
    def __init__(self, samp_rate=20e6, rate=1000000, size=217, idnum=0, cc=True, rs=True, crc=True, rand=True, key='\x00'*32, encrypt=False, auth=False, rolloff=0.2, constellation='QPSK', idleframes=10000, max_in_flight=100):
        gr.hier_block2.__init__(
            self, "SRS-4 Transmitter",
            gr.io_signature(0, 0, 0),
            gr.io_signature(1, 1, gr.sizeof_gr_complex),
        )
        self.message_port_register_hier_in("csp_in")
        self.message_port_register_hier_in("ip_in")

        # Parameters
        self.samp_rate = samp_rate
        self.rate = rate
        self.size = size
        self.idnum = idnum
        self.cc = cc
        self.rs = rs
        self.crc = crc
        self.rand = rand
        self.key = key
        self.encrypt = encrypt
        self.auth = auth
        self.rolloff = rolloff
        self.idleframes = idleframes
        self.max_in_flight = max_in_flight

        # Variables
        self.sps = 4
        self.papr_db = 5.6

        # Selected constellation
        self.modulation = constellation.upper()
        if self.modulation == 'BPSK':
            self.constellation = satlab.ccsds_constellation_bpsk()
        elif self.modulation == 'QPSK':
            self.constellation = satlab.ccsds_constellation_qpsk()
        else:
            raise ValueError('Invalid constellation, must be either BPSK or QPSK')

        # RRC pulse shaping filter
        self.rrc_filt = firdes.root_raised_cosine(self.sps, self.sps, 1, self.rolloff, int(9 * self.sps))

        # Convolutional encoder object
        self.cc_encoder = fec.cc_encoder_make(128, 7, 2, [79, -109], 0, fec.CC_STREAMING, False)

        # Blocks
        self.enframer = satlab.srs3_enframe(self.size, self.idnum)
        self.encrypter = satlab.srs3_encrypt(self.key, self.encrypt, self.auth)
        self.idle_frame_gen = satlab.idle_frame_gen(satlab.frame_size(self.size, idnum=self.idnum, crypt=self.encrypt, auth=self.auth, crc=False, rs=False, cc=False, cc_term=False, include_asm=False), self.idleframes, self.max_in_flight)
        self.tag_to_trigger = satlab.tag_to_trigger('packet_len')
        self.encoder = satlab.srs3_encode(0, self.crc, False, self.rand, self.rs, False)
        self.pdu_to_tagged_stream = blocks.pdu_to_tagged_stream(blocks.byte_t, 'packet_len')
        self.packed_to_unpacked = blocks.packed_to_unpacked_bb(1, gr.GR_MSB_FIRST)
        self.fec_encoder = fec.extended_encoder(encoder_obj_list=self.cc_encoder, threading='capillary', puncpat='11')
        self.pack_k_bits = blocks.pack_k_bits_bb(self.constellation.bits_per_symbol())
        self.mapper = digital.chunks_to_symbols_bc(self.constellation.points(), self.constellation.dimensionality())
        self.rrc = filter.interp_fir_filter_ccc(self.sps, self.rrc_filt)
        self.rrc.declare_sample_delay(0)
        self.multiply_const = blocks.multiply_const_cc(10**(-self.papr_db / 20))
        self.upsampler = filter.mmse_resampler_cc(0, self.sps / (self.samp_rate / self.rate))

        # Connect message blocks
        self.msg_connect((self, 'csp_in'), (self.enframer, 'csp'))
        self.msg_connect((self, 'ip_in'), (self.enframer, 'ip'))
        self.msg_connect((self.enframer, 'output'), (self.encrypter, 'encrypt'))
        self.msg_connect((self.encrypter, 'output'), (self.idle_frame_gen, 'frame'))
        self.msg_connect((self.idle_frame_gen, 'pdu'), (self.encoder, 'input'))
        self.msg_connect((self.tag_to_trigger, 'trigger'), (self.idle_frame_gen, 'trigger'))
        self.msg_connect((self.encoder, 'output'), (self.pdu_to_tagged_stream, 'pdus'))

        # Connect main demodulation chain
        self.connect(self.pdu_to_tagged_stream,
                     self.packed_to_unpacked)
        if cc:
            self.connect(self.packed_to_unpacked,
                         self.fec_encoder,
                         self.pack_k_bits)
        else:
            self.connect(self.packed_to_unpacked,
                         self.pack_k_bits)

        self.connect(self.pack_k_bits,
                     self.mapper,
                     self.rrc,
                     self.multiply_const,
                     self.upsampler, self)

        # Connect idle frame trigger
        self.connect(self.multiply_const, self.tag_to_trigger)

    def get_rate(self):
        return self.rate

    def set_rate(self, rate):
        self.rate = rate
        self.upsampler.set_resamp_ratio(self.sps / (self.samp_rate / self.rate))

    def get_size(self):
        return self.size

    def set_size(self, size):
        self.size = size
        self.idle_frame_gen.set_size(self.size + 2 + (2 if self.idnum else 0) + (12 if self.encrypt or self.auth else 0) + (16 if self.auth else 0))
        self.enframer.set_size(self.size)

    def get_id(self):
        return self.idnum

    def set_id(self, idnum):
        self.idnum = idnum
        self.idle_frame_gen.set_size(self.size + 2 + (2 if self.idnum else 0) + (12 if self.encrypt or self.auth else 0) + (16 if self.auth else 0))
        self.enframer.set_id(self.idnum)

    def get_cc(self):
        return self.cc

    def set_cc(self, cc):
        self.cc = cc

    def get_rs(self):
        return self.rs

    def set_rs(self, rs):
        self.rs = rs
        self.encoder.set_rs(self.rs)

    def get_crc(self):
        return self.crc

    def set_crc(self, crc):
        self.crc = crc
        self.encoder.set_crc(self.crc)

    def get_rand(self):
        return self.rand

    def set_rand(self, rand):
        self.rand = rand
        self.encoder.set_rand(self.rand)

    def get_key(self):
        return self.key

    def set_key(self, key):
        self.key = key
        self.encrypter.set_key(self.key)

    def get_auth(self):
        return self.auth

    def set_auth(self, auth):
        self.auth = auth
        self.idle_frame_gen.set_size(self.size + 2 + (2 if self.idnum else 0) + (12 if self.encrypt or self.auth else 0) + (16 if self.auth else 0))
        self.encrypter.set_auth(self.auth)

    def get_encrypt(self):
        return self.encrypt

    def set_encrypt(self, encrypt):
        self.encrypt = encrypt
        self.idle_frame_gen.set_size(self.size + 2 + (2 if self.idnum else 0) + (12 if self.encrypt or self.auth else 0) + (16 if self.auth else 0))
        self.encrypter.set_encrypt(self.encrypt)

    def get_rolloff(self):
        return self.rolloff

    def set_rolloff(self, rolloff):
        self.rolloff = rolloff
        self.rrc_filt = firdes.root_raised_cosine(self.sps, self.sps, 1, self.rolloff, int(9 * self.sps))
        self.rrc.set_taps(self.rrc_filt)

    def get_frames(self):
        return self.enframer.get_frames()

    def get_bit_rate(self):
        frame_size = satlab.frame_size(self.size, self.crc, self.rs, self.cc, self.idnum, self.encrypt, self.auth, cc_term=False, include_asm=True)
        return self.rate * self.constellation.bits_per_symbol() * self.size / frame_size

    def get_modulation(self):
        return self.modulation
