#!/usr/bin/env python

'''
File name: visualize_dataset.py
Author: Ishaan Sood, KUL (ishaan.sood@kuleuven.be)
Date: 24.02.2020

Description:
Visualization functions of Total Control Turbine Box Datasets
'''

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

def plot_slices(Vel, i, j, k, metadata, var, nt):

    if var == 'u':
        ttl = 'Streamwise velocity at t=' + '{:.2f}'.format(metadata['t'][nt]) + 's '
    elif var == 'v':
        ttl = 'Spanwise velocity at t=' + '{:.2f}'.format(metadata['t'][nt]) + 's '
    elif var == 'w':
        ttl = 'Vertical velocity at t=' + '{:.2f}'.format(metadata['t'][nt]) + 's '

    # i = int(np.floor(metadata['Nx']/2))
    # j = int(np.floor(metadata['Ny']/2))
    # k = int(np.floor(metadata['Nz']/2))

    um = [np.floor(np.min(Vel)), np.ceil(np.max(Vel))]

    figx, ax = plt.subplots(figsize=(6.,4.))
    figx.subplots_adjust(bottom=0.15, top=0.9, left=0.15, right=0.95)
    pcm, ax = plot_plane(ax, Vel[i, :, :], metadata, 'x', i, ttl+'(Front view x='+ '{:.2f}'.format(metadata['x'][i]) +'m)', um)
    axcb, _ = mpl.colorbar.make_axes(ax, shrink=1)
    cb = plt.colorbar(pcm, axcb)
    cb.set_label(var+' [m/s]', fontsize=12)
    cb.ax.tick_params(labelsize=12)
    plt.savefig('front_view.png')
    plt.show()


    figy, ax = plt.subplots(figsize=(6.,4.))
    figy.subplots_adjust(bottom=0.15, top=0.9, left=0.15, right=0.95)
    pcm, ax = plot_plane(ax, Vel[:, j, :], metadata, 'y', j, ttl+'(Side view y='+ '{:.2f}'.format(metadata['y'][j]) +'m)', um)
    axcb, _ = mpl.colorbar.make_axes(ax, shrink=1)
    cb = plt.colorbar(pcm, axcb)
    cb.set_label(var+' [m/s]', fontsize=12)
    cb.ax.tick_params(labelsize=12)
    plt.savefig('side_view.png')
    plt.show()

    figz, ax = plt.subplots(figsize=(6.,4.))
    figz.subplots_adjust(bottom=0.15, top=0.9, left=0.15, right=0.95)
    pcm, ax = plot_plane(ax, Vel[:, :, k], metadata, 'z', k, ttl+'(Top view z='+ '{:.2f}'.format(metadata['z'][k]) +'m)', um)
    axcb, _ = mpl.colorbar.make_axes(ax, shrink=1)
    cb = plt.colorbar(pcm, axcb)
    cb.set_label(var+' [m/s]', fontsize=12)
    cb.ax.tick_params(labelsize=12)
    plt.savefig('top_view.png')
    plt.show()

    return [figx,figy,figz]

def plot_tplanes(V, metadata, var, ind, q, i):

    if var == 'u':
        vttl = 'Streamwise velocity at '+q+'=' + '{:.2f}'.format(metadata[q][i]) + 'm and '
    elif var == 'v':
        vttl = 'Spanwise velocity at '+q+'=' + '{:.2f}'.format(metadata[q][i]) + 'm and '
    elif var == 'w':
        vttl = 'Vertical velocity at '+q+'=' + '{:.2f}'.format(metadata[q][i]) + 'm and '

    lfig = []
    um = [0, 15]
    for k in range(len(V)):
        ttl = vttl+'t='+'{:.1f}'.format(metadata['t'][ind[k]]) +'s.'
        fig, ax = plt.subplots(figsize=(10.,4.))#, ncols=len(V), sharey=True)
        fig.subplots_adjust(bottom=0.15, top=0.9, left=0.1, right=0.95)
        pcm, ax = plot_plane(ax, V[k], metadata, q, i, ttl, um)
        axcb, _ = mpl.colorbar.make_axes(ax, shrink=1)
        cb = plt.colorbar(pcm, axcb)
        cb.set_label(var + ' [m/s]', fontsize=12)
        cb.ax.tick_params(labelsize=12)
        lfig.append(fig)
        plt.show()


    return lfig

def plot_plane(ax, Vel, metadata, q, i, ttl, um):
    X, Y, Z = np.meshgrid(metadata['x'], metadata['y'], metadata['z'])
    if q == 'x':
        x = np.squeeze(Y[:, i, :])
        y = np.squeeze(Z[:, i, :])
        plt.gca().set(xlabel='y [m]', ylabel='z [m]',
                      xlim=[min(metadata['y']),max(metadata['y'])],
                      ylim=[min(metadata['z']),max(metadata['z'])],
                      xticks=np.linspace(min(metadata['y']), max(metadata['y']),5),
                      yticks=np.linspace(min(metadata['z']), max(metadata['z']),5))
    elif q == 'y':
        x = np.squeeze(X[i, :, :])
        y = np.squeeze(Z[i, :, :])
        plt.gca().set(xlabel='x [m]', ylabel='z [m]',
                      xlim=[min(metadata['x']),max(metadata['x'])],
                      ylim=[min(metadata['z']),max(metadata['z'])],
                      xticks=np.linspace(min(metadata['x']), max(metadata['x']),5),
                      yticks=np.linspace(min(metadata['z']), max(metadata['z']),5))
    elif q == 'z':
        x = np.transpose(np.squeeze(X[:, :, i]),(1,0))
        y = np.transpose(np.squeeze(Y[:, :, i]),(1,0))
        ax.set(xlabel='x [m]', ylabel='y [m]',
                      xlim=[min(metadata['x']),max(metadata['x'])],
                      ylim=[min(metadata['y']),max(metadata['y'])],
                      xticks=np.linspace(min(metadata['x']), max(metadata['x']),5),
                      yticks=np.linspace(min(metadata['y']), max(metadata['y']),5))
    pcm = ax.pcolormesh(x, y, Vel, vmin=um[0], vmax=um[1], shading='gouraud')
    ax.set_title(ttl)
    return pcm, ax

def plot_tserie(V,metadata,ind):

    # Color scale
    colors = plt.cm.Blues_r(np.linspace(0.2,0.6,len(V)-1))

    # Create figure
    fig, ax = plt.subplots(figsize=(10.,4.))
    fig.subplots_adjust(bottom=0.15, top=0.9, left=0.1, right=0.95)

    # Add data lines
    for m,k in zip(['u','v','w'],range(3)):
        ax.plot(V['t'],V[m],color=colors[k])

    # Axis layout
    ax.set_xlim([metadata['t'][0],metadata['t'][-1]])
    ax.set_xticks(np.linspace(metadata['t'][0], metadata['t'][-1], 4))
    ax.set_xlabel(r"$t \, [s]$", fontsize=14)
    ax.set_ylim([-2.0, 14.0])
    ax.set_yticks(np.linspace(-2.0, 14, 9))
    ax.set_ylabel(r"$V_i \, [m/s]$", fontsize=14)
    ax.tick_params(axis='both', which='major', labelsize=12)
    ax.grid()
    # Title
    x = metadata['x'][ind[0]]
    y = metadata['y'][ind[1]]
    z = metadata['z'][ind[2]]
    title_str = r'Velocity time series ('+str(metadata['Nt'])+r' samples) at $(x,y,z) \approx ['+'{:3.0f}'.format(x)+'m,'+'{:3.0f}'.format(y)+'m,'+'{:3.0f}'.format(z)+'m]$'
    ax.set_title(title_str,fontsize=14)
    # Legend
    ax.legend([r'$u$',r'$v$',r'$w$'], loc=1, fontsize=12)
    plt.savefig('time_series.png')
    plt.show()

    return fig