# Copyright 2019 The OpenRadar Authors. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import codecs
import socket
import struct
import serial
import time
import matplotlib.pyplot as plt
# import pyqtgraph as pg
# from mpl_toolkits.mplot3d import Axes3D
from enum import Enum
# from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
# from matplotlib.figure import Figure
# from matplotlib.lines import Line2D
# import matplotlib.cbook as cbook
import numpy as np

configFileName = "IWR6843_cfg.cfg"


class CMD(Enum):
    RESET_FPGA_CMD_CODE = '0100'
    RESET_AR_DEV_CMD_CODE = '0200'
    CONFIG_FPGA_GEN_CMD_CODE = '0300'
    CONFIG_EEPROM_CMD_CODE = '0400'
    RECORD_START_CMD_CODE = '0500'
    RECORD_STOP_CMD_CODE = '0600'
    PLAYBACK_START_CMD_CODE = '0700'
    PLAYBACK_STOP_CMD_CODE = '0800'
    SYSTEM_CONNECT_CMD_CODE = '0900'
    SYSTEM_ERROR_CMD_CODE = '0a00'
    CONFIG_PACKET_DATA_CMD_CODE = '0b00'
    CONFIG_DATA_MODE_AR_DEV_CMD_CODE = '0c00'
    INIT_FPGA_PLAYBACK_CMD_CODE = '0d00'
    READ_FPGA_VERSION_CMD_CODE = '0e00'

    def __str__(self):
        return str(self.value)


# MESSAGE = codecs.decode(b'5aa509000000aaee', 'hex')
CONFIG_HEADER = '5aa5'
CONFIG_STATUS = '0000'
CONFIG_FOOTER = 'aaee'
ADC_PARAMS = {
    'chirps': 96,
    'rx': 4,
    'tx': 3,
    'samples': 96,
    'IQ': 2,
    'bytes': 2
}
# STATIC
MAX_PACKET_SIZE = 2097152
# 默认数据包长度
BYTES_IN_PACKET = 1456
# DYNAMIC
BYTES_IN_FRAME = (ADC_PARAMS['chirps'] * ADC_PARAMS['rx'] * ADC_PARAMS['tx'] *
                  ADC_PARAMS['IQ'] * ADC_PARAMS['samples'] *
                  ADC_PARAMS['bytes'])
BYTES_IN_FRAME_CLIPPED = (BYTES_IN_FRAME // BYTES_IN_PACKET) * BYTES_IN_PACKET
PACKETS_IN_FRAME = BYTES_IN_FRAME / BYTES_IN_PACKET
PACKETS_IN_FRAME_CLIPPED = BYTES_IN_FRAME // BYTES_IN_PACKET
UINT16_IN_PACKET = BYTES_IN_PACKET // 2
UINT16_IN_FRAME = BYTES_IN_FRAME // 2
# a = []
# b = []


class DCA1000:
    """Software interface to the DCA1000 EVM board via ethernet.

    Attributes:
        static_ip (str): IP to receive data from the FPGA
        adc_ip (str): IP to send configuration commands to the FPGA
        data_port (int): Port that the FPGA is using to send data
        config_port (int): Port that the FPGA is using to read configuration commands from


    General steps are as follows:
        1. Power cycle DCA1000 and XWR1xxx sensor
        2. Open mmWaveStudio and setup normally until tab SensorConfig or use lua script
        3. Make sure to connect mmWaveStudio to the board via ethernet
        4. Start streaming data
        5. Read in frames using class

    Examples:
        >>> dca = DCA1000()
        >>> adc_data = dca.read(timeout=.1)
        >>> frame = dca.organize(adc_data, 128, 4, 256)

    """
    def __init__(self,
                 configfilename,
                 static_ip='192.168.33.30',
                 adc_ip='192.168.33.180',
                 data_port=4098,
                 config_port=4096):
        # Save network data
        # self.static_ip = static_ip
        # self.adc_ip = adc_ip
        # self.data_port = data_port
        # self.config_port = config_port

        # Create configuration and data destinations
        self.cfg_dest = (adc_ip, config_port)
        self.cfg_recv = (static_ip, config_port)
        self.data_recv = (static_ip, data_port)

        # Create sockets
        self.config_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
                                           socket.IPPROTO_UDP)
        self.data_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
                                         socket.IPPROTO_UDP)

        # Bind data socket to fpga
        self.data_socket.bind(self.data_recv)

        # Bind config socket to fpga
        self.config_socket.bind(self.cfg_recv)

        self.CLIport = serial.Serial('COM7', 115200)
        self.Dataport = serial.Serial('COM8', 921600)

        self.CfgFilename = configfilename

        self.data = []
        self.packet_count = []
        self.byte_count = []

        self.frame_buff = []

        self.curr_buff = None
        self.last_frame = None
        self.lost_packets = None

    def _serialConfig(self):
        # Raspberry pi
        # CLIport = serial.Serial('/dev/ttyACM0', 115200)
        # Dataport = serial.Serial('/dev/ttyACM1', 921600)

        # Windows
        # self.CLIport = serial.Serial('COM11', 115200)
        # self.Dataport = serial.Serial('COM12', 921600)

        config = [line.rstrip('\r\n') for line in open(self.CfgFilename)]
        for i in config:
            if i == '' or i[0] == "#" or i[0] == "%":
                continue
            self.CLIport.write((i + '\n').encode())
            print('sent->' + i)
            time.sleep(0.1)
            reply = self.CLIport.read(self.CLIport.in_waiting).decode()
            print(reply)
        return self.CLIport, self.Dataport

    def _sensor_start(self):
        self.CLIport.write("sensorStart\n".encode())
        print("sent->" + "sensorStart")
        time.sleep(0.1)
        reply = self.CLIport.read(self.CLIport.in_waiting).decode()
        print(reply)

    def _sensor_stop(self):
        self.CLIport.write("sensorStop\n".encode())
        print("sent->" + "sensorStop")
        time.sleep(0.1)
        reply = self.CLIport.read(self.CLIport.in_waiting).decode()
        print(reply)

    def start(self):
        # RECORD_START_CMD_CODE
        # 5a a5 05 00 00 00 aa ee
        print("->网口发送记录开始指令：\n{}".format(
            self._send_command(CMD.RECORD_START_CMD_CODE)))
        self._sensor_start()

        print('commands send -> Start ')

    def stop(self):
        self._sensor_stop()
        # RECORD_START_CMD_CODE
        # 5a a5 06 00 00 00 aa ee
        print("->网口发送记录终止指令：\n{}".format(
            self._send_command(CMD.RECORD_STOP_CMD_CODE)))
        self._close()

    def _close(self):
        # self.data_socket.close()
        self.config_socket.close()
        self.CLIport.close()
        self.Dataport.close()

    def configure(self):
        # # SYSTEM_CONNECT_CMD_CODE
        # # 5a a5 09 00 00 00 aa ee
        # print("->连接确认:")
        # print("{}".format(self.send_command(CMD.SYSTEM_CONNECT_CMD_CODE)))

        # CONFIG_FPGA_GEN_CMD_CODE
        # 5a a5 01 00 00 00 aa ee
        print("->重置FPGA：")
        print("{}".format(self._send_command(CMD.RESET_FPGA_CMD_CODE)))

        # CONFIG_FPGA_GEN_CMD_CODE
        # 5a a5 02 00 00 00 aa ee
        print("->重置AR：")
        print("{}".format(self._send_command(CMD.RESET_AR_DEV_CMD_CODE)))

        # CONFIG_FPGA_GEN_CMD_CODE
        # 5a a5 03 00 06 00 01 02 01 02 03 1e aa ee
        print("->配置FPGA：")
        print("{}".format(
            self._send_command(CMD.CONFIG_FPGA_GEN_CMD_CODE, '0600',
                               '01020102031e')))

        # CONFIG_PACKET_DATA_CMD_CODE
        # 5a a5 0b 00 06 00 be 05 35 0c 00 00 aa ee
        print("->配置PACKET：")
        print("{}".format(
            self._send_command(CMD.CONFIG_PACKET_DATA_CMD_CODE, '0600',
                               'c005350c0000')))

        # # READ_FPGA_VERSION_CMD_CODE
        # # 5a a5 0e 00 00 00 aa ee
        # print("->读取FPGA版本：")
        # print("{}".format(self.send_command_web(CMD.READ_FPGA_VERSION_CMD_CODE)))

    def close(self):
        """Closes the sockets that are used for receiving and sending data

        Returns:
            None

        """
        self.data_socket.close()
        self.config_socket.close()

    def read(self, timeout=1):
        """ Read in a single packet via UDP

        Args:
            timeout (float): Time to wait for packet before moving on

        Returns:
            Full frame as array if successful, else None

        """
        # Configure

        self.data_socket.settimeout(timeout)

        # Frame buffer
        ret_frame = np.zeros(UINT16_IN_FRAME, dtype=np.int16)

        # Wait for start of next frame
        while True:
            packet_num, byte_count, packet_data = self._read_data_packet()
            # a.append(packet_num) # 获取包数组
            # print(packet_num,byte_count)
            if (packet_num - 1) % PACKETS_IN_FRAME_CLIPPED == 0:
                packets_read = 1
                ret_frame[0:UINT16_IN_PACKET] = packet_data
                break

        # Read in the rest of the frame
        while True:
            packet_num, byte_count, packet_data = self._read_data_packet()
            # print(packet_num,byte_count)
            # b.append(packet_num) # 获取包数组
            packets_read += 1
            if (packet_num - 1) % PACKETS_IN_FRAME_CLIPPED == 0:
                self.lost_packets = PACKETS_IN_FRAME_CLIPPED - packets_read
                # 返回-1 基本不丢包
                # print("lost_packets:%d" %self.lost_packets)
                return ret_frame

            curr_idx = ((packet_num - 1) % PACKETS_IN_FRAME_CLIPPED)
            try:
                ret_frame[curr_idx * UINT16_IN_PACKET:(curr_idx + 1) *
                          UINT16_IN_PACKET] = packet_data
            except:
                pass

            if packets_read > PACKETS_IN_FRAME_CLIPPED:
                packets_read = 0

    def _send_command(self, cmd, length='0000', body='', timeout=1):
        """Helper function to send a single commmand to the FPGA

        Args:
            cmd (CMD): Command code to send to the FPGA
            length (str): Length of the body of the command (if any)
            body (str): Body information of the command
            timeout (int): Time in seconds to wait for socket data until timeout

        Returns:
            str: Response message

        """
        # Create timeout exception
        self.config_socket.settimeout(timeout)

        # Create and send message
        resp = ''
        msg = codecs.decode(
            ''.join((CONFIG_HEADER, str(cmd), length, body, CONFIG_FOOTER)),
            'hex')
        try:
            self.config_socket.sendto(msg, self.cfg_dest)
            resp, addr = self.config_socket.recvfrom(MAX_PACKET_SIZE)
        except socket.timeout as e:
            print(e)
        return resp

    def _read_data_packet(self):
        """Helper function to read in a single ADC packet via UDP

        Returns:
 
           int: Current packet number, byte count of data that has already been read, raw ADC data in current packet
        """
        data, addr = self.data_socket.recvfrom(MAX_PACKET_SIZE)
        packet_num = struct.unpack('<1l', data[:4])[0]
        byte_count = struct.unpack('>Q', b'\x00\x00' + data[4:10][::-1])[0]
        packet_data = np.frombuffer(data[10:], dtype=np.int16)
        return packet_num, byte_count, packet_data

    def _listen_for_error(self):
        """Helper function to try and read in for an error message from the FPGA

        Returns:
            None

        """
        self.config_socket.settimeout(None)
        msg = self.config_socket.recvfrom(MAX_PACKET_SIZE)
        if msg == b'5aa50a000300aaee':
            print('stopped:', msg)

    def _stop_stream(self):
        """Helper function to send the stop command to the FPGA

        Returns:
            str: Response Message

        """
        return self._send_command(CMD.RECORD_STOP_CMD_CODE)

    @staticmethod
    def organize(raw_frame, num_chirps, num_rx, num_samples):
        """Reorganizes raw ADC data into a full frame

        Args:
            raw_frame (ndarray): Data to format
            num_chirps: Number of chirps included in the frame
            num_rx: Number of receivers used in the frame
            num_samples: Number of ADC samples included in each chirp

        Returns:
            ndarray: Reformatted frame of raw data of shape (num_chirps, num_rx, num_samples)

        """

        # ret = np.zeros(len(raw_frame) // 2, dtype=complex)

        # # Separate IQ data
        # ret[0::2] = raw_frame[0::4] + 1j * raw_frame[2::4]
        # ret[1::2] = raw_frame[1::4] + 1j * raw_frame[3::4]
        # return ret.reshape((num_chirps, num_rx, num_samples))

        
        sampleI = np.zeros(len(raw_frame) // 2, dtype=complex)
        sampleQ = np.zeros(len(raw_frame) // 2, dtype=complex)
        
        sampleI[0::2] = raw_frame[0::4]
        sampleI[1::2] = raw_frame[1::4]
        sampleQ[0::2] = raw_frame[2::4]
        sampleQ[1::2] = raw_frame[3::4]

        sample = sampleQ  + 1j * sampleI 

        return sample.reshape((num_chirps, num_rx, num_samples))

import mmwave.dsp as dsp
import mmwave.clustering as clu
from demo.visualizer import ellipse_visualize

if __name__ == "__main__":
    DCA = DCA1000(configFileName)
    DCA.configure()
    time.sleep(3)
    DCA._serialConfig()
    DCA.start()
    fig = plt.figure()

    numADCSamples = 96
    numTxAntennas = 3
    numRxAntennas = 4
    numLoopsPerFrame = 96
    numChirpsPerFrame = numTxAntennas * numLoopsPerFrame

    numRangeBins = numADCSamples
    numDopplerBins = numLoopsPerFrame
    numAngleBins = 64

    range_resolution, bandwidth = dsp.range_resolution(numADCSamples)
    doppler_resolution = dsp.doppler_resolution(bandwidth)

    plotRangeDopp = False  
    plot2DscatterXY = True  
    plot2DscatterXZ = False  
    plot3Dscatter = False  
    plotCustomPlt = False

    visTrigger = plot2DscatterXY + plot2DscatterXZ + plot3Dscatter + plotRangeDopp + plotCustomPlt
    assert visTrigger < 2, "Can only choose to plot one type of plot at once"

    singFrameView = False

    # (1.5) Required Plot Declarations
    # if plot2DscatterXY or plot2DscatterXZ:
    #     fig, axes = plt.subplots(1, 2)
    # elif plot3Dscatter:
    #     fig = plt.figure()
    # elif plotRangeDopp:
    #     fig = plt.figure()
    # elif plotCustomPlt:
    #     print("Using Custom Plotting")

    rx_index = [3,1,2,0,11,9,7,5,10,8,6,4];
    sensor_number=12;
    sensor_array = np.array([[1,1,0,0],[1,1,0,0],[1,1,1,1],[1,1,1,1]])

    index = 0;
    frame = np.zeros((100,96,12,96),dtype=complex)
    while True:
        try:
            # (1) Reading in adc data
            # start0 = time.time()
            adc_data = DCA.read(timeout=.5)
            # DCA.organize基本不占用时间
            frame[index] = DCA.organize(adc_data, ADC_PARAMS['chirps'],ADC_PARAMS['rx'] * ADC_PARAMS['tx'],ADC_PARAMS['samples'])
            if index == 99:
                import scipy.io as io
                mat_path = 'rx.mat'
                io.savemat(mat_path, {'rx': frame})
            else:
                index = index + 1
            # end0 = time.time()
            # print('capture data is :%s seconds'%(end0 - start0))

        except KeyboardInterrupt:
            print("结束！")
            DCA.stop()
            plt.close()
            break

    plt.show()
    exit(0)

#
# File:         实现逻辑法航迹起始，NN法数据关联，卡尔曼滤波以及传统航迹管理（2维）
#               主要实现了MOT迭代函数
# Author:       UESTC Yu Xuyao
# Email:        yxy19991102@163.com
#

import scipy.io as scio
import numpy
import scipy.stats as st
import time
from IWR.track import *

# 2自由度卡方分布
threshold = st.chi2.ppf(0.80, 2)
T = 0.1

F = numpy.mat([[1, 0, T, 0], [0, 1, 0, T], [0, 0, 1, 0], [0, 0, 0, 1]])
H = numpy.mat([[1, 0, 0, 0], [0, 1, 0, 0]])
Gamma = numpy.mat([[T * T / 2, 0], [0, T * T / 2], [T, 0], [0, T]])
Q = Gamma * Gamma.T * 0.5
R = 0.1 * H * H.T

vmin = 0.2
vmax = 4

max_velocity = 1.5


def beginTrack(points):
    global threshold, T, F, H, Gamma, Q, R, vmin, vmax

    outside = []

    index = 0
    testroot = []

    Z_k = (points[0])
    Z_k_plus1 = (points[1])

    for j in range(0, Z_k.shape[1]):
        for k in range(0, Z_k_plus1.shape[1]):
            d = max(0, numpy.linalg.norm(Z_k_plus1[:, k] - Z_k[:, j]) - vmax * T) + max(0, -numpy.linalg.norm(
                Z_k_plus1[:, k] - Z_k[:, j]) + vmin * T)
            D = d / numpy.linalg.det(R + R) * d
            if D <= threshold:
                temp = TestTrack(Z_k[:, j])
                temp.addseq(Z_k_plus1[:, k])

                Px0 = 5 * numpy.eye(4)
                P = F * Px0 * F.T + Q
                S = H * P * H.T + R * 0.2
                K = P * H.T * numpy.linalg.inv(S)
                Pkk = (numpy.eye(4) - K * H) * P

                x_init = numpy.concatenate((Z_k_plus1[:, k], (Z_k_plus1[:, k] - Z_k[:, j]) / T), axis=0)
                x_forest = F * x_init

                temp.est = x_init
                temp.pkk = Pkk
                temp.x_predict = x_forest
                temp.M = 2
                temp.N = 2

                testroot.append(temp)

    return testroot


def TestToConfirmed(testroot):
    confirmedroot = ConfirmedTrack(testroot.start)
    confirmedroot.seq = testroot.seq
    confirmedroot.est = testroot.est
    confirmedroot.pkk = testroot.pkk
    confirmedroot.seq = testroot.seq
    confirmedroot.x_predict = testroot.x_predict
    return confirmedroot


def MOT(Z_k, Z_k_prev, confirmedroot, testroot):
    global threshold, T, F, H, Gamma, Q, R, vmin, vmax

    Z_k_present = Z_k

    if Z_k.shape[1] == 0:
        if len(confirmedroot) != 0:
            pos = []
            for kk in range(0, len(confirmedroot)):
                confirmedroot[kk].board += 1
                if confirmedroot[kk].board >= 4:
                    pos.append(kk)
            confirmedroot = numpy.delete(confirmedroot, pos)
            confirmedroot = confirmedroot.tolist()

    if Z_k.shape[1] != 0:
        if len(confirmedroot) != 0:
            pos = []
            for kk in range(0, len(confirmedroot)):
                if Z_k.shape[1] == 0:
                    break

                P = F * confirmedroot[kk].pkk * F.T + Q
                S = H * P * H.T + R
                K = P * H.T * numpy.linalg.inv(S)
                outside = H * confirmedroot[kk].x_predict
                V = []
                for i in range(0, Z_k.shape[1]):
                    V.append((Z_k[:, i] - outside).T * numpy.linalg.inv(S) * (Z_k[:, i] - outside))

                index = numpy.argmin(V)
                val = numpy.amin(V)
                if val <= threshold:
                    confirmedroot[kk].addseq(Z_k[:, index])
                    est = confirmedroot[kk].x_predict + K * (Z_k[:, index] - outside)
                    confirmedroot[kk].addest(est)
                    confirmedroot[kk].pkk = (numpy.eye(4) - K * H) * P
                    confirmedroot[kk].x_predict = F * est
                    confirmedroot[kk].board = 0

                    Z_k = numpy.delete(Z_k, index, axis=1)
                else:
                    confirmedroot[kk].addseq(outside)
                    confirmedroot[kk].addest(confirmedroot[kk].x_predict)
                    confirmedroot[kk].pkk = P
                    confirmedroot[kk].x_predict = F * confirmedroot[kk].est[:, -1]
                    confirmedroot[kk].board += 1

                if confirmedroot[kk].board >= 4 or numpy.linalg.norm(confirmedroot[kk].est[2:4, -1]) > max_velocity:
                    pos.append(kk)

            confirmedroot = numpy.delete(confirmedroot, pos)
            confirmedroot = confirmedroot.tolist()

    if Z_k.shape[1] != 0:
        if len(testroot) != 0:
            pos = []
            for kk in range(0, len(testroot)):
                if Z_k.shape[1] == 0:
                    break

                P = F * testroot[kk].pkk * F.T + Q
                S = H * P * H.T + R
                K = P * H.T * numpy.linalg.inv(S)
                outside = H * testroot[kk].x_predict
                V = []
                for i in range(0, Z_k.shape[1]):
                    V.append((Z_k[:, i] - outside).T * numpy.linalg.inv(S) * (Z_k[:, i] - outside))

                index = numpy.argmin(V)
                val = numpy.amin(V)
                if val <= threshold:
                    testroot[kk].addseq(Z_k[:, index])
                    est = testroot[kk].x_predict + K * (Z_k[:, index] - outside)
                    testroot[kk].addest(est)
                    testroot[kk].pkk = (numpy.eye(4) - K * H) * P
                    testroot[kk].x_predict = F * est
                    testroot[kk].M += 1
                    testroot[kk].N += 1
                    Z_k = numpy.delete(Z_k, index, axis=1)
                else:
                    testroot[kk].addseq(outside)
                    testroot[kk].addest(testroot[kk].x_predict)
                    testroot[kk].pkk = P
                    testroot[kk].x_predict = F * testroot[kk].est[:, -1]
                    testroot[kk].N += 1

                if testroot[kk].N == 6:
                    if testroot[kk].M >= 4 and numpy.linalg.norm(testroot[kk].est[2:4, -1]) < max_velocity:
                        confirmedroot.append(TestToConfirmed(testroot[kk]))
                    else:
                        pos.append(kk)

            testroot = numpy.delete(testroot, pos)
            testroot = testroot.tolist()

    if Z_k.shape[1] != 0:
        testroot.extend(beginTrack([Z_k_prev, Z_k_present]))

    Z_k_prev = Z_k_present
    return confirmedroot, testroot, Z_k_present

#
# File:         实现逻辑法航迹起始，NN法数据关联，卡尔曼滤波以及传统航迹管理（3维）
#               主要实现了MOT迭代函数
# Author:       UESTC YuXuyao
# Email:        yxy19991102@163.com
#

import scipy.io as scio
import numpy
import scipy.stats as st
import time
from IWR.track import *

threshold = st.chi2.ppf(0.80, 3)
T = 0.1

F = numpy.mat([[1, 0, 0, T, 0, 0], [0, 1, 0, 0, T, 0], [0, 0, 1, 0, 0, T], [0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 1, 0],
               [0, 0, 0, 0, 0, 1]])
H = numpy.mat([[1, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0]])
Gamma = numpy.mat([[T * T / 2, 0, 0], [0, T * T / 2, 0], [0, 0, T * T / 2], [T, 0, 0], [0, T, 0], [0, 0, T]])
Q = Gamma * Gamma.T * 0.5
R = 0.1 * H * H.T

vmin = 0.2
vmax = 4

max_velocity = 2

def beginTrack(points):
    global threshold, T, F, H, Gamma, Q, R, vmin, vmax

    outside = []

    index = 0
    testroot = []

    Z_k = (points[0])
    Z_k_plus1 = (points[1])

    for j in range(0, Z_k.shape[1]):
        for k in range(0, Z_k_plus1.shape[1]):
            d = max(0, numpy.linalg.norm(Z_k_plus1[:, k] - Z_k[:, j]) - vmax * T) + max(0, -numpy.linalg.norm(
                Z_k_plus1[:, k] - Z_k[:, j]) + vmin * T)
            D = d / numpy.linalg.det(R + R) * d
            if D <= threshold:
                temp = TestTrack(Z_k[:, j])
                temp.addseq(Z_k_plus1[:, k])

                Px0 = 5 * numpy.eye(6)
                P = F * Px0 * F.T + Q
                S = H * P * H.T + R * 0.2
                K = P * H.T * numpy.linalg.inv(S)
                Pkk = (numpy.eye(6) - K * H) * P

                x_init = numpy.concatenate((Z_k_plus1[:, k], (Z_k_plus1[:, k] - Z_k[:, j]) / T), axis=0)
                x_forest = F * x_init

                temp.est = x_init
                temp.pkk = Pkk
                temp.x_predict = x_forest
                temp.M = 2
                temp.N = 2

                testroot.append(temp)

    return testroot


def TestToConfirmed(testroot):
    confirmedroot = ConfirmedTrack(testroot.start)
    confirmedroot.seq = testroot.seq
    confirmedroot.est = testroot.est
    confirmedroot.pkk = testroot.pkk
    confirmedroot.seq = testroot.seq
    confirmedroot.x_predict = testroot.x_predict
    return confirmedroot


def MOT(Z_k, Z_k_prev, confirmedroot, testroot):
    global threshold, T, F, H, Gamma, Q, R, vmin, vmax

    Z_k_present = Z_k

    if Z_k.shape[1] == 0:
        if len(confirmedroot) != 0:
            pos = []
            for kk in range(0, len(confirmedroot)):
                confirmedroot[kk].board += 1
                if confirmedroot[kk].board >= 4:
                    pos.append(kk)
            confirmedroot = numpy.delete(confirmedroot, pos)
            confirmedroot = confirmedroot.tolist()

    if Z_k.shape[1] != 0:
        if len(confirmedroot) != 0:
            pos = []
            for kk in range(0, len(confirmedroot)):
                if Z_k.shape[1] == 0:
                    break

                P = F * confirmedroot[kk].pkk * F.T + Q
                S = H * P * H.T + R
                K = P * H.T * numpy.linalg.inv(S)
                outside = H * confirmedroot[kk].x_predict
                V = []
                for i in range(0, Z_k.shape[1]):
                    V.append((Z_k[:, i] - outside).T * numpy.linalg.inv(S) * (Z_k[:, i] - outside))

                index = numpy.argmin(V)
                val = numpy.amin(V)
                if val <= threshold:
                    confirmedroot[kk].addseq(Z_k[:, index])
                    est = confirmedroot[kk].x_predict + K * (Z_k[:, index] - outside)
                    confirmedroot[kk].addest(est)
                    confirmedroot[kk].pkk = (numpy.eye(6) - K * H) * P
                    confirmedroot[kk].x_predict = F * est
                    confirmedroot[kk].board = 0

                    Z_k = numpy.delete(Z_k, index, axis=1)
                else:
                    confirmedroot[kk].addseq(outside)
                    confirmedroot[kk].addest(confirmedroot[kk].x_predict)
                    confirmedroot[kk].pkk = P
                    confirmedroot[kk].x_predict = F * confirmedroot[kk].est[:, -1]
                    confirmedroot[kk].board += 1

                if confirmedroot[kk].board >= 4 or numpy.linalg.norm(confirmedroot[kk].est[3:6, -1]) > max_velocity:
                    pos.append(kk)

            confirmedroot = numpy.delete(confirmedroot, pos)
            confirmedroot = confirmedroot.tolist()

    if Z_k.shape[1] != 0:
        if len(testroot) != 0:
            pos = []
            for kk in range(0, len(testroot)):
                if Z_k.shape[1] == 0:
                    break

                P = F * testroot[kk].pkk * F.T + Q
                S = H * P * H.T + R
                K = P * H.T * numpy.linalg.inv(S)
                outside = H * testroot[kk].x_predict
                V = []
                for i in range(0, Z_k.shape[1]):
                    V.append((Z_k[:, i] - outside).T * numpy.linalg.inv(S) * (Z_k[:, i] - outside))

                index = numpy.argmin(V)
                val = numpy.amin(V)
                if val <= threshold:
                    testroot[kk].addseq(Z_k[:, index])
                    est = testroot[kk].x_predict + K * (Z_k[:, index] - outside)
                    testroot[kk].addest(est)
                    testroot[kk].pkk = (numpy.eye(6) - K * H) * P
                    testroot[kk].x_predict = F * est
                    testroot[kk].M += 1
                    testroot[kk].N += 1
                    Z_k = numpy.delete(Z_k, index, axis=1)
                else:
                    testroot[kk].addseq(outside)
                    testroot[kk].addest(testroot[kk].x_predict)
                    testroot[kk].pkk = P
                    testroot[kk].x_predict = F * testroot[kk].est[:, -1]
                    testroot[kk].N += 1

                if testroot[kk].N == 6:
                    if testroot[kk].M >= 4 and numpy.linalg.norm(testroot[kk].est[3:6, -1]) < max_velocity:
                        confirmedroot.append(TestToConfirmed(testroot[kk]))
                    else:
                        pos.append(kk)

            testroot = numpy.delete(testroot, pos)
            testroot = testroot.tolist()

    if Z_k.shape[1] != 0:
        testroot.extend(beginTrack([Z_k_prev, Z_k_present]))

    Z_k_prev = Z_k_present
    return confirmedroot, testroot, Z_k_present

#
# File:         读取IWR6843处理得到的点云数据（3维）进行目标跟踪，同时调用openGL对第一个航迹进行显示（多个显示未完成）
# Author:       UESTC Yu Xuyao
# Email:        yxy19991102@163.com
#

import serial
import time
import numpy as np
import pyqtgraph as pg
from pyqtgraph.Qt import QtGui
import pyqtgraph.opengl as gl
from IWR.MOT3D import *
import numpy
from sklearn.cluster import dbscan
import pandas

# Change the configuration file name
configFileName = 'profile_iwr6843_ods_3d.cfg'

CLIport = {}
Dataport = {}
byteBuffer = np.zeros(2 ** 15, dtype='uint8')
byteBufferLength = 0

testroot = []
confirmedroot = []
Z_k_prev = numpy.mat([])


# ------------------------------------------------------------------
# Function to configure the serial ports and send the data from
# the configuration file to the radar
def serialConfig(configFileName):
    global CLIport
    global Dataport
    # Open the serial ports for the configuration and the data ports

    # Linux
    # CLIport = serial.Serial('/dev/ttyACM0', 115200)
    # Dataport = serial.Serial('/dev/ttyACM1', 921600)

    # Windows
    CLIport = serial.Serial('COM7', 115200)
    Dataport = serial.Serial('COM8', 921600)

    # Read the configuration file and send it to the board
    config = [line.rstrip('\r\n') for line in open(configFileName)]
    for i in config:
        print(i)
        CLIport.write((i + '\n').encode())
        # wait for reply
        time.sleep(0.1)
        reply = CLIport.read(CLIport.in_waiting).decode()
        print(reply)

    return CLIport, Dataport

# ------------------------------------------------------------------

# Funtion to read and parse the incoming data
def readAndParseData(Dataport):
    global byteBuffer, byteBufferLength

    # Constants
    OBJ_STRUCT_SIZE_BYTES = 12
    BYTE_VEC_ACC_MAX_SIZE = 2 ** 15
    MMWDEMO_UART_MSG_DETECTED_POINTS = 1
    MMWDEMO_UART_MSG_RANGE_PROFILE = 2
    maxBufferSize = 2 ** 15
    tlvHeaderLengthInBytes = 8
    pointLengthInBytes = 16
    magicWord = [2, 1, 4, 3, 6, 5, 8, 7]

    # Initialize variables
    magicOK = 0  # Checks if magic number has been read
    dataOK = 0  # Checks if the data has been read correctly
    frameNumber = 0
    detObj = {}

    readBuffer = Dataport.read(Dataport.in_waiting)
    byteVec = np.frombuffer(readBuffer, dtype='uint8')
    byteCount = len(byteVec)

    # Check that the buffer is not full, and then add the data to the buffer
    if (byteBufferLength + byteCount) < maxBufferSize:
        byteBuffer[byteBufferLength:(byteBufferLength + byteCount)] = byteVec[0:byteCount]
        byteBufferLength = byteBufferLength + byteCount

    # Check that the buffer has some data
    if byteBufferLength > 16:

        # Check for all possible locations of the magic word
        possibleLocs = np.where(byteBuffer == magicWord[0])[0]

        # Confirm that is the beginning of the magic word and store the index in startIdx
        startIdx = []
        for loc in possibleLocs:
            check = byteBuffer[loc:loc + 8]
            if np.all(check == magicWord):
                startIdx.append(loc)

        # Check that startIdx is not empty
        if startIdx:

            # Remove the data before the first start index
            if 0 < startIdx[0] < byteBufferLength:
                byteBuffer[:byteBufferLength - startIdx[0]] = byteBuffer[startIdx[0]:byteBufferLength]
                byteBuffer[byteBufferLength - startIdx[0]:] = np.zeros(len(byteBuffer[byteBufferLength - startIdx[0]:]),
                                                                       dtype='uint8')
                byteBufferLength = byteBufferLength - startIdx[0]

            # Check that there have no errors with the byte buffer length
            if byteBufferLength < 0:
                byteBufferLength = 0

            # Read the total packet length
            totalPacketLen = int.from_bytes(byteBuffer[12:12 + 4], byteorder='little')

            # Check that all the packet has been read
            if (byteBufferLength >= totalPacketLen) and (byteBufferLength != 0):
                magicOK = 1

    # If magicOK is equal to 1 then process the message
    if magicOK:

        # Initialize the pointer index
        idX = 0

        # Read the header
        magicNumber = byteBuffer[idX:idX + 8]
        idX += 8
        version = format(int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little'), 'x')
        idX += 4
        totalPacketLen = int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little')
        idX += 4
        platform = format(int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little'), 'x')
        idX += 4
        frameNumber = int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little')
        idX += 4
        timeCpuCycles = int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little')
        idX += 4
        numDetectedObj = int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little')
        idX += 4
        numTLVs = int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little')
        idX += 4
        subFrameNumber = int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little')
        idX += 4

        # Read the TLV messages
        for tlvIdx in range(numTLVs):

            # Check the header of the TLV message
            tlv_type = int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little')
            idX += 4
            tlv_length = int.from_bytes(byteBuffer[idX:idX + 4], byteorder='little')
            idX += 4

            # Read the data depending on the TLV message
            if tlv_type == MMWDEMO_UART_MSG_DETECTED_POINTS:

                # Initialize the arrays
                x = np.zeros(numDetectedObj, dtype=np.float32)
                y = np.zeros(numDetectedObj, dtype=np.float32)
                z = np.zeros(numDetectedObj, dtype=np.float32)
                velocity = np.zeros(numDetectedObj, dtype=np.float32)

                for objectNum in range(numDetectedObj):
                    # Read the data for each object
                    x[objectNum] = byteBuffer[idX:idX + 4].view(dtype=np.float32)
                    idX += 4
                    y[objectNum] = byteBuffer[idX:idX + 4].view(dtype=np.float32)
                    idX += 4
                    z[objectNum] = byteBuffer[idX:idX + 4].view(dtype=np.float32)
                    idX += 4
                    velocity[objectNum] = byteBuffer[idX:idX + 4].view(dtype=np.float32)
                    idX += 4

                # Store the data in the detObj dictionary
                detObj = {"numObj": numDetectedObj, "x": x, "y": y, "z": z, "velocity": velocity}
                dataOK = 1

        # Remove already processed data
        if 0 < idX < byteBufferLength:
            shiftSize = totalPacketLen

            byteBuffer[:byteBufferLength - shiftSize] = byteBuffer[shiftSize:byteBufferLength]
            byteBuffer[byteBufferLength - shiftSize:] = np.zeros(len(byteBuffer[byteBufferLength - shiftSize:]),
                                                                 dtype='uint8')
            byteBufferLength = byteBufferLength - shiftSize

            # Check that there are no errors with the buffer length
            if byteBufferLength < 0:
                byteBufferLength = 0

    return dataOK, frameNumber, detObj


# ------------------------------------------------------------------

# Funtion to update the data and display in the plot
def update():
    dataOk = 0
    global detObj, confirmedroot, testroot, Z_k_prev, LinePlot
    x = []
    y = []
    z = []
    # Read and parse the received data
    dataOk, frameNumber, detObj = readAndParseData(Dataport)

    if dataOk:
        x = detObj["x"]
        y = detObj["y"]
        z = detObj["z"]

        if len(x) != 0:
            core_samples, cluster_ids = dbscan(np.vstack((x, y, z)).T, eps=0.3, min_samples=3)
            df = pandas.DataFrame(np.c_[np.vstack((x, y, z)).T, cluster_ids],
                                  columns=['x', 'y', 'z', 'cluster_id'])
            df = df[df.cluster_id != -1]
            count = df['cluster_id'].value_counts()

            x = []
            y = []
            z = []
            if count.shape[0] >= 1:
                for i in range(0, count.shape[0]):
                    df1 = df[df.cluster_id == i]
                    mean = np.mean([df1['x'], df1['y'], df1['z']], 1)
                    x.append(mean[0])
                    y.append(mean[1])
                    z.append(mean[2])

        Z_k = numpy.mat([x, y, z])

        confirmedroot, testroot, Z_k_prev = MOT(Z_k, Z_k_prev, confirmedroot, testroot)

        if len(confirmedroot) > 0:
            tmppos = confirmedroot[0].est[0:3, :].T.A
            for i in range(1,len(confirmedroot)):
                tmppos = numpy.vstack((tmppos,confirmedroot[i].est[0:3, :].T.A))
            LinePlot.setData(pos=tmppos)
            QtGui.QApplication.processEvents()
        else:
            LinePlot.setData(pos=np.array([100, 100, 100]))
            QtGui.QApplication.processEvents()

    return dataOk


# -------------------------    MAIN   -----------------------------------------

# Configurate the serial port
CLIport, Dataport = serialConfig(configFileName)

# START QtAPPfor the plot
app = QtGui.QApplication([])

app = pg.mkQApp()

view = gl.GLViewWidget()
view.setBackgroundColor('k')
view.show()

axis = gl.GLAxisItem()
axis.scale(10, 10, 10)

view.addItem(axis)

LinePlot = gl.GLScatterPlotItem()
view.addItem(LinePlot)

# Main loop
detObj = {}
frameData = {}
currentIndex = 0
while True:
    try:
        # Update the data and check if the data is okay
        dataOk = update()

        if dataOk:
            # Store the current frame into frameData
            frameData[currentIndex] = detObj
            currentIndex += 0

        time.sleep(0.01)  # Sampling frequency of 20 Hz

    # Stop the program and close everything if Ctrl + c is pressed
    except KeyboardInterrupt:
        CLIport.write(('sensorStop\n').encode())
        print('sensorStop\n')
        CLIport.close()
        Dataport.close()
        win.close()
        break