import numpy as np
import re, struct

# hyst as a float
# always put positive voltages first

PARAMS = ["ZPiezoconst", "LockinRC", "Vertmandelay", "VertSpecBack", "LockinAmpl", "Current[A]", "SetPoint", "DSP_Clock",
            "DAC-Type", "Gainpreamp", "Length x[A]", "Length y[A]"]
re_string = f"({'|'.join(PARAMS)})=(.+)\\n"
re_string = re_string.replace('[', r'\[').replace(']', r'\]')
param_re = re.compile(re_string)

class STM:
    def __init__(self, DAC_Voltage_Range = 20000.0):
        self.DAC_Voltage_Range = DAC_Voltage_Range #in mV
    def __str__(self):
        ret = 'DAC_Voltage_Range: %.1f mV' %self.DAC_Voltage_Range
        return ret

def shift(arr, n):
    # shifts a 1D array n positions to the right. Repeats starting/ending values for positions outside original array bounds.
    # ex:  [1, 2, 3, 4, 5, 6, 7] shifted n=2
    # out: [1, 1, 1, 2, 3, 4, 5]
    n = int(n)
    start, end = arr[0], arr[-1]
    out = np.zeros(arr.shape)
    if n > 0:
        out[n:] = arr[:-n]
        out[:n] = start
    elif n < 0:
        out[:n] = arr[-n:]
        out[n:] = end
    else:
        out = arr
    return out

class Spectra:
    def __init__(self, V, I, dIdV):
        self.label = None
        self.V, self.I, self.dIdV = V, I, dIdV
        self.N = len(V)
        self.z = np.empty(self.N)

        self.ZPiezoconst = 1.0
        self.LockinRC = 0.0
        self.Speclength = 0
        self.VertSpecBack = 0
        self.LockinAmpl = 0
        self.Current = 1
        self.hyst = 0
        self.x = 0
        self.y = 0
        self.Length_x = 0
        self.Length_y = 0

    #use command print(spectra_object) to print out all relevant parameters
    def __str__(self):
        ret = f'Spectra object at label: {self.label}\n'
        for key in self.__dict__.keys():
            if key.startswith('__'):
                continue
            x = getattr(self, key)
            if isinstance(x, int) or isinstance(x, float) or isinstance(x, str) or x is None:
                ret += f'{key}: {x}\n'
        return ret

    def average(self, hyst_cor=True):
        # this function averages trace and retraces together, editing the attributes 'V', 'I', 'dIdV', and 'z' to reflect this.
        VSB = self.VertSpecBack
        HSYT = self.hyst

        # Ensure at least 1 retrace exists.
        if VSB + 1 <= 1:
            raise Exception('VertSpecBack is %d, which means no averaging or hystersis correction is possible.' %n)
        n = self.N // (VSB + 1)
        self.V = self.V[0:n]

        # Ensure # data points is a multiple of # of trace/retraces. Removes points from the end accordingly.
        dif = self.N % (VSB + 1)
        if dif != 0:
             self.I = self.I[:-dif]
             self.dIdV = self.dIdV[:-dif]
             self.z = self.z[:-dif]
             self.N = len(self.V)

        # corrects for hysteresis (constant shift)
        if hyst_cor:
            if HSYT % 2 != 0:
                HSYT += 1
            self.I = shift(self.I, -HSYT/2)
            self.dIdV = shift(self.dIdV, -HSYT/2)
            self.z = shift(self.z, -HSYT/2)

        # Reshape single-row data arrays into a matrix of multiple rows, with single row containing exactly one trace.
        self.I = self.I.reshape((VSB + 1, n))
        self.dIdV = self.dIdV.reshape((VSB + 1, n))
        self.z = self.z.reshape((VSB + 1, n))

        # Reverse "back-traces", which are odd-numbered rows. (Forward/backward scans always alternate)
        for i in range(1, VSB + 1, 2): # i.e. odd numbers only
            self.I[i, :] = self.I[i, :][::-1]
            self.dIdV[i, :] = self.dIdV[i, :][::-1]
            self.z[i, :] = self.z[i, :][::-1]

        # sum down column, divide by number of rows.
        self.I = (1 / (VSB + 1)) * np.sum(self.I, axis=0)
        self.dIdV = (1 / (VSB + 1)) * np.sum(self.dIdV, axis=0)
        self.z = (1 / (VSB + 1)) * np.sum(self.z, axis=0)

        if self.V[0] < 0:
            self.V = self.V[::-1]
            self.I = self.I[::-1]
            self.dIdv = self.dIdV[::-1]
            self.z = self.z[::-1]

        
    def normalize(self, kappa):
        #normalize the spectra using an input value of kappa in Ang^-1
        self.I = self.I * np.exp(-2 * kappa * self.z)
        self.dIdV = self.dIdV * np.exp(-2* kappa * self.z)

def load_VERT_file(filename, stm = STM()):
    # Now uses regex to search for parameters of interest. Regex object is defined above

    # find all the relevant parameters
    f = open(filename, 'r')
    version = f.readline()
    file = f.read()
    f.close()

    if "ParVERT32" in version:
        Version = '3.2' #STMAFM Software version
        LockinRC_factor = 7.989578/0.00251 #mutliplicative factor to get LockinRC param in Hz
        Current_index = 4 #column indexes for extracting current and dIdV
        dIdV_index = 5
    if "ParVERT30" in version:
        Version = '3.0'
        LockinRC_factor = 7960.
        Current_index = 3
        dIdV_index = 4
    
    file = file.split('DATA') # makes list of len 2: first elem is params, second elem is data
    params = param_re.findall(file[0])
    p_dict = dict()
    for p, v in params:
        try:
            p_dict[p] = float(v)
        except ValueError:
            if p == 'DAC-Type':
                p_dict[p] = float(v[:-3])  # removing 'bit' from end of string
            else:
                raise Exception(f'Parameter {p} gives an improper value {v}')

    data = np.fromstring(file[1], sep='\t')
    N, x, y = data[0], data[1], data[2]
    final_elem = int(N * 5)  # to remove comments at the end
    data = data[4:final_elem + 4].reshape((int(N), 5))
                   
    z = data[:, 2] * p_dict["ZPiezoconst"] / 1000.0 # Angstroms
    V = data[:, 1] / 1000.0 # Volts
    I = data[:, Current_index] / 2**p_dict["DAC-Type"] * stm.DAC_Voltage_Range / 10**p_dict["Gainpreamp"] / 1e6 # current in amps 
    dIdV = data[:, dIdV_index] / 2**p_dict["DAC-Type"] * stm.DAC_Voltage_Range # conductance in meV

    ret = Spectra(V, I, dIdV)
    ret.z = z
    
    label = filename[-19:-5]
    if 'L' in label or 'R' in label:
        label = filename[-25:-5]
    ret.label = label

    ret.__dict__['ZPiezoconst'] = p_dict["ZPiezoconst"]
    ret.__dict__['LockinRC'] = p_dict["LockinRC"] * LockinRC_factor # in Hz
    ret.__dict__['Speclength'] = p_dict["Vertmandelay"] * N / p_dict["DSP_Clock"]
    ret.__dict__['VertSpecBack'] = int(p_dict["VertSpecBack"])
    ret.__dict__['LockinAmpl'] = p_dict["LockinAmpl"]
    try:
        ret.__dict__['Current'] = p_dict["Current[A]"]
    except KeyError:
        ret.__dict__['Current'] = p_dict['SetPoint']
    ret.N = int(N)
    ret.__dict__['x'] = x * 10./2.**(20-1) * 40 * 10 # value in DAC units * Voltage range / 2^(bit count - 1 (account for positive vs negative)) * piezo constant * high voltage gain (hard coded for now)
    ret.__dict__['y'] = y * 10./2.**(20-1) * 40 * 10
    ret.__dict__['Length_x'] = p_dict["Length x[A]"]
    ret.__dict__['Length_y'] = p_dict["Length y[A]"]

    #hysteresis correction factor
    hyst = int(N / (ret.__dict__['Speclength'] * p_dict["LockinRC"] * LockinRC_factor * np.pi))
    ret.__dict__['hyst'] = hyst

    return ret

# UNCHANGED BELOW

class specgrid:
    def __init__(self, V, I, dIdV):
        self.label = None

        #data should include V, z, I, and dIdV in that order
        
        #Assign values to spectra object
        self.ZPiezoconst = 1.0
        self.LockinRC = 0.0
        self.Speclength = 0
        self.VertSpecBack = 0
        self.LockinAmpl = 0
        self.Current = 1

        self.hyst = 0

        self.z = np.ones(len(V))
        self.V = V
        
        self.I = I
        self.dIdV = dIdV
        
        #Number of data points, needed for averaging and hysteresis correction
        N = len(V)
        self.N = N

        self.nx, self.ny = I.shape[0], I.shape[1]

    def average(self, hyst_cor = True):

        N = self.N 
        if self.VertSpecBack + 1 <= 1:
            raise Exception('VertSpecBack is %d, which means no averaging or hystersis correction is possible.' %n)
        
        I = np.empty((self.nx, self.ny, self.N/(self.VertSpecBack + 1)))
        dIdV = np.empty((self.nx, self.ny, self.N/(self.VertSpecBack + 1)))
        for i in range(self.nx):
            for j in range(self.ny):
                spec = spectra(self.V, self.I[i, j, :], self.dIdV[i, j, :])
                spec.z = self.z
                spec.hyst = self.hyst
                spec.VertSpecBack = self.VertSpecBack
                spec.average(hyst_cor = hyst_cor)
                I[i, j, 0:len(spec.I)] = spec.I
                dIdV[i, j, 0:len(spec.dIdV)] = spec.dIdV
        self.z = spec.z
        self.V = spec.V
        self.I = I
        self.dIdV = dIdV

    def normalize(self, kappa):
        I = np.empty((self.nx, self.ny, self.N))
        dIdV = np.empty((self.nx, self.ny, self.N))
        for i in range(self.nx):
            for j in range(self.ny):
                spec = spectra(self.V, self.I[i, j, :], self.dIdV[i, j, :])
                spec.z = self.z
                spec.normalize(kappa)
                I[i, j, :] = spec.I
                dIdV[i, j, :] = spec.dIdV
        self.I = I
        self.dIdV = dIdV
        
    #use command print(spectra_object) to print out a few relevant parameters
    def __str__(self):
        ret = 'label: ' + self.label +'\n'
        ret += 'nx, ny: %d, %d\n' %(self.nx, self.ny)
        ret += 'data points: %d\n' %self.N
        ret += 'VertSpecBack: %d\n' %self.VertSpecBack
        ret += 'LockinAmpl: %.1f mV\n' %self.LockinAmpl
        ret += 'Setpoint: %.1e A\n' %self.Current
        ret += 'Hystersis Correction: %d' %self.hyst
        return ret

def read_item(content, loc, format_):
    size = struct.calcsize(format_)
    return struct.unpack(format_, content[loc:loc+size])[0]
        
def load_specgrid_file(filename, stm = STM()):
    #find all the relevant parameters
    f = open(filename + '.dat', 'r')
    for line in f:
        if "ParVERT32" in line:
            Version = '3.2' #STMAFM Software version
            LockinRC_factor = 7960. #mutliplicative factor to get LockinRC param in HzLockinRC_factor = 7960.
            Current_index = 5 #column indexes for extracting current and dIdV
            dIdV_index = 4
        if "ParVERT30" in line:
            Version = '3.0'
            LockinRC_factor = 7960.
            Current_index = 4
            dIdV_index = 3
        if "ZPiezoconst" in line:
            ZPiezoconst = float(line[line.find('=')+1:-1])
        if "LockinRC" in line:
            LockinRC = float(line[line.find('=')+1:-1])
        if "Vertmandelay" in line:
            Vertmandelay = float(line[line.find('=')+1:-1])
        if "VertSpecBack" in line:
            VertSpecBack = int(line[line.find('=')+1:-1])
        if "LockinAmpl" in line:
            LockinAmpl = float(line[line.find('=')+1:-1]) #mV
        if "Current[A]" in line:
            Current = float(line[line.find('=')+1:-1]) #Amps
        if "DSP_Clock" in line:
            DSP_Clock = float(line[line.find('=')+1:-1]) #DSP Units to determine Speclength
        if "DAC-Type" in line:
            DAC_Type = float(line[line.find('=')+1:-4])
        if "Gainpreamp" in line:
            Gainpreamp = float(line[line.find('=')+1:-1])
    f.close()

    with open(filename, mode='rb') as file: # b is important -> binary
        content = file.read()
    version = read_item(content, 0, 'i')    
    nx, ny = read_item(content, 4, 'i'), read_item(content, 8, 'i')
    dx, dy = read_item(content, 12, 'i'), read_item(content, 16, 'i')
    specxgrid, specygrid = read_item(content, 20, 'i'), read_item(content, 24, 'i')
    vertpoints = read_item(content, 28, 'i')
    vertmandelay = read_item(content, 32, 'i')
    vertmangain = read_item(content, 36, 'i')
    biasvoltage = read_item(content, 40, 'f')
    tunnelcurrent = read_item(content, 44, 'f')
    imagedatasize = read_item(content, 48, 'i')
    specgriddatasize = read_item(content, 52, 'i') 
    specgridchan = read_item(content, 56, 'i')
    specgridchannelselectval = read_item(content, 60, 'i')
    specgriddatasize64 = read_item(content, 64, 'q')
    if version >= 4:
        xstart, xend = read_item(content, 72, 'i'), read_item(content, 76, 'i')
        ystart, yend = read_item(content, 80, 'i'), read_item(content, 84, 'i')
    else:
        xstart, ystart = 1, 1
        xend = int(nx/specxgrid) + 1
        yend = int(ny/specygrid) + 1
    V = np.empty(vertpoints)
    z = np.empty(vertpoints)
    for i in range(vertpoints):
        V[i] = read_item(content, 1024 + (i*8), 'f')
        z[i] = read_item(content, 1028 + (i*8), 'f')
    
    current = np.empty((xend-xstart + 1, yend-ystart + 1, vertpoints))
    conductance = np.empty((xend-xstart + 1, yend-ystart + 1, vertpoints))
    n = 0
    for i in range(current.shape[0]):
        for j in range(current.shape[1]):
            I = np.empty(vertpoints)
            dIdV = np.empty(vertpoints)
            for k in range(vertpoints):
                I[k] = read_item(content, 1032 + n*vertpoints*8 + k*8, 'f')
                dIdV[k] = read_item(content, 1028 + n*vertpoints*8 + k*8, 'f')
            current[i, j, :] = I
            conductance[i, j, :] = dIdV
            n += 1

    current *= stm.DAC_Voltage_Range / 2**DAC_Type / 10**Gainpreamp /1e6 #current in amps
    conductance *= stm.DAC_Voltage_Range /2**DAC_Type #meV
    V /= 1000.0 #volts
    ret = specgrid(V, current, conductance)
    ret.z = z * ZPiezoconst/1000.0 #Angstroms

    Speclength = Vertmandelay * vertpoints / DSP_Clock
    
    ret.label = filename[-23:-9]
    ret.ZPiezoconst = ZPiezoconst
    ret.LockinRC = LockinRC * LockinRC_factor #Hz
    ret.Speclength = Speclength
    ret.VertSpecBack = VertSpecBack
    ret.LockinAmpl = LockinAmpl
    ret.Current = Current
    ret.N = vertpoints
    #hysteresis correction factor
    hyst = int(vertpoints/Speclength/LockinRC/LockinRC_factor/np.pi)
    ret.hyst = hyst

    return ret
