import numpy as np
import matplotlib.pyplot as plt
import glob

def readSensorArray(basename):
    files = np.array(glob.glob('%s-*'%basename))
    if len(files) == 0:
        print("No files found for input '%s'"%basename)
        raise SystemExit
    # sort files correctly
    It = np.array([int(f.split('-')[-1]) for f in files]) # time index is last part in filename-n
    Is = np.argsort(It)
    files = files[Is]
    steps = np.size(files)

       
    # column 1: origElemNum
    # colum 2-4: global coordinates (x, y, z)
    # column 5-7: real mechanic Displacement (x, y, z)
    # column 8-10: imaginary mechanic Displacement (x, y, z)
    # column 11-13: local coordinates (xi, leta, zeta)
    sensorCoordinates = np.loadtxt(files[0], delimiter=',', skiprows=1, usecols=(1,2,3)) 
    sensorResults = np.empty((steps,np.shape(sensorCoordinates)[0],3),dtype=complex)
    for i,f in enumerate(files):
        sensorResults[i,:,:] = np.loadtxt(f, delimiter=',', skiprows=1, usecols=(4,5,6)) + 1j * np.loadtxt(files[0], delimiter=',', skiprows=1, usecols=(7, 8, 9)) 
    return (sensorCoordinates,sensorResults)

if __name__ == "__main__":
    filename_reduced = "results_txt/centerLine_reduced"
    filename = "results_txt/centerLine"

    (sensorCoordinates_reduced,sensorResults_reduced) = readSensorArray(filename_reduced)
    (sensorCoordinates,sensorResults) = readSensorArray(filename)

    idxCenterX_reduced = np.argwhere(sensorCoordinates_reduced[:,0] == np.max(sensorCoordinates_reduced[:,0])/2)
    idxCenterX = np.argwhere(sensorCoordinates[:,0] == np.max(sensorCoordinates[:,0])/2)

    centerDisplacementX_reduced =np.squeeze(sensorResults_reduced[:,idxCenterX_reduced,1])
    centerDisplacementX =np.squeeze(sensorResults[:,idxCenterX,1])

    frequencies_reduced = np.linspace(1,1000,sensorResults_reduced.shape[0])
    frequencies = np.linspace(1,1000,sensorResults.shape[0])

    fig, ax = plt.subplots()
    ax.loglog(frequencies_reduced, np.abs(centerDisplacementX_reduced),label="reduced model")
    ax.loglog(frequencies, np.abs(centerDisplacementX),label="full model")
    ax.legend()
    ax.set_title("Transferfunction from the pressure to the x-displacement \n of the centerpoint on the centerline")
    ax.set_ylabel(r'$|\frac{\Delta x_c}{p}|$ in $\frac{\mathrm{mm}}{\mathrm{bar}}$')
    ax.set_xlabel("Frequency in Hz")
    plt.savefig("transferfunction.png")
    plt.show()
    plt.close()