from ast import Continue
import socket
import sys
import platform
import binascii
import datetime
import csv

from queue import Queue
import threading
import time

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np

class CH10Packet:
    def __init__(self, packetData):
        offset = 0;
        self.msgFormat = packetData[offset] & 0x0F
        self.msgType = (packetData[offset] >> 4) & 0x0F
        offset += 1
        self.msgSeqNo = (packetData[offset+2] << 16) | (packetData[offset+1] << 8) | (packetData[offset])
        offset += 3
        self.uSync = (packetData[offset+1] << 8) | (packetData[offset])
        offset += 2
        self.channelID = (packetData[offset+1] << 8) | (packetData[offset])
        offset += 2
        self.pktLen = (packetData[offset+3] << 24) | (packetData[offset+2] << 16) | (packetData[offset+1] << 8) | (packetData[offset])
        offset += 4
        self.dataLen = (packetData[offset+3] << 24) | (packetData[offset+2] << 16) | (packetData[offset+1] << 8) | (packetData[offset])
        offset += 4
        self.dataTypeVer = (packetData[offset])
        offset += 1
        self.seqNumber = (packetData[offset])
        offset += 1
        self.pktFlags = (packetData[offset])
        offset += 1
        self.dataType = (packetData[offset])

        self.bIncludesSecondaryTimeHeader = False
        if (self.pktFlags / 128.0 >= 1.0):
            self.bIncludesSecondaryTimeHeader = True

        if self.dataType != 0X21:
            return
        
        if self.dataType == 0x00 or self.dataType == 0x01:
            return

        offset += 1
        self.aubyRelTime = (packetData[offset+5] << 40) | (packetData[offset+4] << 32) | (packetData[offset+3] << 24) | (packetData[offset+2] << 16) | (packetData[offset+1] << 8) | (packetData[offset])
        
        offset += 6
        self.checksum = (packetData[offset+1] << 8) | (packetData[offset])
        
        offset += 2
        if(self.bIncludesSecondaryTimeHeader):
            self.secondaryTimeNs = (packetData[offset+3] << 24) | (packetData[offset+2] << 16) | (packetData[offset+1] << 8) | (packetData[offset])
            offset += 4
            self.secondaryTimeSec = (packetData[offset+3] << 24) | (packetData[offset+2] << 16) | (packetData[offset+1] << 8) | (packetData[offset])
            offset += 4
            offset += 2 #rsv
            self.checksumSecondary = (packetData[offset+1] << 8) | (packetData[offset])
            offset += 2

        self.Ch10AnalogMode = packetData[offset] & 0x03
        self.Ch10AnalogLength = (packetData[offset] >> 2) & 0x3F
        offset += 1
        self.Ch10AnalogSubChannelID = (packetData[offset])
        offset += 1
        self.Ch10AnalogSubChannelCount = (packetData[offset])
        offset += 1
        self.Ch10AnalogFactor = packetData[offset] & 0x0F
        self.Ch10AnalogSAME = (packetData[offset] >> 4) & 0x0F
        offset += 1
        
        self.bIsTSRAIR = False
        if(self.Ch10AnalogSubChannelCount==18):
            self.bIsTSRAIR = True

        self.sampleData = list()
        self.sampleCount = ((self.dataLen - 4) / self.Ch10AnalogSubChannelCount)/2.0
        i = 1
        while i <= self.sampleCount:
            thisSample = list()
            ii = 1
            if(self.bIncludesSecondaryTimeHeader):
                val = float(self.secondaryTimeSec) + (float(self.secondaryTimeNs)/1000000000.0) + (i/8000.0)
                
                thisSample.append(val)
                    # datetime.datetime.utcfromtimestamp(self.secondaryTimeSec)
                    # .strftime('%Y-%m-%d %H:%M:%S') + f':{self.secondaryTimeNs}')
            else:
                thisSample.append(self.aubyRelTime)
            
            while ii <= self.Ch10AnalogSubChannelCount:
                thisSample.append(toSigned16(self.bIsTSRAIR,(packetData[offset+1] << 8) | (packetData[offset])))
                offset+=2
                ii += 1
            
            self.sampleData.append(thisSample)
            i += 1

class CH10TMATPacket:
    def __init__(self, packetData):
        offset = 0;
        self.Version = packetData[offset] & 0xFF
        if self.Version != 0x11:
            return
        offset += 1
        self.SetupRecordConfigurationChange = (packetData[offset]) & 0x01
        self.Format = (packetData[offset] & 0x02) >> 1
        offset = 40
        self.tmats = packetData[offset:len(packetData)+1].decode("utf-8")

multicast_group = '239.1.2.10'
interface_ip = '192.168.4.199'
multicast_port_range = [8400]
fig, ax = plt.subplots()
line, = ax.plot([], [])  # Empty line object for updating data
lines = [line]

data_queue = Queue()
plotwidth = 50000
xdata = []
ydata = []

def toSigned16(bIsTSRAIR, n):
    if bIsTSRAIR:
        n = n & 0xffff
        return (n ^ 0x8000) - 0x8000
    else:
        return n - 0x8000

def create_socket(interface_ip, multicast_port, multicast_group):
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 32)
    
    # linux binds to multicast address, windows to interface address
    ip_bind = interface_ip if platform.system() == "Windows" else multicast_group
    sock.bind((ip_bind, multicast_port))
    
     # socket.IPPROTO_IP works on Linux and Windows
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_IF, socket.inet_aton(interface_ip))
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, socket.inet_aton(multicast_group) + socket.inet_aton(interface_ip))
    return sock

def initPlot():
    global xdata 
    global ydata
    ax.set_xlim(0, plotwidth)
    ax.set_ylim(-32767, 32767) 
    line.set_data([], [])
    line.set_linewidth(1)
    ax.autoscale(axis='y')

    xdata = [i * 1.0 for i in range(plotwidth)]
    ydata = [0.0 for i in range(plotwidth)]

    return line,

def animatePlot(frame):
    global xdata 
    global ydata
    items = []

    while data_queue.qsize() > 0:
        items.append(data_queue.get_nowait())

    times = [np.floor(r[0]*1000) / 1000 for r in items]
    xdata.extend(times)
    xdata = xdata[0-plotwidth:]

    new = [r[1] for r in items]
    ydata.extend(new)
    ydata = ydata[0-plotwidth:]
    
    ax.set_xlim(min(xdata),max(xdata))
    # ax.set_autoscalex_on(True)
    
    line.set_data(xdata, ydata)

    return line,

def read_packets():
    read_list = []
    for multicast_port in multicast_port_range:
        read_list.append(create_socket(interface_ip, multicast_port, multicast_group))
    bDisplayPacketMetadata = False
    bDisplayPacketData = False
    bDisplayPacketTmats = True
    while True:
        for sock in read_list:
            #print(f'Reading --- {sock.getsockname()[1]}')
            data, srcAddress = sock.recvfrom(1500)

            ch10packet = CH10Packet(data)
            # tmatsPacket = CH10TMATPacket(data)
            #print(f'Received packet of {len(data)} bytes from {srcAddress[0]}:{sock.getsockname()[1]}')
            if ch10packet.dataType == 0X21:
                for i in ch10packet.sampleData:
                    data_queue.put(i)

                if bDisplayPacketData:
                    my_array = np.array(ch10packet.sampleData)

                    hexdata = binascii.hexlify(my_array, " ", 2)
                    #hexdata = ''.join(['{:02x} '.format(b) for b in ch10packet.sampleData])
                    print('Data = %s' % hexdata)
                
                if bDisplayPacketMetadata:            
                    print(f'msgFormat: {ch10packet.msgFormat}')
                    print(f'msgType: {ch10packet.msgType}')
                    print(f'msgSeqNo: {ch10packet.msgSeqNo}')
                    print(f'uSync: {ch10packet.uSync}')
                    print(f'channelID: {ch10packet.channelID}')
                    print(f'pktLen: {ch10packet.pktLen}')
                    print(f'dataLen: {ch10packet.dataLen}')
                    print(f'dataTypeVer: {ch10packet.dataTypeVer}')
                    print(f'seqNumber: {ch10packet.seqNumber}')
                    print(f'pktFlags: {ch10packet.pktFlags}')
                    print(f'dataType: {ch10packet.dataType}')
                    print(f'aubyRelTime: {ch10packet.aubyRelTime}')
                    print(f'checksum: {ch10packet.checksum}')
                    if ch10packet.bIncludesSecondaryTimeHeader:
                        print(f'secondaryTimeNs: {ch10packet.secondaryTimeNs}')
                        print(f'secondaryTimeSec: {ch10packet.secondaryTimeSec}')
                        print('UTC Time:',datetime.datetime.utcfromtimestamp(ch10packet.secondaryTimeSec).strftime('%Y-%m-%d %H:%M:%S'), f':{ch10packet.secondaryTimeNs}')
                    print(f'Ch10AnalogMode: {ch10packet.Ch10AnalogMode}')
                    print(f'Ch10AnalogLength: {ch10packet.Ch10AnalogLength}')
                    print(f'Ch10AnalogSubChannelID: {ch10packet.Ch10AnalogSubChannelID}')
                    print(f'Ch10AnalogSubChannelCount: {ch10packet.Ch10AnalogSubChannelCount}')
                    print(f'Ch10AnalogFactor: {ch10packet.Ch10AnalogFactor}')
                    print(f'Ch10AnalogSAME: {ch10packet.Ch10AnalogSAME}')
                    print(f'Channel 1: {toSigned16(ch10packet.bIsTSRAIR,ch10packet.sampleData[0][1])}')
        
                # with open(f'{srcAddress[0]}.{sock.getsockname()[1]}.{epochTimeNow}.csv', 'a', newline='') as csvfile:
                #     wr = csv.writer(csvfile, quoting=csv.QUOTE_ALL)
                #     for sample in ch10packet.sampleData:
                #         wr.writerow(sample)
def main():
    epochTimeNow = (datetime.datetime.now(datetime.timezone.utc) - datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc)).total_seconds()
    consumer_thread = threading.Thread(target=read_packets)
    consumer_thread.daemon = True
    consumer_thread.start()
    
    time.sleep(0.5)

    ani = animation.FuncAnimation(
        fig, animatePlot, 
        init_func=initPlot,
        interval=20, 
        # blit=True,
        cache_frame_data=False
        )


    plt.show()
    # input('Press enter to exit')

if __name__ == '__main__':
    main()
