import numpy as np
import matplotlib.pyplot as plt
import struct
import os, sys
import re
import copy

class Matrix:
    """
    Class to Read and Hangle Matrix files
    """
    def __init__(self,Path): # Give the Path of the folder containing all the mtrx files
        # Read PATH and open file
        self.Path = Path
        self.fp = None # file variable
        for x in os.listdir(Path): # List the folder and look for the _0001.mtrx file
            if x[-10:] == "_0001.mtrx":
                self.fp = open(self.Path+"/"+x, "rb")
        if self.fp == None:
            print("Matrix file not found!")
            sys.exit(1)
        if self.fp.read(8) != b"ONTMATRX": # header of the file
            print("Unknown header! Wrong Matrix file format")
            sys.exit(2)
        self.version = self.fp.read(4) # should be 0101
        self.IDs = {}
        self.params = {} # dictionary to list all the parameters
        self.images = {} # images[x] are the parameters used during the record for file named x

        # Parse the file and read the block
        while True: # While not EOF scan files and read block
            r = self.read_block()
            if r == False:
                break

    def read_string(self):
        """
        Strings are stored as UTF-16. First 32-bits is the string length
        """
        N = struct.unpack("<L", self.fp.read(4))[0] # string length
        if N == 0:
            return ""
        s = self.fp.read(N*2).decode('utf-16')
        return s

    def plotSTS(self, ID, num=1): # plot STS file called xxx--ID_num.I(V)_mtrx
        x, y = self.getSTS(ID, num)
        plt.plot(x, y)
        plt.show()
    
    def getUpDown(self, X, Y, NPTS):
        """
        Split data in Up and Down measurement, pad them with NaN if necessary and return them in increasing order.
        The returned value are X,Yup, Ydown
        If Up/Down data are missing an empty array will be returned
        """
        if len(Y) < NPTS: # Missing data
            Y = np.pad(Y, NPTS, 'constant', constant_values=np.nan)
        elif len(Y) > NPTS: # Forward and backward scans
            if len(Y) < 2*NPTS: # Missing data
                Y = np.pad(Y, 2*NPTS, 'constant', constant_values=np.nan)
            if X[NPTS-1] < X[0]:
                return X[NPTS:], [Y[NPTS:], Y[NPTS-1::-1]]
            else:
                return X[:NPTS], [Y[:NPTS], Y[-1:NPTS-1:-1]]
        if X[-1] < X[0]:
            return X[::-1], [np.empty(NPTS), Y[::-1], np.empty(NPTS)]
        return X, [Y, np.empty(NPTS)]

        

    def getSTSData(self, ID, nums=[1]):
        if not ID in self.IDs or len(nums) < 1:
            return None
        # retrieve the spectroscopy data (V, I and an object IM containing the parameters)
        V, I, IM = self.getSTS(ID, nums[0], params=True)
        NPTS = int(IM['Spectroscopy']['Device_1_Points']['value'])
        hasDI = self.IDs[ID]['hasDI']
        # Call the function to split and flip data if it's UP/Down measurements
        V, I = self.getUpDown(V, I, NPTS)
        for num in nums[1:]: # Skip first num as it's already parsed above
            X, Y = self.getUpDown(*self.getSTS(ID, num), NPTS=NPTS)
            if not np.array_equal(V, X):
                raise Exception("Bias axis differs between measurements?!?")
            for i in range(2): # i=0: Up scan, i=1: Down scan
                I[i] = np.vstack((I[i], Y[i]))
        Im = [np.nan]*2   # Store the mean of I
        Ims = [np.nan]*2  # Store StDev of I
        for i in range(2): # i=0: Up scan, i=1: Down scan
            Im[i] = I[i].mean(axis=0)
            Ims[i] = I[i].std(axis=0)