#!/usr/bin/env python3
# -*- coding: utf-8 -*-

###############################################################################
###############################################################################
########    V1.4 2021/08/25     francescopiscitelli      ######################
########    script to read the pcapng file from VMM readout
###############################################################################
###############################################################################

# import argparse
import numpy as np
import pcapng as pg
import os
import time
import sys
from lib import libPlotting as plo

###############################################################################
###############################################################################

            
class readouts():            
    def __init__(self): 
        self.Ring    = np.zeros((0), dtype = 'float64')
        self.Fen     = np.zeros((0), dtype = 'float64')
        self.VMM     = np.zeros((0), dtype = 'float64')
        self.hybrid  = np.zeros((0), dtype = 'float64')
        self.ASIC    = np.zeros((0), dtype = 'float64')
        self.Channel = np.zeros((0), dtype = 'float64')
        self.ADC     = np.zeros((0), dtype = 'float64')
        self.timeStamp  = np.zeros((0), dtype = 'float64')
        self.BC      = np.zeros((0), dtype = 'float64')
        self.OTh     = np.zeros((0), dtype = 'float64')
        self.TDC     = np.zeros((0), dtype = 'float64')
        self.GEO     = np.zeros((0), dtype = 'float64')
               
    def transformInReadouts(self, data):
        self.Ring     = data[:,0]
        self.Fen      = data[:,1]
        self.VMM      = data[:,2]
        self.hybrid   = data[:,3]
        self.ASIC     = data[:,4]
        self.Channel  = data[:,5]
        self.ADC      = data[:,6]
        self.timeStamp  = data[:,7]
        self.BC   = data[:,8]
        self.OTh  = data[:,9]
        self.TDC  = data[:,10]
        self.GEO  = data[:,11]
                
    # def list(self):
    #     print("Rings {}".format(self.Ring))
    #     print("Fens {}".format(self.Fen))
    
    def append(self, reado):
        
        self.Ring     = np.concatenate((self.Ring, reado.Ring), axis=0)
        self.Fen      = np.concatenate((self.Fen, reado.Fen), axis=0)
        self.VMM      = np.concatenate((self.VMM, reado.VMM), axis=0)
        self.hybrid   = np.concatenate((self.hybrid, reado.hybrid), axis=0)
        self.ASIC     = np.concatenate((self.ASIC, reado.ASIC), axis=0)
        self.Channel  = np.concatenate((self.Channel, reado.Channel), axis=0)
        self.ADC      = np.concatenate((self.ADC, reado.ADC), axis=0)
        self.timeStamp  = np.concatenate((self.timeStamp, reado.timeStamp), axis=0)
        self.BC   = np.concatenate((self.BC, reado.BC), axis=0)
        self.OTh  = np.concatenate((self.OTh, reado.OTh), axis=0)
        self.TDC  = np.concatenate((self.TDC, reado.TDC), axis=0)
        self.GEO  = np.concatenate((self.GEO, reado.GEO), axis=0)
              
    def concatenateReadoutsInArrayForDebug(self):
        
        leng = len(self.timeStamp)
        
        readoutsArray = np.zeros((leng,7),dtype = 'float64')
        
        readoutsArray[:,0] = self.timeStamp
        readoutsArray[:,1] = self.Ring
        readoutsArray[:,2] = self.Fen
        readoutsArray[:,3] = self.hybrid
        readoutsArray[:,4] = self.ASIC
        readoutsArray[:,5] = self.Channel
        readoutsArray[:,6] = self.ADC
        
        return readoutsArray

###############################################################################
###############################################################################

class checkInstrumentID():
    def __init__(self, ID):
        self.FREIAID = 72
        self.EstiaID = 76
        
        self.printa = True
        
        if ID == self.FREIAID:
             print('found Freia data stream')
        elif ID == self.EstiaID:
             print('found Estia data stream')
        else:
             print('found some other data stream')
             
        print('loading ... [0%]',end=' ')
              
        self.printa = False
        
        
class VMM3A():
    def __init__(self, buffer, timeResolution, timeResolutionType):
        
        self.timeResolutionType = timeResolutionType
        
        # decode into little endian integers
        PhysicalRing = int.from_bytes(buffer[0:1], byteorder='little')
        self.Fen  = int.from_bytes(buffer[1:2], byteorder='little')
        Length    = int.from_bytes(buffer[2:4], byteorder='little')
        TimeHI    = int.from_bytes(buffer[4:8], byteorder='little')
        TimeLO    = int.from_bytes(buffer[8:12], byteorder='little')
        self.BC   = int.from_bytes(buffer[12:14], byteorder='little')
        OTADC     = int.from_bytes(buffer[14:16], byteorder='little')
        G0GEO     = int.from_bytes(buffer[16:17], byteorder='little')
        self.TDC  = int.from_bytes(buffer[17:18], byteorder='little')
        self.VMM  = int.from_bytes(buffer[18:19], byteorder='little')
        self.Channel = int.from_bytes(buffer[19:20], byteorder='little')
        
        #######################
        #  IMPORTANT NOTE: phys ring is 0 and 1 for logical ring 0 etc. Always 12 logical rings 
        self.Ring = np.floor(PhysicalRing/2)
        #######################

        self.ADC = OTADC & 0x3FF  #extract only 10 LSB
        self.OTh      = OTADC >> 15    #extract only 1 MSB

        self.G0       = G0GEO >> 7
        self.GEO      = G0GEO & 0x3F
        
        self.ASIC     = self.VMM & 0x1           #extract only LSB
        self.hybrid   = (self.VMM & 0xE) >> 1    #extract only 1110 and shift right by one 

        # if bD0 == 0: # normal mode
        #    pass
        # elif bD0 == 1: # calibration mode
        #    pass

        #  in seconds
        if self.timeResolutionType == 'coarse':
            self.timeStamp =  TimeHI + TimeLO * timeResolution # coarse time resolution
        elif self.timeResolutionType == 'fine':
            self.timeStamp =  TimeHI + TimeLO*timeResolution + ( timeResolution*2 - (self.TDC*(60/255))*1e-9 )  #fine time resolution
            
###############################################################################
###############################################################################
        

class pcapng_reader():
    def __init__(self, filePathAndFileName):
        
        if os.path.exists(filePathAndFileName) is False:
            temp2 = os.path.split(filePathAndFileName)
            filePath = temp2[0]+'/'
            fileName = [temp2[1]]
            print('\n \033[1;31m---> File: '+fileName+' DOES NOT EXIST \033[1;37m')
            print('\n ---> in folder: '+filePath+' \n')
            print(' ---> Exiting ... \n')
            print('------------------------------------------------------------- \n')
            sys.exit()

        self.ff = open(filePathAndFileName, 'rb')
        
        self.readouts = readouts()
        
        self.fileSize   = os.path.getsize(filePathAndFileName) #bytes
        print('data is {} kbytes'.format(self.fileSize/1e3))
        
        self.debug = False

        self.offset = 25            #bytes Num of bytes after the word (cookie) ESS = 0x 45 53 53
        self.ESSheaderSize    = 30  #bytes
        self.dataPacketLength = 20  #bytes
        
        # self.timeResolution = 11.25e-9  #s per tick
        
        self.timeResolution = 11.35686096362965e-9  #s per tick ESS 
        
        self.numOfPacketsPerTransfer = 400 
        
        # self.numOfPacketsPerTransfer = 447

        self.expectedESSpacketSize = self.numOfPacketsPerTransfer*self.dataPacketLength+self.ESSheaderSize #8970 bytes
        self.preallocLength    =  round(self.fileSize*1.2/self.expectedESSpacketSize)*self.numOfPacketsPerTransfer
        
        # I add a 20% *1.2 for safety

        self.packetCount       = 0
        self.truePacketCount   = 0
        self.nonESSPacketCount = 0
        self.totalReadoutCount = 0      

    def __del__(self):
        try:
            self.ff.close()
        except:
            pass

    def dprint(self, msg):
        if self.debug:
            print("{}".format(msg))

    def read(self, timeResolutionType='coarse'):    
    
        self.timeResolutionType = timeResolutionType
        
        scanner = pg.FileScanner(self.ff)
        
        data = np.zeros((self.preallocLength,12), dtype='float64') 

        for block in scanner:
            self.packetCount += 1
            readoutCount      = 0

            try:
                packetLength = block.packet_len
                packetData   = block.packet_data
            except:
                continue

            self.truePacketCount += 1
            self.dprint("packet {} - length {}".format(self.packetCount, packetLength))

            indexESS = packetData.find(b'ESS')

            if indexESS == -1:
                self.nonESSPacketCount += 1
                continue

            if self.truePacketCount == 1:
                checkInstrumentID(packetData[indexESS+3])
            
            indexDataStart = indexESS + 2 + self.offset + 1
            
            ESSlength = int.from_bytes(packetData[indexESS+4:indexESS+6], byteorder='little') # bytes

            # check that ESS is always in the same place
            # tempIndexDataStart.append(indexDataStart)

            readoutCount = (packetLength - indexDataStart) / self.dataPacketLength
            self.dprint("readoutCount {}".format(readoutCount))

            if readoutCount.is_integer() is not True:
                print('something wrong with data bytes dimensions')
                break
            else:
                readoutCount = int(readoutCount)
                self.totalReadoutCount += readoutCount
                
                for currentReadout in range(readoutCount):
                
                    indexStart = indexDataStart + self.dataPacketLength * currentReadout
                    indexStop  = indexDataStart + self.dataPacketLength * (currentReadout + 1)

                    vmm3 = VMM3A(packetData[indexStart:indexStop], self.timeResolution, self.timeResolutionType)
                    
                    # self.data.append(vmm3)
                    # NOTE this append at every cycle is not efficient for speed so better to allocate the array and fill it, then append outside inner loop
                   
                    index = (self.truePacketCount-1)*self.numOfPacketsPerTransfer+currentReadout
                    
                    data[index, 0] = vmm3.Ring
                    data[index, 1] = vmm3.Fen
                    data[index, 2] = vmm3.VMM
                    data[index, 3] = vmm3.hybrid
                    data[index, 4] = vmm3.ASIC
                    data[index, 5] = vmm3.Channel
                    data[index, 6] = vmm3.ADC
                    data[index, 7] = vmm3.timeStamp
                    data[index, 8] = vmm3.BC
                    data[index, 9] = vmm3.OTh
                    data[index, 10] = vmm3.TDC
                    data[index, 11] = vmm3.GEO

                    self.dprint(" \t Packet: {} ({} bytes), Readout: {}, Ring {}, FEN {}, VMM {}, hybrid {}, ASIC {}, Ch {}, Time {} s, BC {}, OverTh {}, ADC {}, TDC {}, GEO {} " \
                                .format(self.truePacketCount,ESSlength,currentReadout+1,vmm3.Ring,vmm3.Fen,vmm3.VMM,vmm3.hybrid,vmm3.ASIC,vmm3.Channel,vmm3.timeStamp,vmm3.BC,vmm3.OTh,vmm3.ADC,vmm3.TDC,vmm3.GEO))

                
                    ###########
           
            # check 
            packetLength = readoutCount*self.dataPacketLength + self.ESSheaderSize  # bytes
            if packetLength != ESSlength and self.truePacketCount == 1:
               print('something wrong with this packet: exp size {} bytes, found {} bytes.'.format(ESSlength,packetLength))
               
            roughNumOfPackets   = round(self.fileSize/ESSlength) 
            steps = int(roughNumOfPackets/4)+1
            if np.mod(self.truePacketCount,steps) == 0 or np.mod(self.truePacketCount,steps) == 0:
                    percents = int(round(100.0 * self.truePacketCount / float(roughNumOfPackets), 1))
                    print('['+format(percents,'01d') + '%]',end=' ')
     
        print('[100%]',end=' ') 

        # here I remove  the rows that have been preallocated but no filled 
        datanew = np.delete(data,np.arange(self.totalReadoutCount,self.preallocLength),axis=0)
        
        self.readouts.transformInReadouts(datanew)
        
        # check 
        if datanew.shape[0] != self.totalReadoutCount:
            print('\nsomething wrong ... mismatch between data exp. length {} and what was read {}'.format(self.totalReadoutCount,self.data.Ring.shape[0]))
          
        print('\ndata loaded - found {} readouts ({} kbytes) - Packets: valid {}, nonESS {}, All {})'.format(self.totalReadoutCount,self.truePacketCount*ESSlength/1e3,self.truePacketCount,self.nonESSPacketCount,self.packetCount))    
          
        self.__del__()
        
###############################################################################
###############################################################################

if __name__ == '__main__':
   # parser = argparse.ArgumentParser()
   # parser.add_argument("-f", metavar='file', help = "pcap file",
   #                     type = str, default = "VMM3a_Freia.pcapng")
   # parser.add_argument('-d', action='store_true', help = "add debug print")

   tProfilingStart = time.time()

   # arg = parser.parse_args()
   # filePath = './'+"VMM3a_Freia.pcapng"
   # filePath = './'+"VMM3a.pcapng"
   
   filePath = './data/'+'freia_1k_pkts_ng.pcapng'

   # pr = pcap_reader(filePath)
   # pr.debug = True
   pr.read(resolutionType='coarse')
   vmm3 = pr.readouts
   
   aa = vmm3.concatenateReadoutsInArrayForDebug()
   
   # ppp = plo.plottingReadouts(vmm3, config)
   # ppp.plotChRaw(parameters.cassettes.cassettes)
   # ppp.plotTimeStamps(parameters.cassettes.cassettes)


   # aa = pr.d
   # bb = pr.e
   
   # aaa = aa[446900:,5:9]
   # bbb = bb[446900:,5:9]
   
   # for k in range(446900,447000,1):
   #      print(" \t Ring {}, FEN {}, VMM {}, hybrid {}, ASIC {}, Ch {}, Time {} s, BC {}, OverTh {}, ADC {}, TDC {}, GEO {} " \
   #                               .format(vmm3.Ring[k],vmm3.Fen[k],vmm3.VMM[k],vmm3.hybrid[k],vmm3.ASIC[k],vmm3.Channel[k],vmm3.timeStamp[k],vmm3.BC[k],vmm3.OTh[k],vmm3.ADC[k],vmm3.TDC[k],vmm3.GEO[k]))
   
   tElapsedProfiling = time.time() - tProfilingStart
   print('\n Data Loading Completed in %.2f s' % tElapsedProfiling) 