import numpy as np
import matplotlib.pyplot as plt
import h5py
from hdf5_tools import *

def readSensorArray(filename):
    data = np.loadtxt(filename, delimiter=',', skiprows=1)    
    # 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 = data[:,1:4]
    sensorResults = data[:,4:7]+1j*data[:,7:10]
    return (sensorCoordinates,sensorResults)

if __name__ == "__main__":
    # Displacement of the centerline
    frequencies = np.linspace(1,1000,10)
    quadraticError = np.zeros((10,),dtype=complex)
    for i in range(10):
        # Read the sensorArray for the reduced and full system
        filenameReduced = f"results_txt/centerLine_reduced-{i+1}"
        sensorDataReduced = readSensorArray(filenameReduced)
        filename = f"results_txt/centerLine-{i+1}"
        sensorData = readSensorArray(filename)

        # Plot y-displacement
        fig, ax = plt.subplots()
        ax.plot(1E3*sensorDataReduced[0][:,0],1E3*np.real(sensorDataReduced[1][:,1]),label="reduced Model")
        ax.plot(1E3*sensorDataReduced[0][:,0],1E3*np.real(sensorData[1][:,1]), label="full Model")
        ax.legend()
        ax.set_title(f"Displacement of the centerline at {frequencies[i]} Hz")
        ax.set_ylabel("y-displacement (real) in mm")
        ax.set_xlabel("x-coordindate in mm")
        plt.savefig(f"harmonic_results_{frequencies[i]}.png")
        plt.close()
    

    # Get the node Results from the Simulation and calculate the quadratic error
    try:
        h5f_r = h5py.File('results_hdf5/bendingBeam_reducedHarmonic.cfs','r')
        h5f = h5py.File('results_hdf5/bendingBeam_harmonic.cfs','r')
        for i in range(10):
            U_reduced = get_result(h5f_r,'mechDisplacement',region='V_beam',multistep=1,step=i+1) # all displacements
            U = get_result(h5f,'mechDisplacement',region='V_beam',multistep=1,step=i+1) # all displacements
            quadraticError[i] = np.mean(np.square(U-U_reduced))
        h5f_r.close()
        h5f.close()

        #Plot the quadratic error
        fig, ax = plt.subplots()
        ax.plot(frequencies,abs(quadraticError))
        ax.set_title(f"Quadratic error of all node displacements")
        ax.set_ylabel("quadratic error in m^2")
        ax.set_xlabel("Frequency in Hz")
        plt.savefig(f"quadratic_error.png")
        plt.close()
    except:
        # the h5 file should be closed if the data access does not work
        h5f_r.close()
        h5f.close()
        raise