############################################################################################
# Python script to calculate the wave-induced strain on the iceberg.
# NOTE: Before running this script, make sure that the following files are present
#           [YOURDIR]/dispVsFreq.py
#           [YOURDIR]/strainVsFreq.py
# These are generated by the getSolutionInterp.edp FreeFem script.
# See the bash script run-script.sh for more information.
#
# Run:
#          python3 strain2psd.py ICEBERG
#
# First compute the Pierson-Moskowitz spectrum (PSD) for approximating the ocean wave data
# Load the frequency domain data for the displacement and the strain.
# Compute the strain^2*PSD curve.
# Compute the area under the strain^2*PSD curve in a frequency bin to compute the RMS strain
#############################################################################################

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import sys
import time
from scipy.signal import find_peaks
font = {'weight' : 'normal',
        'size'   : 15}
matplotlib.rc('font', **font)
pi = np.pi

filePath = sys.argv[1]

dirs = [filePath+str(1), filePath+str(1.5)]
youngs = [1e9, 1.5e9]
g=9.8
omegap=0.092*2*pi
omega0=1.14*omegap
alpha=0.0081
beta=0.74

omega0s = [0.039, 0.047]
omega1s = np.add(omega0s, 0.003)

i = 0
plt.figure(figsize=[10,4])
for d in dirs:
    strainVsFreq = np.loadtxt(d+"/strainVsFreq.dat")

    # Pierson Moskowitz spectrum
    omeganew = 2*pi*np.linspace(0.01,0.125,len(strainVsFreq))
    S=alpha*g**2/omeganew**5*np.exp(-beta*(omega0/omeganew)**4) + 1e-7;

    strain2S = (strainVsFreq[:,0]**2)*S
    omega0 = 2*pi*omega0s[i]
    omega1 = 2*pi*omega1s[i]

    plt.semilogy(omeganew/(2*pi), strain2S, label="E = "+str(youngs[i]/1e9)+" GPa")
    plt.fill_between(omeganew/(2*pi), strain2S, 0, where=(omeganew > omega0) & (omeganew < omega1), color='red', alpha=0.5)
    plt.legend()
    plt.grid()

    omeganew1 = omeganew[(omeganew > omega0) & (omeganew < omega1)]
    strain2S1 = strain2S[(omeganew > omega0) & (omeganew < omega1)]
    area2 = np.trapz(strain2S1, omeganew1/(2*pi))
    area = np.sqrt(area2)
    print(area/1e-3)

    i += 1

plt.xlabel("$\\frac{\omega}{2\pi}$")
plt.ylabel("$\\varepsilon^2\,PSD$")
plt.show()
