#################################################################################
#
# Tutorial file how to read in hdf5 data from SP-Wind precursor simulations, 
# generated in the context of the TotalControl project.
# 
# A thorough description of the simulations and data can be found at in the deliverable report
#
# author: Wim Munters
# date: 08 Feb 2019
# email: wim.munters@gmail.com
#
#################################################################################
import numpy as np
import h5py 
import matplotlib.pyplot as plt

plt.close('all')

# open the files
fsnap = h5py.File('precursor_snapshot.h5', 'r')
ftavg = h5py.File('precursor_timeavg.h5', 'r')
ftser = h5py.File('precursor_timeseries.h5', 'r')

###########################################################################################################################
# 
# FILE: precursor_timeseries.h5
#
#   group: metmast_timeseries
#       description:    measurements of virtual metmast at 2 Hz resolution (at x = 16 km, y = 0 km)
#       dataset: z (elevations of metmast measurements)
#       subgroups: u, v, w, (theta)
#           datasets:  t_0000#.5 (e.g. at t = 0.5 sec)
#
#   group: inflowplane_timese#ries
#       description:    cross# sectional plane of y x z = 16 km x 1.5 km (at x = 16 km)
#       dataset: y (spanwise grid location)
#       dataset: z (vertical grid location)
#       subgroups: u, v, w, (theta)
#           datasets:  t_0000.5 (e.g. at t = 0.5 sec)
#
#   group: 3Dbox_timeseries
#       description:    time series of 3D box located at [15 900 to 16 000] x [0 to 100] x [0 500]
#       dataset: x (streamwise grid)
#       dataset: y (spanwise grid)
#       dataset: z (vertical grid)
#       subgroup: u, v, w, theta
#           dataset: t_0002.0 (e.g. at t = 2 sec)
#
###########################################################################################################################

#-----------------------------------------------------------------------
# METMAST DATA: example
#-----------------------------------------------------------------------
metmast_group = ftser["metmast_timeseries"]
groupname = "u"
T  = 4500
Nz = 225
dt = 0.5
Nt = np.int(T/dt)

time   = np.linspace(0, T, Nt)
height = metmast_group["z"]

u_metmast_array = np.empty((Nz, Nt))

# load all the data in memory
for ti in range(Nt):
    t = (ti+1)*dt

    setname = "t_{:05.1f}".format(t)
    try:
        u_metmast_array[:,ti] = metmast_group[groupname][setname]
    except KeyError:
        print('Running out of bounds on metmast data at t = {:05.1f}'.format(t))
        exit

# plot full metmast data
plt.figure()
plt.pcolormesh(time, height, u_metmast_array)
plt.xlabel('Time [s]') 
plt.ylabel('Height [m]')
plt.title('Metmast streamwise velocity [m/s]')
plt.colorbar()

# plot timeseries of point measurement
heights= [10, 100, 150, 200] #m


plt.figure()
for h in heights:
    index = np.searchsorted(height, h)
    plt.plot(time, u_metmast_array[index,:], label="z = {:d} m".format(h))
plt.xlabel('Time [s]')
plt.ylabel('u [m/s]')
plt.title('Metmast point measurements')
plt.legend()

#-----------------------------------------------------------------------
# INFLOW PLANE:example
#-----------------------------------------------------------------------
inflow_group = ftser["inflowplane_timeseries"]

y, z = inflow_group["y"], inflow_group["z"]
component = "u"

z = np.linspace(0, 1.5e3, 225)

t = 300

setname = "t_{:04.1f}".format(t)

data = inflow_group[component][setname]
plt.figure()
plt.imshow(np.flipud(data), extent=(0,np.max(y),0,np.max(z)))


#-----------------------------------------------------------------------
# 3D BOX:example
#-----------------------------------------------------------------------
box_group = ftser["3Dbox_timeseries"]
component = "u"
x,y,z = box_group["x"], box_group["y"], box_group["z"]

t = 300

plt.figure()
dataset = box_group[component]["t_{:05.1f}".format(t)]
plt.subplot(131)
plt.title('topview')
plt.imshow(np.flipud(np.transpose(dataset[0,:,:])), extent=(0, np.max(x), 0, np.max(y)))
plt.xlabel('x [m]')
plt.ylabel('y [m]')

plt.subplot(132)
plt.title('sideview')
plt.imshow(np.flipud(dataset[:,0,:]), extent=(0, np.max(x), 0, np.max(z)))
plt.xlabel('x [m]')
plt.ylabel('z [m]')

plt.subplot(133)
plt.title('frontview')
plt.imshow(np.flipud(dataset[:,:,0]), extent=(0, np.max(y), 0, np.max(z)))
plt.xlabel('y [m]')
plt.ylabel('z [m]')

plt.tight_layout()



###########################################################################################################################
#
# FILE: precursor_timeavg.h5
#
#   group: kx/ky_spectra_timeavg
#       description:    spectra of velocity fluctuations in the streamwise/spanwise direction as a function of height. 
#                       averaged over time and spanwise/streamwise direction.
#       dataset: kx/ky  (wavenumbers)
#       dataset: z      (vertical locations of spectral measurements)
#       dataset: uu     (streamwise velocity)
#       dataset: vv     (spanwise velocity)
#       dataset: ww     (vertical velocity)
#
#   group: profiles_timeavg
#       description:    profiles of flow variables. averaged over horizontal directions and time
#       dataset: z      (vertical locations of measurements)
#       dataset: U/V/W  (time-averaged velocity in stream, spanwise, and vertical directions)
#       dataset: theta  (only for CN cases, should be multiplied with theta_ref
#       dataset: uu/uv/uw/vv/vw/ww (Reynolds stresses)
#
###########################################################################################################################

#-----------------------------------------------------------------------
# FLOW PROFILES: example
#-----------------------------------------------------------------------
profiles = ftavg["profiles_timeavg"] 
plt.figure()
plt.subplot(131)
plt.plot(profiles["U"], profiles["z"], label='U')  # Streamwise velocity
plt.plot(profiles["V"], profiles["z"], label='V')  # Spanwise velocity
plt.legend()
plt.xlabel(r'Velocity [m/s]')
plt.ylabel(r'z [m]'); plt.ylim((0, 1500))

#plt.subplot(132)
#theta_ref = 273 + 15
#plt.plot(np.array(profiles["theta"])*theta_ref, profiles["z"])  # Temperature
#plt.xlabel(r'V [m/s]')
#plt.ylim((0, 1500))

plt.subplot(133)
plt.plot(-np.array(profiles["uw"]), profiles["z"]) # Reynolds shear stress
plt.xlabel(r"-u'w' [m2/s2]")
plt.ylim((0, 1500))

plt.tight_layout()

#-----------------------------------------------------------------------
# SPECTRA: example
#-----------------------------------------------------------------------
spectra_groups = ["kx_spectra_timeavg", "ky_spectra_timeavg"]
plt.figure()
for igroup, group in enumerate(spectra_groups):
    plt.subplot(1,2,igroup+1)
    g = ftavg[group]

    z = g["z"]
    k = g["k"]
    spec_u = g["uu"]

    heights = [10, 100, 150, 200]
    for h in heights: 
        index = np.searchsorted(z, h)

        plt.loglog(k, spec_u[index,:k.size], label="z = {:d} m".format(h))

    plt.title(group)
    plt.xlabel('k [1/m]')
    plt.ylabel('E_uu [m2/s2]')

    if igroup==0:
        plt.legend()

plt.tight_layout()



#
# FILE: precursor_snapshot.h5
#
#   group: fullfield_snapshot
#       description:    snapshot of the full precursor field 
#       dataset: x (streamwise grid)
#       dataset: y (spanwise grid)
#       dataset: z (vertical grid)
#       dataset: u (axial velocity)
#       dataset: v (span velocity)
#       dataset: w (vert velocity)
#       dataset: theta (temperature)
#

#-----------------------------------------------------------------------
# FULLFIELD
#-----------------------------------------------------------------------
snapgroup = fsnap["fullfield_snapshot"]
x, y, z = snapgroup["x"], snapgroup["y"], snapgroup["z"]
u, v, w = snapgroup["u"], snapgroup["v"], snapgroup["w"]
# theta  = snapgroup["theta"]
plt.figure()

ip, jp = 200, 200

plt.subplot(231)
plt.title('u'); plt.ylabel('y [m]'); 
plt.imshow(np.transpose(u[:,:,17]), extent=(0, np.max(x), 0, np.max(y)))
plt.axvline(x=x[ip],linestyle='--', color='k', ymax=y[jp]/np.max(y))
plt.axhline(y=y[ip],linestyle='--', color='k', xmax=x[ip]/np.max(x))
plt.subplot(232)
plt.title('v'); 
plt.imshow(np.transpose(v[:,:,17]), extent=(0, np.max(x), 0, np.max(y)))
plt.axvline(x=x[ip],linestyle='--', color='k', ymax=y[jp]/np.max(y))
plt.axhline(y=y[ip],linestyle='--', color='k', xmax=x[ip]/np.max(x))
plt.subplot(233)
plt.title('w'); 
plt.imshow(np.transpose(w[:,:,17]), extent=(0, np.max(x), 0, np.max(y)))
plt.axvline(x=x[ip],linestyle='--', color='k', ymax=y[jp]/np.max(y))
plt.axhline(y=y[ip],linestyle='--', color='k', xmax=x[ip]/np.max(x))


plt.subplot(234)
plt.title('u (zoom)'); plt.ylabel('y [m]'); plt.xlabel('x [m]')
plt.imshow(np.transpose(u[:ip,:jp,17]), extent=(0, x[ip], 0, y[jp]))
plt.subplot(235)
plt.title('v (zoom)'); plt.xlabel('x [m]')
plt.imshow(np.transpose(v[:ip,:jp,17]), extent=(0, x[ip], 0, y[jp]))
plt.subplot(236)
plt.title('w (zoom)'); plt.xlabel('x [m]')
plt.imshow(np.transpose(w[:ip,:jp,17]), extent=(0, x[ip], 0, y[jp]))
plt.tight_layout()

