import ctypes
import datetime
import ipaddress
import socket
import struct
import time

import numpy as np

import data_structures as msg_types
from data_structures import CpiHeader

TIMING_ENGINE_ADDR = 0x40051000
DIG_RX_ADDR = 0x20000000
DIG_RX_STRIDE = 0x10000
WAVEFORM_GEN_ADDR = 0x40053000

NUM_RX = 2


ADC_SAMPLE_RATE = 187.5e6
DAC_SAMPLE_RATE = 187.5e6

def form_chirp(pulsewidth, bw, sample_rate, win=None, ):

    n = int(np.round(pulsewidth * sample_rate))
    d_t = 1 / sample_rate
    f = np.linspace(-bw/2, bw/2, n)

    phi = np.cumsum(2 * np.pi * f * d_t)

    x = np.exp(-1j * phi)

    return x.astype(np.complex64)

class RadarManager:
    def __init__(self,
                 host="192.168.1.200",
                 port=5001):

        self.host = host
        self.port = port
        self.s = None
        self.CONNECTED = False

        self.connect()

        # Update UDP packet size
        self.packet_size = 4096
        self.axi_write_register(0x4005001C, self.packet_size)

        self.reset_10g_udp()


        self.stop_running()

    def connect(self):
        self.CONNECTED = False

        retry_cnt = 0
        while True:
            try:
                self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.s.settimeout(5)
                self.s.connect((self.host, self.port))
                self.s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                self.s.settimeout(5)
                self.CONNECTED = True
                break
            except socket.error as e:
                print("Connection Error %s" % e)
                # print("Connection Timeout, trying again %d" % retry_cnt)
                self.CONNECTED = False
                retry_cnt += 1
                if retry_cnt >= 1:
                    break

        return self.CONNECTED

    def disconnect(self):
        self.s.close()
        self.CONNECTED = False

    def is_connected(self):
        return self.CONNECTED

    def send_bytes(self, msg_bytes):
        self.s.sendall(msg_bytes)

    def send_message(self, msg, update_hdr=True, enable_ack=True, wait_for_ack=True, timeout=10):
        self.s.settimeout(timeout)
        ret = True
        msg.header.flags = 0

        # Request a completion ACK
        if enable_ack:
            msg.header.flags = msg_types.HDR_FLAG_REQ_ACK

        # Serialize message and send
        self.send_bytes(bytes(msg))

        # Need to wait for response
        if wait_for_ack:
            recv_bytes, err = self.get_response(ctypes.sizeof(msg_types.AckType))
            resp = msg_types.AckType.from_buffer_copy(recv_bytes)
            ret = True

        return ret

    def get_response(self, size):
        recv_data = bytearray()
        while len(recv_data) < size:
            try:
                recv_data += self.s.recv(size, 0)
            except socket.timeout as e:
                return 0, -1
        return recv_data, 0

    def axi_write_register(self, address, data):
        # Make sure address is word aligned
        address -= (address % 4)

        # Form message
        msg = msg_types.WriteRegType()
        msg.address = address
        msg.data = data

        self.send_message(msg)

        return

    def axi_write_register_burst(self, address, data):
        # Make sure address is word aligned
        address -= (address % 4)

        # Form message
        msg = msg_types.WriteRegBurstType()
        msg.address = address
        msg.length = len(data)
        for i in range(len(data)):
            msg.data[i] = data[i]

        self.send_message(msg)

        return

    def axi_read_register(self, address):
        # Make sure address is word aligned
        address -= (address % 4)

        # Form message
        msg = msg_types.ReadRequestType()
        msg.address = address

        self.send_message(msg, enable_ack=False, wait_for_ack=False)

        # Get response
        recv_bytes, _ = self.get_response(ctypes.sizeof(msg_types.ReadResponseType))
        resp = msg_types.ReadResponseType.from_buffer_copy(recv_bytes)

        return resp.data

    def set_dac_nco(self, channel, frequency):
        # Form message
        msg = msg_types.DacNcoConfigType()
        msg.channel = channel
        msg.frequency = frequency

        self.send_message(msg)

        return

    def set_adc_nco(self, channel, frequency):
        # Form message
        msg = msg_types.AdcNcoConfigType()
        msg.channel = channel
        msg.frequency = frequency

        self.send_message(msg)

        return

    def rf_spi_write(self, dev_sel, num_bits, data):
        # Form message
        msg = msg_types.RfSpiWriteType()
        msg.dev_sel = dev_sel
        msg.num_bits = num_bits
        msg.data = data

        self.send_message(msg)

        return


    def load_waveform(self, ch, amp, bw, pw):
        addr = 0x0010000 + 0x0010000 * ch
        print('Load', hex(addr))
        num_samples = pw
        wf = form_chirp(pw, bw, 1)
        wf = wf * amp

        bram_address = addr

        iq = wf * 0x7FFF
        iq_real = iq.real.astype(np.uint16)
        iq_imag = iq.imag.astype(np.uint16)
        iq_real = iq_real.astype(np.uint32)
        iq_imag = iq_imag.astype(np.uint32)
        data = iq_real | (iq_imag << 16)

        num_bursts = num_samples / msg_types.MAX_BURST_LENGTH
        num_bursts = int(np.ceil(num_bursts))

        for i in range(num_bursts):
            start_ind = i * msg_types.MAX_BURST_LENGTH
            stop_ind = start_ind + msg_types.MAX_BURST_LENGTH
            stop_ind = min(stop_ind, num_samples)
            burst_data = data[start_ind:stop_ind]
            self.axi_write_register_burst(bram_address + i * 4 * msg_types.MAX_BURST_LENGTH, burst_data)

    def reset_10g_udp(self):
        val = self.axi_read_register(0x40050008)
        self.axi_write_register(0x40050008, val | (1 << 15))
        self.axi_write_register(0x40050008, val)
        time.sleep(2)

    def setup_cpi_header(self, pri, inter_cpi, num_pulses, num_samples, start_sample,
                      tx_num_samples, tx_start_sample, rx_lo_offset, tx_lo_offset):
        # Load Up header BRAM
        header = CpiHeader()
        header.sync = 0xAABBCCDD
        header.num_pulses = num_pulses
        header.num_samples = num_samples
        header.start_sample = start_sample
        header.tx_num_samples = tx_num_samples
        header.tx_start_sample = tx_start_sample
        header.pri = pri
        header.inter_cpi = inter_cpi
        header.tx_lo_offset = tx_lo_offset
        header.rx_lo_offset = rx_lo_offset

        data = np.frombuffer(bytes(header), dtype=np.uint32)
        self.axi_write_register_burst(TIMING_ENGINE_ADDR + 0xC00, data)

    def setup_timing_engine(self, pri, num_pulses, inter_cpi):
        self.axi_write_register(TIMING_ENGINE_ADDR + 0x4, pri - 1)
        self.axi_write_register(TIMING_ENGINE_ADDR + 0x8, num_pulses)
        self.axi_write_register(TIMING_ENGINE_ADDR + 0x10, inter_cpi - 1)

    def setup_rx(self, num_samples, start_sample):
        for i in range(NUM_RX):
            self.axi_write_register(DIG_RX_ADDR + i*DIG_RX_STRIDE + 0x4, num_samples >> 2)
            self.axi_write_register(DIG_RX_ADDR + i*DIG_RX_STRIDE + 0x8, start_sample >> 2)

            # Setup RX Strobe
            self.axi_write_register(TIMING_ENGINE_ADDR + 0x88 + i * 8, start_sample)
            self.axi_write_register(TIMING_ENGINE_ADDR + 0x8C + i * 8, num_samples)

    def setup_tx(self, num_samples, start_sample):
        self.axi_write_register(WAVEFORM_GEN_ADDR + 0x4, num_samples >> 2)
        self.axi_write_register(WAVEFORM_GEN_ADDR + 0x8, start_sample >> 2)

        # Setup TX Strobe
        self.axi_write_register(TIMING_ENGINE_ADDR + 0x80, start_sample)
        self.axi_write_register(TIMING_ENGINE_ADDR + 0x84, num_samples)

    def start_running(self):
        for i in range(NUM_RX):
            self.axi_write_register(DIG_RX_ADDR + i*DIG_RX_STRIDE + 0x0, 0)  # RX Reset
        self.axi_write_register(WAVEFORM_GEN_ADDR + 0x0, 0)  # TX Reset
        self.axi_write_register(TIMING_ENGINE_ADDR + 0x0, 0)  # Timing Reset (do this last)

    def stop_running(self):
        self.axi_write_register(TIMING_ENGINE_ADDR + 0x0, 1)  # Timing Engine Reset
        for i in range(NUM_RX):
            self.axi_write_register(DIG_RX_ADDR + i*DIG_RX_STRIDE + 0x0, 1)  # RX Reset
        self.axi_write_register(WAVEFORM_GEN_ADDR + 0x0, 1)  # TX Reset

    def setup_rf_attenuators(self, rf_atten):
        self.rf_spi_write((1 << 0), 6, rf_atten[0])  # TX0 RF (ADRF5730)
        self.rf_spi_write((1 << 1), 6, rf_atten[1])  # TX1 RF (ADRF5730)

        self.rf_spi_write((1 << 2), 6, rf_atten[2])  # RX0 RF (ADRF5721)
        self.rf_spi_write((1 << 3), 6, rf_atten[3])  # RX0 IF (HMC624)

        self.rf_spi_write((1 << 4), 6, rf_atten[4])  # RX1 RF (ADRF5721)
        self.rf_spi_write((1 << 5), 6, rf_atten[5])  # RX1 IF (HMC624)


    def configure_cpi(self, pri, inter_cpi, num_pulses, num_samples, start_sample,
                      tx_num_samples, tx_start_sample, rx_lo_offset, tx_lo_offset):
        self.load_waveform(0, 1, 0.1, tx_num_samples)
        self.load_waveform(1, 1, 0.1, tx_num_samples)

        rf_atten = [1, 2, 3, 4, 5, 6]
        self.setup_rf_attenuators(rf_atten)

        # DAC at 5.25 GHz is in second nyquist
        # ADC would be in 3rd nyquist
        f_dac = 9e9
        f_adc = 3e9
        tx_lo = 5.25e9 % f_dac
        rx_lo = 5.25e9 % f_adc
        for i in range(4):
            self.set_adc_nco(i, rx_lo)
        self.set_dac_nco(0, tx_lo)
        self.set_dac_nco(1, tx_lo)
        self.set_dac_nco(2, tx_lo + tx_lo_offset)
        self.set_dac_nco(3, tx_lo + rx_lo_offset)
        self.setup_timing_engine(pri, num_pulses, inter_cpi)
        self.setup_rx(num_samples, start_sample)
        self.setup_tx(tx_num_samples, tx_start_sample)
        self.setup_cpi_header(pri, inter_cpi, num_pulses, num_samples, start_sample,
                      tx_num_samples, tx_start_sample, rx_lo_offset, tx_lo_offset)