#!/usr/bin/python

import threading
import time
import HVPM
import struct
import time
import math
from calibrationData import calibrationData
import Operations as ops
from copy import deepcopy
import numpy as np

class channels:
    timeStamp = 0
    MainCurrent = 1
    USBCurrent = 2
    AuxCurrent = 3
    MainVoltage = 4
    USBVoltage = 5

class SampleEngine:
    def __init__(self, Monsoon,bulkProcessRate=128):
        """Declares global variables.
        During testing, we found the garbage collector would slow down sampling enough to cause a lot of dropped samples.
        We've tried to combat this by allocating as much as possible in advance."""
        self.monsoon = Monsoon
        self.__mainCal = calibrationData()
        self.__usbCal = calibrationData()
        self.__auxCal = calibrationData()
        self.__padding = np.zeros((64))
        self.__fineThreshold = Monsoon.fineThreshold
        self.__auxFineThreshold = Monsoon.auxFineThreshold
        self.__ADCRatio = (float)(62.5 / 1e6); #Each tick of the ADC represents this much voltage
        self.__mainVoltageScale = Monsoon.mainvoltageScale
        self.__usbVoltageScale = Monsoon.usbVoltageScale
        self.dropped = 0
        self.bulkProcessRate = 128
        self.__packetSize = 64
        self.__startTime = time.time()
        #Indices
        self.__mainCoarseIndex = 0
        self.__mainFineIndex = 1
        self.__usbCoarseIndex = 2
        self.__usbFineIndex = 3
        self.__auxCoarseIndex = 4
        self.__auxFineIndex = 5
        self.__mainVoltageIndex = 6
        self.__usbVoltageIndex = 7
        self.__timestampIndex = 10
       
        #Output lists
        self.__mainCurrent = []
        self.__usbCurrent = []
        self.__auxCurrent = []
        self.__usbVoltage = []
        self.__mainVoltage = []
        self.__timeStamps = []

        #Output controls
        self.__outputConsoleMeasurements = False
        self.__outputTimeStamp = True
        self.__collectMainMeasurements = True
        self.__collectUSBMeasurements = False
        self.__collectAuxMeasurements = False
        self.__collectMainVoltage = True
        self.__collectUSBVoltage = False
        self.__channels = [self.__outputTimeStamp, self.__collectMainMeasurements,self.__collectUSBMeasurements,self.__collectAuxMeasurements,self.__collectMainVoltage,self.__collectUSBVoltage]
        self.__channelnames = ["Time(ms)","Main(mA)", "USB(mA)", "Aux(mA)", "Main Voltage(V)", "USB Voltage(V)"]
        self.__channelOutputs = [self.__mainCurrent,self.__usbCurrent,self.__auxCurrent,self.__mainVoltage,self.__usbVoltage]
        self.sampleCount = 0
        self.__CSVOutEnable = False

        #output writer
        self.__f = None

        pass

    def ConsoleOutput(self, bool):
        """Enables or disables the display of realtime measurements"""
        self.__outputConsoleMeasurements = bool

    def enableChannel(self,channel):
        """Enables a channel.  Takes sampleEngine.channel class value as input."""
        self.__channels[channel] = True 
        pass

    def disableChannel(self,channel):
        """Disables a channel.  Takes sampleEngine.channel class value as input."""
        self.__channels[channel] = False
        pass
    def enableCSVOutput(self, filename):
        """Opens a file and causes the sampleEngine to periodically output samples when taking measurements
        filename: The file measurements will be output to."""
        self.__f = open(filename,"w")
        self.__CSVOutEnable = True

    def disableCSVOutput(self):
        """Closes the CSV file if open and disables CSV output."""
        if(self.__f is not None):
            self.__f.close()
            self.__f = None
        self.__CSVOutEnable = False


    def __Clear(self):
        """Wipes away all of the old output data."""
        self.__mainCurrent = []
        self.__usbCurrent = []
        self.__auxCurrent = []
        self.__usbVoltage = []
        self.__mainVoltage = []
        self.sampleCount = 0
        self.__mainCal.clear()
        self.__usbCal.clear()
        self.__auxCal.clear()

    def __isCalibrated(self):
        """Returns true if every channel has sufficient calibration samples."""
        A = self.__mainCal.calibrated
        B = self.__usbCal.calibrated
        C = self.__auxCal.calibrated
        return A and B and C

    def __vectorProcess(self,measurements):
        """Translates raw ADC measurements into current values."""

       #Currents
        if(self.__isCalibrated()):
            measurements = np.array(measurements)
            output = []
            sDebug = ""
            if(self.__channels[channels.MainCurrent]):
            #Main Coarse
                scale = self.monsoon.statusPacket.mainCoarseScale
                zeroOffset = self.monsoon.statusPacket.mainCoarseZeroOffset
                calRef = self.__mainCal.getRefCal(True)
                calZero = self.__mainCal.getZeroCal(True)
                zeroOffset += calZero
                if(calRef - zeroOffset != 0):
                    slope = scale / (calRef - zeroOffset)
                else:
                    slope = 0
                Raw = measurements[:,self.__mainCoarseIndex] - zeroOffset
                mainCoarseCurrents = Raw * slope 
        
                #Main Fine
                scale = self.monsoon.statusPacket.mainFineScale
                zeroOffset = self.monsoon.statusPacket.mainFineZeroOffset
                calRef = self.__mainCal.getRefCal(False)
                calZero = self.__mainCal.getZeroCal(False)
                zeroOffset += calZero
                if(calRef - zeroOffset != 0):
                    slope = scale / (calRef - zeroOffset)
                else:
                    slope = 0
                Raw = measurements[:,self.__mainFineIndex] - zeroOffset
                mainFinecurrents = Raw * slope / 1000
                mainCurrent = np.where(measurements[:,self.__mainFineIndex] < self.__fineThreshold, mainFinecurrents, mainCoarseCurrents)
                self.__mainCurrent.append(mainCurrent)
                sDebug = "Main Current: " + repr(round(mainCurrent[0],2))

            if(self.__channels[channels.USBCurrent]):
                #USB Coarse
                scale = self.monsoon.statusPacket.usbCoarseScale
                zeroOffset = self.monsoon.statusPacket.usbCoarseZeroOffset
                calRef = self.__usbCal.getRefCal(True)
                calZero = self.__usbCal.getZeroCal(True)
                zeroOffset += calZero
                if(calRef - zeroOffset != 0):
                    slope = scale / (calRef - zeroOffset)
                else:
                    slope = 0
                Raw = measurements[:,self.__usbCoarseIndex] - zeroOffset
                usbCoarseCurrents = Raw * slope 

                #USB Fine
                scale = self.monsoon.statusPacket.usbFineScale
                zeroOffset = self.monsoon.statusPacket.usbFineZeroOffset
                calRef = self.__usbCal.getRefCal(False)
                calZero = self.__usbCal.getZeroCal(False)
                zeroOffset += calZero
                if(calRef - zeroOffset != 0):
                    slope = scale / (calRef - zeroOffset)
                else:
                    slope = 0
                Raw = measurements[:,self.__usbFineIndex] - zeroOffset
                usbFineCurrents = Raw * slope/ 1000
                usbCurrent = np.where(measurements[:,self.__usbFineIndex] < self.__fineThreshold, usbFineCurrents, usbCoarseCurrents)
                self.__usbCurrent.append(usbCurrent)
                sDebug = sDebug + " USB Current: " + repr(round(usbCurrent[0], 2))
        
            if(self.__channels[channels.AuxCurrent]):
                #Aux Coarse
                scale = self.monsoon.statusPacket.mainFineScale
                zeroOffset = 0
                calRef = self.__auxCal.getRefCal(True)
                calZero = self.__auxCal.getZeroCal(True)
                zeroOffset += calZero
                if(calRef - zeroOffset != 0):
                    slope = scale / (calRef - zeroOffset)
                else:
                    slope = 0
                Raw = measurements[:,self.__auxCoarseIndex] - zeroOffset
                auxCoarseCurrents = Raw * slope 

                #Aux Fine
                scale = self.monsoon.statusPacket.auxFineScale
                zeroOffset = 0
                calRef = self.__auxCal.getRefCal(False)
                calZero = self.__auxCal.getZeroCal(False)
                zeroOffset += calZero
                if(calRef - zeroOffset != 0):
                    slope = scale / (calRef - zeroOffset)
                else:
                    slope = 0
                Raw = measurements[:,self.__auxFineIndex] - zeroOffset
                auxFineCurrents = Raw * slope / 1000
                auxCurrent = np.where(measurements[:,self.__auxFineIndex] < self.__auxFineThreshold, auxFineCurrents, auxCoarseCurrents)
                self.__auxCurrent.append(auxCurrent)
                sDebug = sDebug + " Aux Current: " + repr(round(auxCurrent[0], 2))

            #Voltages
            if(self.__channels[channels.MainVoltage]):
                mainVoltages = measurements[:,self.__mainVoltageIndex] * self.__ADCRatio * self.__mainVoltageScale
                self.__mainVoltage.append(mainVoltages)
                
                sDebug = sDebug + " Main Voltage: " + repr(round(mainVoltages[0],2))
                

            if(self.__channels[channels.USBVoltage]):
                usbVoltages = measurements[:,self.__usbVoltageIndex] * self.__ADCRatio * self.__usbVoltageScale
                self.__usbVoltage.append(usbVoltages)
                sDebug = sDebug + " USB Voltage: " + repr(round(usbVoltages[0],2))
            timeStamp = measurements[:,self.__timestampIndex]
            self.__timeStamps.append(timeStamp)
            sDebug = sDebug + " Dropped: " + repr(self.dropped)
            if(self.__outputConsoleMeasurements):
                print(sDebug)


    def __processPacket(self, measurements):
        """Separates received packets into ZeroCal, RefCal, and measurement samples."""
        Samples = []
        for measurement in measurements:
            self.dropped = measurement[0]
            flags = measurement[1]
            numObs = measurement[2]
            offset = 3
            for i in range(0,numObs):
                sample = measurement[offset:offset+10]
                sample.append(measurement[len(measurement)-1])
                sampletype = sample[8] & 0x30
                if(sampletype == ops.SampleType.ZeroCal):
                    self.__processZeroCal(sample)
                elif(sampletype == ops.SampleType.refCal):
                   self.__processRefCal(sample)
                elif(sampletype == ops.SampleType.Measurement):
                    Samples.append(sample)
                    
                offset += 10
        return Samples

    def __processZeroCal(self,meas):
        """Adds raw measurement data to the zeroCal tracker"""
        self.__mainCal.addZeroCal(meas[self.__mainCoarseIndex], True)
        self.__mainCal.addZeroCal(meas[self.__mainFineIndex], False)
        self.__usbCal.addZeroCal(meas[self.__usbCoarseIndex], True)
        self.__usbCal.addZeroCal(meas[self.__usbFineIndex], False)
        self.__auxCal.addZeroCal(meas[self.__auxCoarseIndex], True)
        self.__auxCal.addZeroCal(meas[self.__auxFineIndex], False)
        return True
    def __processRefCal(self, meas):
        """Adds raw measurement data to the refcal tracker"""
        self.__mainCal.addRefCal(meas[self.__mainCoarseIndex], True)
        self.__mainCal.addRefCal(meas[self.__mainFineIndex], False)
        self.__usbCal.addRefCal(meas[self.__usbCoarseIndex], True)
        self.__usbCal.addRefCal(meas[self.__usbFineIndex], False)
        self.__auxCal.addRefCal(meas[self.__auxCoarseIndex], True)
        self.__auxCal.addRefCal(meas[self.__auxFineIndex], False)
        return True

    def getSamples(self):
        """Returns samples in a Python list.  Format is [timestamp, main, usb, aux, mainVolts,usbVolts].  Only includes enabled channels."""
        result = self.__arrangeSamples()
        return result

    def __outputToCSV(self):
        """This is intended to be called periodically during sampling.  
        The alternative is to store measurements in an array or queue, which will overflow allocated memory within a few hours depending on system settings.
        Writes measurements to a CSV file"""

        output = self.__arrangeSamples()
        for i in range(len(output[0])):
            sOut = ""
            for j in range(len(output)):
                sOut = sOut + repr(output[j][i]) + ","
            sOut = sOut + "\n"
            self.__f.write(sOut)

    def __arrangeSamples(self):
        """Arranges output lists so they're a bit easier to process."""
        output = []
        times = []
        for data in self.__timeStamps:
            for measurement in data:
                times.append(measurement)
        output.append(times)
        self.__timeStamps = []
        if(self.__channels[channels.MainCurrent]):
            main = []
            for data in self.__mainCurrent:
                for measurement in data:
                    main.append(measurement)
            output.append(main)
            self.__mainCurrent = []
        if(self.__channels[channels.USBCurrent]):
            usb = []
            for data in self.__usbCurrent:
                for measurement in data:
                    usb.append(measurement)
            output.append(usb)
            self.__usbCurrent = []
        if(self.__channels[channels.AuxCurrent]):
            Aux = []
            for data in self.__auxCurrent:
                for measurement in data:
                    Aux.append(measurement)
            output.append(Aux)
            self.__auxCurrent = []
        if(self.__channels[channels.MainVoltage]):
            volts = []
            for data in self.__mainVoltage:
                for measurement in data:
                    volts.append(measurement)
            output.append(volts)
            self.__mainVoltage = []
        if(self.__channels[channels.USBVoltage]):
            volts = []
            for data in self.__usbVoltage:
                for measurement in data:
                    volts.append(measurement)
            output.append(volts)
            self.__usbVoltage = []
        return output
    def __outputCSVHeaders(self):
        """Creates column headers in the CSV output file for each enabled channel."""
        for i in range(len(self.__channelnames)):
            if(self.__channels[i]):
                self.__f.write((self.__channelnames[i] + ","))
        self.__f.write("\n")
    def startSampling(self, samples=5000,granularity = 1):
        """Starts sampling."""
        """samples:  The number of samples you would like to collect before stopping."""
        """granularity: Controls the resolution at which samples are stored.  1 = all samples stored, 10 = 1 out of 10 samples stored, etc."""
        self.__Clear()
        self.monsoon.StartSampling(1250,0xFFFFFFFF)
        Samples = [[0 for x in range(self.__packetSize+1)] for y in range(self.bulkProcessRate)]
        S = 0
        debugcount = 0
        minutes = 0
        granularity_index = 0
        self.__startTime = time.time()
        if(self.__CSVOutEnable):
            self.__outputCSVHeaders()
        while self.sampleCount < samples:
            buf = self.monsoon.BulkRead()
            Sample = self.monsoon.swizzlePacket(buf)
            numSamples = Sample[2]
            granularity_index += numMeasurements
            self.sampleCount += numMeasurements
            if(granularity_index >= granularity):
                granularity_index = 0
                Sample.append(time.time() - self.__startTime)
                Samples[S] = Sample
                S += 1
                if(S >= self.bulkProcessRate):
                    bulkPackets = self.__processPacket(Samples)
                    if(len(bulkPackets) > 0):
                        self.__vectorProcess(bulkPackets)
                    S = 0
                if(S % self.bulkProcessRate/2 == 0 and self.__CSVOutEnable):
                    self.__outputToCSV()
        self.monsoon.stopSampling()
        if(self.__CSVOutEnable):
            self.__outputToCSV()
            self.disableCSVOutput()
        pass

