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

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


class srs4_rx(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, decrypt=False, auth=False, rolloff=0.2, swtol=0, constellation='QPSK'):
        gr.hier_block2.__init__(
            self, "SRS-4 Receiver",
            gr.io_signature(1, 1, gr.sizeof_gr_complex),
            gr.io_signature(1, 1, gr.sizeof_gr_complex),
        )
        self.message_port_register_hier_out("csp_out")
        self.message_port_register_hier_out("ip_out")

        # 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.decrypt = decrypt
        self.auth = auth
        self.rolloff = rolloff
        self.swtol = swtol

        # Variables
        self.sps = 4
        self.max_freq_offset_hz = 10e3
        self.sym_rate_ppm = 1000
        self.asm = 0x1acffc1d
        self.ss_loopbw = 5e-4
        self.ss_tedgain = 0.13621019192947206
        self.nfilts = 32

        # 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()
        elif self.modulation == '8PSK':
            self.constellation = satlab.ccsds_constellation_8psk()
        else:
            raise ValueError('Invalid constellation, must be either BPSK, QPSK or 8PSK')

        # Calculate soft decision LUTs
        lut_precision = 8
        self.constellation.gen_soft_dec_lut(lut_precision)

        # RRC polyphase matched filter
        self.rrc_filt = firdes.root_raised_cosine(self.nfilts, self.nfilts * self.rate * self.sps, self.rate, self.rolloff, int(5 * self.sps * self.nfilts))

        # Convolutional decoder object
        self.cc_decoder = fec.cc_decoder.make(128, 7, 2, [79, -109], 0, -1, fec.CC_STREAMING, False)

        # Blocks
        self.agc = analog.agc_cc(1e-4, 1.0, 1.0)
        self.agc.set_max_gain(65536)
        self.downsampler = filter.mmse_resampler_cc(0, (self.samp_rate / self.rate) / self.sps)
        self.symsync = digital.symbol_sync_cc(
            digital.TED_GARDNER,
            self.sps,
            self.ss_loopbw,
            1.0,
            self.ss_tedgain,
            self.sps * (self.sym_rate_ppm / 1e6),
            1,
            self.constellation,
            digital.IR_PFB_MF,
            self.nfilts,
            self.rrc_filt)
        self.costas = digital.costas_loop_cc((2 * math.pi / 2000), len(self.constellation.points()), False)
        self.multiply_const = blocks.multiply_const_cc(math.e ** (-1j * math.pi / 8) if (self.modulation == '8PSK') else 1)
        self.phase_stepper = satlab.phase_stepper(int(self.size * 8 * 10), 2 * math.pi / self.constellation.rotational_symmetry())
        self.demapper = digital.constellation_soft_decoder_cf(self.constellation)
        self.fec_decoder = fec.extended_decoder(decoder_obj_list=self.cc_decoder, threading='capillary', ann=None, puncpat='11', integration_period=10000)
        self.slicer = digital.binary_slicer_fb()
        self.swdetect = satlab.correlate_access_code_tag_fixed('{:032b}'.format(self.asm), self.swtol, 'packet_len', satlab.frame_size(self.size, self.crc, self.rs, False, self.idnum, self.decrypt, self.auth) * 8)
        self.unpacked_to_packed = blocks.unpacked_to_packed_bb(1, gr.GR_MSB_FIRST)
        self.tagged_stream_multiply_length = blocks.tagged_stream_multiply_length(gr.sizeof_char * 1, 'packet_len', 1 / 8.0)
        self.tagged_stream_to_pdu = blocks.tagged_stream_to_pdu(blocks.byte_t, 'packet_len')
        self.decoder = satlab.srs3_decode(self.size + 2 + (2 if self.idnum > 0 else 0) + (12 if self.decrypt or self.auth else 0) + (16 if self.auth else 0), self.crc, False, self.rand, self.rs, False)
        self.decryptor = satlab.srs3_decrypt(self.key, self.decrypt, self.auth)
        self.deframer = satlab.srs3_deframe(self.size, self.idnum)
        self.rate_tavg_sig = blocks.probe_signal_f()
        self.freq_err_sig = blocks.probe_signal_f()
        self.null_sink = blocks.null_sink(gr.sizeof_float * 1)

        # Limit Costas loop frequency range
        max_freq_offset_rad_per_sym = self.max_freq_offset_hz * 2 * math.pi / self.rate
        self.costas.set_min_freq(-max_freq_offset_rad_per_sym)
        self.costas.set_max_freq(max_freq_offset_rad_per_sym)

        # Connect main demodulation chain
        self.connect(self,
                     self.agc,
                     self.downsampler,
                     self.symsync,
                     self.costas,
                     self.multiply_const,
                     self.phase_stepper,
                     self.demapper,
                     self.fec_decoder if self.cc else self.slicer,
                     self.swdetect,
                     self.unpacked_to_packed,
                     self.tagged_stream_multiply_length,
                     self.tagged_stream_to_pdu)

        # Connect message blocks
        self.msg_connect((self.tagged_stream_to_pdu, 'pdus'), (self.decoder, 'decode'))
        self.msg_connect((self.tagged_stream_to_pdu, 'pdus'), (self.phase_stepper, 'kick'))
        self.msg_connect((self.decoder, 'output'), (self.decryptor, 'decrypt'))
        self.msg_connect((self.decryptor, 'output'), (self.deframer, 'deframe'))
        self.msg_connect((self.deframer, 'csp'), (self, 'csp_out'))
        self.msg_connect((self.deframer, 'ip'), (self, 'ip_out'))

        # Connect monitor signals
        self.connect((self.symsync, 1), (self.null_sink, 0))
        self.connect((self.symsync, 2), (self.null_sink, 1))
        self.connect((self.symsync, 3), self.rate_tavg_sig)
        self.connect((self.costas, 1), self.freq_err_sig)

        # Constellation debug output
        self.connect((self.costas, 0), (self, 0))

    def get_rate(self):
        return self.rate

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

    def get_size(self):
        return self.size

    def set_size(self, size):
        self.size = size
        self.phase_stepper.set_timeout(int(self.size * 8 * 10))
        self.swdetect.set_packet_len(satlab.frame_size(self.size, self.crc, self.rs, False, self.idnum, self.decrypt, self.auth) * 8)
        self.decoder.set_size(self.size + 2 + (2 if self.idnum else 0) + (12 if self.decrypt or self.auth else 0) + (16 if self.auth else 0))
        self.deframer.set_size(self.size)

    def get_id(self):
        return self.idnum

    def set_id(self, idnum):
        self.idnum = idnum
        self.swdetect.set_packet_len(satlab.frame_size(self.size, self.crc, self.rs, False, self.idnum, self.decrypt, self.auth) * 8)
        self.decoder.set_size(self.size + 2 + (2 if self.idnum > 0 else 0) + (12 if self.decrypt or self.auth else 0) + (16 if self.auth else 0))
        self.deframer.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.swdetect.set_packet_len(satlab.frame_size(self.size, self.crc, self.rs, False, self.idnum, self.decrypt, self.auth) * 8)
        self.decoder.set_rs(self.rs)

    def get_crc(self):
        return self.crc

    def set_crc(self, crc):
        self.crc = crc
        self.swdetect.set_packet_len(satlab.frame_size(self.size, self.crc, self.rs, False, self.idnum, self.decrypt, self.auth) * 8)
        self.decoder.set_crc(self.crc)

    def get_rand(self):
        return self.rand

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

    def get_key(self):
        return self.key

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

    def get_auth(self):
        return self.auth

    def set_auth(self, auth):
        self.auth = auth
        self.swdetect.set_packet_len(satlab.frame_size(self.size, self.crc, self.rs, False, self.idnum, self.decrypt, self.auth) * 8)
        self.decoder.set_size(self.size + 2 + (2 if self.idnum > 0 else 0) + (12 if self.decrypt or self.auth else 0) + (16 if self.auth else 0))
        self.decryptor.set_auth(self.auth)

    def get_decrypt(self):
        return self.decrypt

    def set_decrypt(self, decrypt):
        self.decrypt = decrypt
        self.swdetect.set_packet_len(satlab.frame_size(self.size, self.crc, self.rs, False, self.idnum, self.decrypt, self.auth) * 8)
        self.decoder.set_size(self.size + 2 + (2 if self.idnum > 0 else 0) + (12 if self.decrypt or self.auth else 0) + (16 if self.auth else 0))
        self.decryptor.set_decrypt(self.decrypt)

    def get_rolloff(self):
        return self.rolloff

    def set_rolloff(self, rolloff):
        self.rolloff = rolloff
        self.rrc_filt = firdes.root_raised_cosine(self.nfilts, self.nfilts * self.rate * self.sps, self.rate, self.rolloff, int(5 * self.sps * self.nfilts))
        # symsync can not update taps runtime

    def get_swtol(self):
        return self.swtol

    def set_swtol(self, swtol):
        self.swtol = swtol
        self.swdetect.set_threshold(self.swtol)

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

    def get_constellation_points(self):
        return self.constellation.points()

    def get_freq_err(self):
        return self.freq_err_sig.level() * self.rate / (2 * math.pi)

    def get_sym_err(self):
        return -(self.rate_tavg_sig.level() - self.sps) / self.sps * 1e6

    def get_bit_rate(self):
        frame_size = satlab.frame_size(self.size, self.crc, self.rs, self.cc, self.idnum, self.decrypt, 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
