#########################################################################################################################################################################
# Author: Raphael Maurin, raphael.maurin@imft.fr
# 24/11/2017
#
# Post processing script, to extract and plot the results of sedimentTransportExample scripts, in terms of dimensionless sediment transport,
# and time-averaged depth profiles of streamwise fluid and solid velocity, solid volume fraction and sediment transport rate density.
# For example of applications, see the articles Maurin et al (2015,2016), and the PhD of Maurin (2015), available in the references of YADE.
#
############################################################################################################################################################################

import numpy as np
import os
from matplotlib.pyplot import *
import matplotlib.gridspec as gridspec

#######
# INPUT
timeStepSave = 0.1  # Data saved every timeStepSave seconds in the simulation
startFile = 500  # File where the post-processing begins
endFile = 10000  # File where the post-processing ends

scriptPath = os.path.abspath(os.path.dirname(sys.argv[-1]))  #Path where the script is stored
if os.path.exists(scriptPath + '/data/') == False:  #If the data folder does not exist, no data to extract, exit.
	print('\n There is no data to extract in this folder ! Please first run sedimentTransportExample_1DRANSCoupling.py !\n')
	exit()
else:  #Else, extract the first file in order to get the size of the vectors, ndimz
	exec(open(scriptPath + '/data/0.py').read())
	ndimz = len(phiPartPY)  #Size of the vectors, mesh parameter.

#Initilization of the variable to extract and plot
[qs, time, phiPart, vxPart, vxFluid] = [[], [], np.zeros(ndimz), np.zeros(ndimz), np.zeros(ndimz + 1)]

########
# LOOP over time to EXTRACT the data
########
for i in range(startFile, endFile):
	nameFile = scriptPath + '/data/' + str(i) + '.py'  #Name of the file at the considered time step
	if os.path.exists(nameFile) == False:  #Check if the file exist
		print('\nThe file {0} does not exist, stop the post processing at this stage !\n'.format(nameFile))
		endFile = i - 1  #Last file processed
		break
	#Extract the data from at the time (file) considered. Assign the vectors to qsMean, phiPartPY, vxPartPY, vxFluidPY, and zAxis see the structure of a saved file X.py.
	exec(open(nameFile).read())
	time += [i * timeStepSave]  #Recreate the time scale
	qs += [qsMean]  #Extract the sediment transport rate as a function of time
	phiPart += phiPartPY  #Average over time the solid volume fraction vertical profile
	vxPart += vxPartPY * phiPartPY  #Average over time the solid velocity profile. The spatial averaging is weighted by the particle volume so that it is necessary to re-multiply the averaging by the solid volume fraction before re-normalizing. See the paper Maurin et al (2015) in the references for more details on the averaging.
	vxFluid += vxFluidPY  #Average over time the fluid velocity vertical profile
#Normalization
toto = np.where(phiPart > 0)
vxPart[toto] /= phiPart[toto]  #Avoid division by zero
phiPart /= (endFile - startFile + 1)
vxFluid /= (endFile - startFile + 1)

######
# PLOT

# sediment transport rate as a function of time
fig1 = figure()  #Create a figure
gs = gridspec.GridSpec(1, 1)  #Split the figure in 1 row and 1 column
ax = subplot(gs[0, 0])  #Assign the axis
ax.plot(time, qs, 'ok')  #Plot the sediment transport rate as a function of time
ax.grid()  #Put a grid on the figure
ax.set_xlabel('t (s)')  #Assign the x axis label
ax.set_ylabel(r'$Q_s^*$', rotation='horizontal')  #Assign the y axis label
savefig('transportRate.png')  #Save the figure

# Solid and fluid velocity depth profiles, solid volume fraction and sediment transport rate density depth profiles
fig2 = figure(figsize=(12, 4.5))
gs = gridspec.GridSpec(1, 3)  #Split the figure in 1 row and 3 column
ax1 = subplot(gs[0, 0])  #ax1 = subfigure corresponding to 1st row 1st column
ax2 = subplot(gs[0, 1])  #ax2 = subfigure corresponding to 1st row 2nd column
ax3 = subplot(gs[0, 2])  #ax3 = subfigure corresponding to 1st row 3rd column
ax1.plot(vxPart, zAxis, '-r')  #Plot the solid velocity profile on first subfigure
#ax1.plot(vxFluid[:ndimz],zAxis,'-b')	#Plot the fluid velocity profile on first subfigure
ax2.plot(phiPart, zAxis, '-r', label='NEW')  #Plot the solid volume fraction on second subfigure
ax3.plot(phiPart * vxPart, zAxis, '-r')  #Plot the sediment transport rate density on the third subfigure
ax1.set_xlabel(r'$<v^k_x>$ (m/s)')  #Assign the x axis label of the first subfigure
ax1.set_ylabel(r'$z/d$', rotation='horizontal')  #Assign the y axis label of the first subfigure
ax2.set_xlabel(r'$\phi$')  #Assign the x axis label of the second subfigure
ax3.set_xlabel(r'$q (m/s)$')  #Assign the x axis label of the third subfigure
ax2.set_yticklabels([])  #Remove the vertical scale for the second subfigure, same as first
ax3.set_yticklabels([])  #Remove the vertical scale for the third subfigure, same as first
ax1.grid()  #Put a grid on each subfigure
ax2.grid()
ax3.grid()
fig2.subplots_adjust(
        left=0.04, bottom=0.11, right=0.98, top=0.98, wspace=.1
)  #Adjust the spacing around the full figure (left, bottom, right, top) and between the subfigure (wspace), in percent of the figure size
exec(open('DATAr2d6s2_Maurinetal2016.py').read())
ax1.plot(vxPart_r2d6s2, sig_r2d6s2, 'k')
ax2.plot(phiPart_r2d6s2, sig_r2d6s2, 'k', label='Maurin et al (2016)')
ax3.plot(phiPart_r2d6s2 * vxPart_r2d6s2, sig_r2d6s2, 'k')

ax1.set_ylim([6, 15])
ax2.set_ylim([6, 15])
ax3.set_ylim([6, 15])
ax2.legend()

savefig('depthProfiles.png')  #Save the figure

show()  #Show the two figures created
