import json
import base64
import numpy as np
import scipy.fft as fft
import matplotlib.pyplot as plt

class DataDecoder:
    def __init__(self, filename):
        self.data = ''
        with open(filename, 'r') as file:
            self.data = file.read()
        self.structed_data = json.loads(self.data)

    def getRawData(self, averaging_num=0, data_num=0, channel_num=0):
        for items in self.structed_data:
            if(items['averaging_num'] == averaging_num and items['data_num'] == data_num):
                for channel_data in items['channel_data']:
                    if(channel_data['channel_num'] == channel_num):
                        return channel_data['channel_data']
        return None

    def getDataDecoded(self, averaging_num=0, data_num=0, channel_num=0, points=10):
        rawData = self.getRawData(averaging_num, data_num, channel_num)
        edata = base64.b64decode(rawData.encode('utf-8'))
        arr = np.frombuffer(edata, dtype=np.int16, count=points, offset=0)
        return arr

    def getDataScaled(self, averaging_num=0, data_num=0, channel_num=0, points=10, range=5.0):
        decodedData = self.getDataDecoded(averaging_num, data_num, channel_num, points)
        return range * decodedData / 32768
    
    def getDataRate(self, averaging_num=0, data_num=0):
        for items in self.structed_data:
            if(items['averaging_num'] == averaging_num and items['data_num'] == data_num):
                return items['measurement_rate']

    def getDataSpectrum(self, averaging_num=0, data_num=0, channel_num=0, points=10, range=5.0, zero_fill=0, first_idx=0, last_idx=10):
        scaledData = self.getDataScaled(averaging_num, data_num, channel_num, points, range)
        zerofilledData = np.append(scaledData[first_idx:last_idx], np.zeros(zero_fill))
        rate = self.getDataRate()
        transferedData = fft.rfft(zerofilledData) * 2 / (last_idx-first_idx)
        freqs = fft.rfftfreq((last_idx-first_idx)+zero_fill, 1/rate)
        spectrum = np.abs(transferedData)
        phases = np.angle(transferedData)
        spect_dict = {'freqs': freqs,
                     'phases': phases,
                     'spectrum': spectrum}
        return spect_dict
        
    def getDataPoints(self, averaging_num=0, data_num=0):
        for items in self.structed_data:
            if(items['averaging_num'] == averaging_num and items['data_num'] == data_num):
                return items['measurement_points']


#dec = DataDecoder('fid_test.json')

#print('getRawData(0, 0, 1):')
#print(dec.getRawData(0, 0, 1)) # ::getRawData(averagingIndex, dataTriggerIndex, channelIndex)
#print('\n')

#print('getDataDecoded(0, 0, 0, 100):')
#a = dec.getDataDecoded(0, 0, 0, 100)
#print(a[0:100]) # ::getDataDecoded(averagingIndex, dataTriggerIndex, channelIndex, count)
#print('\n')

#print('getDataScaled(0, 1, 1, 100, 0.1):')
#a = dec.getDataDecoded(0, 0, 1, 80000)
#rate = dec.getDataRate(0, 0)
#print(a[500:1000])
#t = np.arange(0, 80000 / rate, 1 / rate)
#plt.plot(t, a)
#plt.show()
#print('\n')

#print('getDataSpectrum(0, 1, 1, max_points, 5.0, max_points):')
#points = dec.getDataPoints(0, 0)
#spect_dict = dec.getDataSpectrum(0, 0, 1, points[0], 0.1, 0, 5000, 19600)
#plt.plot(spect_dict['freqs'], spect_dict['spectrum'])
#plt.show()
#print('\n')