import ctypes
import time
import numpy as np
from matplotlib import pyplot as plt

import data_structures
import radar_manager
from data_recorder import DataRecorder

# Give 10g eth interface an ip and set MTU for better performance
# sudo ifconfig enp5s0f0 192.168.2.10 up mtu 5000
# sudo ifconfig enp5s0f1 192.168.3.10 up mtu 5000

# Note that increases the size of rmem_max in the linux kernel improves performance for data recording
# this can be done witht the following terminal command
# sudo sysctl -w net.core.rmem_max=1048576

def db20(x):
    return 20*np.log10(np.abs(x))


def db20n(x):
    x = db20(x)
    x = x - np.max(x)
    return x


def main():
    print('Hello')

    radar = radar_manager.RadarManager()

    clk = radar_manager.TIMING_ENGINE_FREQ

    freqs = np.array([16, 21, 13, 3.25, 3.5, 5, 2])*1e9
    pri_lsb = 16e-9
    print(freqs * pri_lsb)

    # Test AD9081 Reg Access
    print(hex(radar.ad9081_read_reg(0x0A0A)))
    radar.ad9081_write_reg(0x0A0A, 0x60)
    print(hex(radar.ad9081_read_reg(0x0A0A)))


    # CPI Parameters (timing values are in clk ticks)
    num_pulses = 128
    # Should be multiple of udp packet size, currently 4096 bytes, or 1024 samples
    num_samples = 5000
    start_sample = 2000
    tx_num_samples = 1024
    tx_start_sample = start_sample
    pri = int(.0004 * clk)
    print(pri)
    inter_cpi = 50
    tx_lo_offset = 10e6
    rx_lo_offset = 0

    pri_float = pri / clk

    print('PRI', pri_float, 'PRF', 1 / pri_float)
    print('Expected Data Rate', num_samples * 4 / pri_float / 1e6)



    recorder0 = DataRecorder("192.168.2.128", 1234, packet_size=radar.packet_size)
    recorder1 = DataRecorder("192.168.3.128", 1235, packet_size=radar.packet_size)
    recorder0.start_recording('test0.bin', True)
    recorder1.start_recording('test1.bin', True)

    radar.configure_cpi(pri, inter_cpi, num_pulses, num_samples, start_sample,
                        tx_num_samples, tx_start_sample, rx_lo_offset, tx_lo_offset)

    print('Start Running')
    radar.start_running()
    # Let it run for a bit
    time.sleep(2)
    # Stop running
    radar.stop_running()
    # Stop the data recorder
    recorder0.stop_recording()
    recorder1.stop_recording()

    # Parse some data

    # Find header, recording buffer could have wrapped depending on data rate and how long we ran for
    recorders = [recorder0, recorder1]
    for recorder in recorders:
        headers = []
        offset = 0
        plot_recorder = recorder
        hdr_sync = False
        while not hdr_sync:
            data = plot_recorder.buffer[offset:offset + 4]
            sync_word = np.frombuffer(data, dtype=np.uint32)[0]
            if sync_word == 0xAABBCCDD:
                hdr_sync = True
                print('Header found at offset', offset)
            else:
                offset += 4

        num_cpi = 1
        for i in range(num_cpi):
            # Get Header
            data = plot_recorder.buffer[offset:offset + ctypes.sizeof(data_structures.CpiHeader)]
            offset += ctypes.sizeof(data_structures.CpiHeader)
            headers.append(data_structures.CpiHeader.from_buffer_copy(data))
            num_pulses = headers[i].num_pulses
            num_samples = headers[i].num_samples

            # Get CPI
            data_size = num_pulses * num_samples * 4
            data = plot_recorder.buffer[offset:offset + data_size]
            offset += data_size

        # Check some header fields
        cpi_times = np.array([x.system_time for x in headers]) / 187.5e6
        pps_frac = np.array([x.pps_frac_sec for x in headers]) / 187.5e6
        pps_sec = np.array([x.pps_sec for x in headers])
        utc_time = pps_sec + pps_frac
        print(pri, inter_cpi, num_pulses * pri + inter_cpi)
        print(cpi_times - cpi_times[0])
        print(pps_frac)
        print(pps_sec - pps_sec[0])

        # Plot last CPI
        data2 = np.frombuffer(data, dtype=np.int16)
        i = data2[0::2]
        q = data2[1::2]
        iq = i + 1j * q
        iq = iq.reshape(-1, num_samples)
        iq = iq + 1e-15

        vmin = -60
        vmax = 0

        fid, axs = plt.subplots(3)
        axs[0].plot(iq.T.real, '-')
        axs[0].plot(iq.T.imag, '--')
        axs[0].grid()

        # axs[1].imshow(db20n(iq), aspect='auto', interpolation='nearest', vmin=vmin, vmax=vmax)
        axs[1].imshow(iq.real, aspect='auto', interpolation='nearest')
        axs[1].set_ylabel('Pulse Count')
        axs[1].set_xlabel('Sample Count')

        iq_freq = np.fft.fftshift(np.fft.fft(iq, axis=1), axes=1)
        freq_axis = (np.arange(num_samples)/num_samples - 0.5) * radar_manager.BASEBAND_SAMPLE_RATE / 1e6
        axs[2].plot(freq_axis, db20n(iq_freq.T))
        axs[2].grid()


    plt.show()


if __name__ == '__main__':

    main()