# -*- coding: utf-8 -*-
"""
Various plots

"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, FFMpegWriter
import xarray as xr
import os

def quiver(data, arrScale = 25.0, threshold = None, nthArr = 1, 
              contourLevels = None, colbar = True, logscale = False,
              aspectratio='equal', colbar_orient = 'vertical', units = None):
    """
    Generates a quiver plot of a 'data' xarray DataArray object (single frame from a dataset)
    Inputs:
        data - xarray DataArray of the type defined in pivpy, one of the frames in the Dataset
            selected by default using .isel(t=0)
        threshold - values above the threshold will be set equal to threshold
        arrScale - use to change arrow scales
        nthArr - use to plot only every nth arrow from the array 
        contourLevels - use to specify the maximum value (abs) of contour plots 
        colbar - True/False wether to generate a colorbar or not
        logscale - if true then colorbar is on log scale
        aspectratio - set auto or equal for the plot's apearence
        colbar_orient - 'horizontal' or 'vertical' orientation of the colorbar (if colbar is True)
    Outputs:
        none
    Usage:
        graphics.quiver(data, arrScale = 0.2, threshold = Inf, n)
    """
    
    data = dataset_to_array(data)
        
    x = data.x
    y = data.y
    u = data.u
    v = data.v
    
    if units is not None:
        lUnits = units[0] # ['m' 'm' 'mm/s' 'mm/s']
        velUnits = units[2]
        tUnits = velUnits.split('/')[1] # make it 's' or 'dt'
    else:
        lUnits, velUnits, tUnits = '', '', ''
    
    
    if threshold is not None:
        data['u'] = xr.where(data['u']>threshold, threshold, data['u'])
        data['v'] = xr.where(data['v']>threshold, threshold, data['v'])
        
    S = np.array(np.sqrt(u**2 + v**2))
    
    fig = plt.get_fignums()
    if len(fig) == 0: # if no figure is open
        fig, ax = plt.subplots() # open a new figure
    else:
        ax = plt.gca()     
    
    if contourLevels is None:
        levels = np.linspace(0, np.max(S.flatten()), 30) # default contour levels up to max of S
    else:
        levels = np.linspace(0, contourLevels, 30)
                
    if logscale:
        c = ax.contourf(x,y,S,alpha=0.8,
                 cmap = plt.get_cmap("Blues"), 
                 levels = levels, norm = plt.colors.LogNorm())
    else:
        c = ax.contourf(x,y,S,alpha=0.8,
                 cmap = plt.get_cmap("Blues"), 
                 levels=levels)
    if colbar:
        cbar = plt.colorbar(c, orientation=colbar_orient)
        cbar.set_label(r'$\left| \, V \, \right|$ ['+ lUnits +' $\cdot$ '+ tUnits +'$^{-1}$]')
        
    ax.quiver(x[::nthArr],y[::nthArr],
               u[::nthArr,::nthArr],v[::nthArr,::nthArr],units='width',
               scale = np.max(S*arrScale),headwidth=2)
    ax.set_xlabel('x (' + lUnits + ')')
    ax.set_ylabel('y (' + lUnits + ')')
    ax.set_aspect(aspectratio)       
    return fig,ax

def histogram(data, normed = False):
    """
    this function will plot a normalized histogram of
    the velocity data.
    Input:
        data : xarray DataSet with ['u','v'] attrs['units']
        normed : (optional) default is False to present normalized
        histogram
        
    """
    
    u = np.asarray(data.u).flatten()
    v = np.asarray(data.v).flatten()
    
    units = data.attrs['units']
    f,ax = plt.subplots(2)
    
    ax[0].hist(u,bins=np.int(np.sqrt(len(u))*0.5),density=normed)
    ax[0].set_xlabel('u ['+units[2]+']')
    
    ax[1] = plt.subplot2grid((2,1),(1,0))
    ax[1].hist(v,bins=np.int(np.sqrt(len(v)*0.5)),density=normed)
    ax[1].set_xlabel('v ['+units[2]+']')
    plt.tight_layout()
    return f, ax


def contour_plot(data, threshold = None, contourLevels = None, 
                    colbar = True,  logscale = False, aspectration='equal', units=None):
    """ contourf ajusted for the xarray PIV dataset, creates a 
        contour map for the data['w'] property. 
        Input:
            data : xarray PIV DataArray, converted automatically using .isel(t=0)
            threshold : a threshold value, default is None (no data clipping)
            contourLevels : number of contour levels, default is None
            colbar : boolean (default is True) show/hide colorbar 
            logscale : boolean (True is default) create in linear/log scale
            aspectration : string, 'equal' is the default
        
    """

    data = dataset_to_array(data)

    
    if units is not None:
        lUnits = units[0] # ['m' 'm' 'mm/s' 'mm/s']
        # velUnits = units[2]
        # tUnits = velUnits.split('/')[1] # make it 's' or 'dt'
    else:
        # lUnits, velUnits = '', ''
        lUnits = ''
        
    f,ax = plt.subplots()    
    
    if threshold is not None:
        data['w'] = xr.where(data['w']>threshold, threshold, data['w'])
        
    m = np.amax(abs(data['w']))
    if contourLevels == None:
        levels = np.linspace(-m, m, 30)
    else:
        levels = np.linspace(-contourLevels, contourLevels, 30)
        
    if logscale:
        c = ax.contourf(data.x,data.y,np.abs(data['w']), levels=levels,
                 cmap = plt.get_cmap('RdYlBu'), norm=plt.colors.LogNorm())
    else:
        c = ax.contourf(data.x,data.y,data['w'], levels=levels,
                 cmap = plt.get_cmap('RdYlBu'))
        
    plt.xlabel('x [' + lUnits + ']')
    plt.ylabel('y [' + lUnits + ']')
    if colbar:
        cbar = plt.colorbar(c)
        cbar.set_label(r'$\omega$ [s$^{-1}$]')
    ax.set_aspect(aspectration)
    return f,ax
        
         
def showf(data, variables=None, units=None, fig=None):
    """ 
    showf(data, var, units)
    Arguments:
        data : xarray.DataSet that contains dimensions of t,x,y
               and variables u,v and maybe w (scalar)
    """

    if variables is None:
        xlabel = ' '
        ylabel = ' '
    else:
        xlabel = variables[0]
        ylabel = variables[1]

    if units is not None:
        xlabel += ' ' + units[0]
        ylabel += ' ' + units[1]


    fig = plt.figure(None if fig is None else fig.number)


    for t in data['t']:
        d = data.isel(t=t)
        plt.quiver(d['x'],d['y'],d['u'],d['v'],d['u']**2 + d['v']**2)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.draw()
        plt.pause(0.1)
        
    plt.show()

def showscal(data, property='ken'):
    """ 
    showf(data, var, units)
    Arguments:
        data : xarray.DataSet that contains dimensions of t,x,y
               and a variable w (scalar)
    """
    # fig = plt.figure(None if fig is None else fig.number)
    # import pdb; pdb.set_trace()
    # xlabel = (None if var is None else var[0]) + ' [' + (None if units is None else units[0])+']'
    # ylabel = (None if var is None else var[1]) + ' [' + (None if units is None else units[1])+']'
    
    
    data = data.piv.vec2scal(property=property)
    contour_plot(data)                
        

def animate(data, arrowscale=1, savepath=None):
    """ animates the quiver plot for the dataset (multiple frames) 
    Input:
        data : xarray PIV type of DataSet
        arrowscale : [optional] integer, default is 1
        savepath : [optional] path to save the MP4 animation, default is None
    
    Output:
        if savepath is None, then only an image display of the animation
        if savepath is an existing path, a file named im.mp4 is saved
    
    """    
    X, Y = data.x, data.y
    U, V = data.u[:,:,0], data.v[:,:,0] # first frame
    fig, ax = plt.subplots(1,1)
    M = np.sqrt(U**2 + V**2)
    
    Q = ax.quiver(X[::3,::3], Y[::3,::3], 
                  U[::3,::3], V[::3,::3], M[::3,::3],
                 units='inches', scale=arrowscale)
    
    cb = plt.colorbar(Q)
    
    units = data.attrs['units']
    
    cb.ax.set_ylabel('velocity (' + units[2] + ')')
    
    text = ax.text(0.2,1.05, '1/'+str(len(data.t)), ha='center', va='center',
                   transform=ax.transAxes)
    
    def update_quiver(num,Q,data,text):
        U,V = data.u[:,:,num],data.v[:,:,num]
        
        M = np.sqrt(U[::3,::3]**2 + V[::3,::3]**2)   
        Q.set_UVC(U,V,M)
        text.set_text(str(num+1)+'/'+str(len(data.t)))
        return Q

    anim = FuncAnimation(fig, update_quiver, fargs=(Q,data,text),
                               frames = len(data.t), blit=False)
    mywriter = FFMpegWriter()
    if savepath:
        p = os.getcwd()
        os.chdir(savepath)
        anim.save('im.mp4', writer=mywriter)
        os.chdir(p)
    else: anim.save('im.mp4', writer=mywriter)  
    
    

def dataset_to_array(data,N=0):
    """ converts xarray Dataset to array """
    if 't' in data.dims:
        print('Warning: function for a single frame, using first frame, supply data.isel(t=N)')
        data = data.isel(t=N)
    return data
[((56, 10, 56, 27), 'matplotlib.pyplot.get_fignums', 'plt.get_fignums', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((102, 11, 102, 26), 'matplotlib.pyplot.subplots', 'plt.subplots', ({(102, 24, 102, 25): '2'}, {}), '(2)', True, 'import matplotlib.pyplot as plt\n'), ((107, 12, 107, 41), 'matplotlib.pyplot.subplot2grid', 'plt.subplot2grid', ({(107, 29, 107, 34): '(2, 1)', (107, 35, 107, 40): '(1, 0)'}, {}), '((2, 1), (1, 0))', True, 'import matplotlib.pyplot as plt\n'), ((110, 4, 110, 22), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((139, 11, 139, 25), 'matplotlib.pyplot.subplots', 'plt.subplots', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((157, 4, 157, 36), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(157, 15, 157, 35): "('x [' + lUnits + ']')"}, {}), "('x [' + lUnits + ']')", True, 'import matplotlib.pyplot as plt\n'), ((158, 4, 158, 36), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(158, 15, 158, 35): "('y [' + lUnits + ']')"}, {}), "('y [' + lUnits + ']')", True, 'import matplotlib.pyplot as plt\n'), ((186, 10, 186, 57), 'matplotlib.pyplot.figure', 'plt.figure', ({(186, 21, 186, 56): 'None if fig is None else fig.number'}, {}), '(None if fig is None else fig.number)', True, 'import matplotlib.pyplot as plt\n'), ((197, 4, 197, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((230, 14, 230, 31), 'matplotlib.pyplot.subplots', 'plt.subplots', ({(230, 27, 230, 28): '1', (230, 29, 230, 30): '1'}, {}), '(1, 1)', True, 'import matplotlib.pyplot as plt\n'), ((231, 8, 231, 28), 'numpy.sqrt', 'np.sqrt', ({(231, 16, 231, 27): 'U ** 2 + V ** 2'}, {}), '(U ** 2 + V ** 2)', True, 'import numpy as np\n'), ((237, 9, 237, 24), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({(237, 22, 237, 23): 'Q'}, {}), '(Q)', True, 'import matplotlib.pyplot as plt\n'), ((256, 15, 256, 29), 'matplotlib.animation.FFMpegWriter', 'FFMpegWriter', ({}, {}), '()', False, 'from matplotlib.animation import FuncAnimation, FFMpegWriter\n'), ((51, 20, 51, 71), 'xarray.where', 'xr.where', ({(51, 29, 51, 48): "data['u'] > threshold", (51, 50, 51, 59): 'threshold', (51, 61, 51, 70): "data['u']"}, {}), "(data['u'] > threshold, threshold, data['u'])", True, 'import xarray as xr\n'), ((52, 20, 52, 71), 'xarray.where', 'xr.where', ({(52, 29, 52, 48): "data['v'] > threshold", (52, 50, 52, 59): 'threshold', (52, 61, 52, 70): "data['v']"}, {}), "(data['v'] > threshold, threshold, data['v'])", True, 'import xarray as xr\n'), ((54, 17, 54, 37), 'numpy.sqrt', 'np.sqrt', ({(54, 25, 54, 36): 'u ** 2 + v ** 2'}, {}), '(u ** 2 + v ** 2)', True, 'import numpy as np\n'), ((58, 18, 58, 32), 'matplotlib.pyplot.subplots', 'plt.subplots', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((60, 13, 60, 22), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((65, 17, 65, 50), 'numpy.linspace', 'np.linspace', ({(65, 29, 65, 30): '0', (65, 32, 65, 45): 'contourLevels', (65, 47, 65, 49): '30'}, {}), '(0, contourLevels, 30)', True, 'import numpy as np\n'), ((76, 15, 76, 57), 'matplotlib.pyplot.colorbar', 'plt.colorbar', (), '', True, 'import matplotlib.pyplot as plt\n'), ((142, 20, 142, 71), 'xarray.where', 'xr.where', ({(142, 29, 142, 48): "data['w'] > threshold", (142, 50, 142, 59): 'threshold', (142, 61, 142, 70): "data['w']"}, {}), "(data['w'] > threshold, threshold, data['w'])", True, 'import xarray as xr\n'), ((146, 17, 146, 39), 'numpy.linspace', 'np.linspace', ({(146, 29, 146, 31): '-m', (146, 33, 146, 34): 'm', (146, 36, 146, 38): '30'}, {}), '(-m, m, 30)', True, 'import numpy as np\n'), ((148, 17, 148, 63), 'numpy.linspace', 'np.linspace', ({(148, 29, 148, 43): '-contourLevels', (148, 45, 148, 58): 'contourLevels', (148, 60, 148, 62): '30'}, {}), '(-contourLevels, contourLevels, 30)', True, 'import numpy as np\n'), ((160, 15, 160, 30), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({(160, 28, 160, 29): 'c'}, {}), '(c)', True, 'import matplotlib.pyplot as plt\n'), ((191, 8, 191, 69), 'matplotlib.pyplot.quiver', 'plt.quiver', ({(191, 19, 191, 25): "d['x']", (191, 26, 191, 32): "d['y']", (191, 33, 191, 39): "d['u']", (191, 40, 191, 46): "d['v']", (191, 47, 191, 68): "(d['u'] ** 2 + d['v'] ** 2)"}, {}), "(d['x'], d['y'], d['u'], d['v'], d['u'] ** 2 + d['v'] ** 2)", True, 'import matplotlib.pyplot as plt\n'), ((192, 8, 192, 26), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(192, 19, 192, 25): 'xlabel'}, {}), '(xlabel)', True, 'import matplotlib.pyplot as plt\n'), ((193, 8, 193, 26), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(193, 19, 193, 25): 'ylabel'}, {}), '(ylabel)', True, 'import matplotlib.pyplot as plt\n'), ((194, 8, 194, 18), 'matplotlib.pyplot.draw', 'plt.draw', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((195, 8, 195, 22), 'matplotlib.pyplot.pause', 'plt.pause', ({(195, 18, 195, 21): '(0.1)'}, {}), '(0.1)', True, 'import matplotlib.pyplot as plt\n'), ((249, 12, 249, 50), 'numpy.sqrt', 'np.sqrt', ({(249, 20, 249, 49): 'U[::3, ::3] ** 2 + V[::3, ::3] ** 2'}, {}), '(U[::3, ::3] ** 2 + V[::3, ::3] ** 2)', True, 'import numpy as np\n'), ((258, 12, 258, 23), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((259, 8, 259, 26), 'os.chdir', 'os.chdir', ({(259, 17, 259, 25): 'savepath'}, {}), '(savepath)', False, 'import os\n'), ((261, 8, 261, 19), 'os.chdir', 'os.chdir', ({(261, 17, 261, 18): 'p'}, {}), '(p)', False, 'import os\n'), ((81, 23, 81, 41), 'numpy.max', 'np.max', ({(81, 30, 81, 40): '(S * arrScale)'}, {}), '(S * arrScale)', True, 'import numpy as np\n'), ((98, 8, 98, 26), 'numpy.asarray', 'np.asarray', ({(98, 19, 98, 25): 'data.u'}, {}), '(data.u)', True, 'import numpy as np\n'), ((99, 8, 99, 26), 'numpy.asarray', 'np.asarray', ({(99, 19, 99, 25): 'data.v'}, {}), '(data.v)', True, 'import numpy as np\n'), ((151, 38, 151, 55), 'numpy.abs', 'np.abs', ({(151, 45, 151, 54): "data['w']"}, {}), "(data['w'])", True, 'import numpy as np\n'), ((69, 24, 69, 45), 'matplotlib.pyplot.get_cmap', 'plt.get_cmap', ({(69, 37, 69, 44): '"""Blues"""'}, {}), "('Blues')", True, 'import matplotlib.pyplot as plt\n'), ((70, 41, 70, 61), 'matplotlib.pyplot.colors.LogNorm', 'plt.colors.LogNorm', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((73, 24, 73, 45), 'matplotlib.pyplot.get_cmap', 'plt.get_cmap', ({(73, 37, 73, 44): '"""Blues"""'}, {}), "('Blues')", True, 'import matplotlib.pyplot as plt\n'), ((152, 24, 152, 46), 'matplotlib.pyplot.get_cmap', 'plt.get_cmap', ({(152, 37, 152, 45): '"""RdYlBu"""'}, {}), "('RdYlBu')", True, 'import matplotlib.pyplot as plt\n'), ((152, 53, 152, 73), 'matplotlib.pyplot.colors.LogNorm', 'plt.colors.LogNorm', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((155, 24, 155, 46), 'matplotlib.pyplot.get_cmap', 'plt.get_cmap', ({(155, 37, 155, 45): '"""RdYlBu"""'}, {}), "('RdYlBu')", True, 'import matplotlib.pyplot as plt\n')]