# -*- coding: iso-8859-15 -*-
'''
Created on 14.09.2012

@author: Jakob
'''
from __future__ import print_function
#
#import matplotlib.cm as cm

def plot_matrix(m):
    import matplotlib.pyplot as plt
    plt.imshow(m, interpolation='none', cmap = plt.cm.Spectral) # @UndefinedVariable
    plt.show()
    
    
def subplot_matrix(m):
    import matplotlib.pyplot as plt    
    
    for i in range(m.shape[2]):
        plt.subplot(m.shape[2]*100 + 11 + i)
        plt.imshow(m[:,:,i], interpolation='none', cmap = plt.cm.Spectral) # @UndefinedVariable
        
    plt.show()
    

def plot_logfit_ensemble(e, cfg):
    import numpy as np
    import matplotlib.pyplot as plt
    
    # prepare data in lists
    
    depth = e.depth
    z_max = depth * cfg['logheight']
    z = []
    v = []
    
    for c in e.cells:
        z.append(depth + c.z_position)
        v.append(c.velocity.magn2d() if cfg['component'] == 3 else  c.velocity[cfg['component']] )
    

    # fitted velocity profile with more points to look nicer
    z_dense = np.linspace(0, depth)
    v_dense = (2.5 * np.log(30 * z_dense / e.ks)) * e.v_shear 
    
    # actual "raw" data
    plt.plot(v,z)
    plt.plot(v_dense, z_dense)
    plt.grid(True)
    plt.show()


def plot_logfit_ensembles(p, cfg, initial_ensemble=100):
    '''
    like plot_logfit_ensemble() but with interactiveness 
    based on: http://matplotlib.org/mpl_examples/widgets/slider_demo.py
    '''
    import matplotlib.pyplot as plt    
    import numpy as np
    from matplotlib.widgets import Slider
    from pylab import axes
    # prepare data in lists
    
    depth = p.ensembles[initial_ensemble].depth
    z_max = depth * cfg['logheight']
    z,v = [],[]
    
    for c in p.ensembles[initial_ensemble].cells:
        z.append(depth + c.z_position)
        v.append(c.velocity.magn2d()() if cfg['component'] == 3 else  c.velocity[cfg['component']] )
    

    # fitted velocity profile with more points to look nicer
    z_dense = np.linspace(0, depth)
    v_dense = (2.5 * np.log(30 * z_dense / p.ensembles[initial_ensemble].ks)) * p.ensembles[initial_ensemble].v_shear 
    
    # actual "raw" data
    p1, = plt.plot(v,z)
    p2, = plt.plot(v_dense, z_dense)
    plt.grid(True)
    
    # update function. 
    def update(en):
        en = int(en)
        if not p.ensembles[en].void:
            z = []
            v = []
            for c in p.ensembles[en].cells:
                z.append(depth + c.z_position)
                v.append(c.velocity.magn2d() if cfg['component'] == 3 else  c.velocity[cfg['component']] )
            z_dense = np.linspace(0, depth)
            v_dense = (2.5 * np.log(30 * z_dense / p.ensembles[en].ks)) * p.ensembles[en].v_shear 
            
            p1.set_data(v,z)
            p2.set_data(v_dense, z_dense)

    ens_slider = Slider(axes([0.25, 0.1, 0.65, 0.03]), label='Ensemble', valmin=0, valmax=len(p.ensembles)-1, valinit=initial_ensemble, valfmt='%1.0f')
    ens_slider.on_changed(update)
    plt.show()
    
    
def plot_logfit_profile(p, cfg):
    '''
    plot some data regarding the estimated roughness
    kwargs:
    p ...    the profile in question
    cfg ... config that was used for logfit_profile()
    '''
    import numpy as np
    import matplotlib.pyplot as plt    
    from matplotlib.font_manager import FontProperties
    
    data = np.zeros(shape=(len(p.ensembles), 8))
    data_mask = np.zeros(shape=(len(p.ensembles), 8))
    '''
    array indices:
    0    ks
    1    tau_shear
    2    v_shear
    3    a    a from ln(z) = a*v + b
    4    b    b from ln(z) = a*v + b
    5    r    pearson correlation coefficient
    6    n    numer of cells taken for fitting
    7    depth
    '''

    for i in range(len(p.ensembles)):
        if hasattr(p.ensembles[i], 'ks'):
            data[i,0] = p.ensembles[i].ks
            data[i,1] = p.ensembles[i].tau_shear
            data[i,2] = p.ensembles[i].v_shear
            data[i,3] = p.ensembles[i].logfit_debug[0]
            data[i,4] = p.ensembles[i].logfit_debug[1]
            data[i,5] = p.ensembles[i].logfit_debug[2]
            data[i,6] = p.ensembles[i].logfit_debug[4]
        else:
            data_mask[i,:] = 1
        data[i,7] = p.ensembles[i].depth if  p.ensembles[i].depth > 0 else 0
            
    data_masked = np.ma.masked_array(data, data_mask)

    x = list(range(len(p.ensembles)))
        
    plt.figure(1)
    plt.subplot(311)
    
    
    cell_hi = []
    cell_lo = []
    cell_lg = []        # depth of last cell used for logfit
    
    for e in p.ensembles:
        if not e.void:
            temp =  [-i.z_position for i in e.cells]
            cell_hi.append(min(temp)) # topmost cells
            cell_lo.append(max(temp)) # lowest cells
            
            z_max = e.depth * (1 - cfg['logheight'] )
            cell_lg.append(max([z if z<z_max else 0 for z in temp]))
#            cell_lg.append((e.cells[0] - e.cells[1]) * data[] )
        else:
            # if no void value
            cell_hi.append(0)
            cell_lo.append(0)
            cell_lg.append(0)
        
    cell_hi = np.array(cell_hi)
    cell_lo = np.array(cell_lo)
    cell_lg = np.array(cell_lg)
    
    plt.plot(x, data[:,7], color='0', linewidth=2, label='depth')
    plt.plot(x, data_masked[:,6], label='# of cells used')
    plt.gca().invert_yaxis()
    plt.bar(np.array(x)-0.5, height=cell_lo-cell_hi, width=[1]*len(x),  bottom=cell_hi, color='0.95', linewidth=0, label='existing cell')
    plt.bar(np.array(x)-0.5, height=-data_masked[:,6]*p.cell_height, bottom=cell_lo, width=[1]*len(x), color='0.8', linewidth=0, label='cell used for logfit')
    #plt.fill_between(x, cell_lo, cell_hi, color='0.95', label='existing cell')
    #plt.fill_between(x, cell_lo, cell_lg, color='0.8', label='cell used for logfit')
    plt.xlim(0,max(x))
    plt.ylabel('depth [m], # of used cells')
    plt.grid(True)
    
    # small legend font
    fontP = FontProperties()
    fontP.set_size('small')
    plt.legend(loc=2, bbox_to_anchor=(1,1), prop=fontP)
    
    
    
    # plot #2: log things
    plt.subplot(312)
    p1 = plt.scatter(x, abs(data_masked[:,0]), marker='x', color='0', label='ks')
    plt.ylim(0,100)
    plt.xlim(0,max(x))
    plt.ylabel('ks [m]')
#    plt.yscale('log')
    plt.grid(True)
    
    plt.twinx()
    p2 = plt.scatter(x, data_masked[:,1], marker='x', color='.5', label='tau_shear')
    plt.ylabel('tau_shear [N/m^2]')
    plt.legend([p1, p2], ['ks', 'tau_shear'], loc=2, bbox_to_anchor=(1.03,1), prop=fontP)
    plt.xlim(0,max(x))
    
    
    
    
    # plot #3: linear data
    plt.subplot(313)
    plt.scatter(x, abs(data_masked[:,5]),  marker='x', color='0')
    plt.ylabel('correlation coefficient')
    plt.xlim(0,max(x))
    plt.grid(True)
    
    
    
#    i = 0
#    print(('median: {}, min: {}, max: {}'.format(np.ma.median(data_masked[:,i]), data_masked[:,i].min(), data_masked[:,i].max())))
    
    plt.show()


def plot_profile_3d(p, cfg={}, vmin=None, vmax=None, axes=None, ymin=None, ymax=None):
    '''
    general profile plotting function
    
    cfg vars:
    datatype    string        determines what to plot. available types:
                            - "velocity" requires variable $component
                            - "custom" requires variable $cellattr
                            - "backscatter" (not implemented yet)
    components    list        a list of of velocity components, that will be used. valid components: "x", "y", "z"
    cellattr    string        the attribute of cell the cells that will be plotted. 
    style        string        visualization method. implemented methods: "contour", "gradient", "vector"
    saveas        string        path to output file, optional
    title        string        title of plot, optional
    '''
    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib import rc
    from outliers import get_cell_matrix, get_valuematrix_from_cellmatrix

    if 'style' in cfg and cfg['style'] not in ["contour", "gradient", "vector", "stream"]:
        raise NotImplementedError('plot_profile(): "{}" is not a valid plot style!'.format(cfg['style']))
    
    
    #
    # default values
    cfg_datatype = cfg['datatype'] if 'datatype' in cfg else 'velocity'
    cfg_components = cfg['components'] if 'components' in cfg else ['x', 'y']
    cfg_style = cfg['style'] if 'style' in cfg else 'gradient'
    
    
    #
    # produce profile coordinates
    
    
    # max number of cells, for the matrix size..
#    ymax = max([len(e.cells) for e in p.ensembles])

    
    # x
    e_positions = [e.position for e in p.ensembles]
    x = [0]
    for i in range(1, len(e_positions)):
        x.append(x[i-1] + (e_positions[i] - e_positions[i-1]).magn())
            
            
    # z
    # z with possible irregular values (when extrapolated)
    
    # ensemble with most cells, to steal the z positions of the cells
    ncells = [len(e.cells) if hasattr(e, 'cells') else 0 for e in p.ensembles]
    longestensemble = ncells.index(max(ncells))
    z = [c.z_position for c in p.ensembles[longestensemble].cells]        
    

    # profile depth
    z_depth = np.ma.masked_array([-e.depth for e in p.ensembles], [e.void for e in p.ensembles])
    z_depths = -np.ma.masked_array([e.four_depths for e in p.ensembles],[[e.void, e.void, e.void, e.void] for e in p.ensembles])
    
    # matrix with ProcesscedCell objects
    cm = get_cell_matrix(p)
    
    
    #
    # fetch data to be displayed
    data = []
    if cfg_datatype == 'velocity':
        for c in cfg_components:
            if c in ['x', 'y', 'z']:
                eval = '.velocity.'+c
                data.append(get_valuematrix_from_cellmatrix(cm, eval, outputs=1, nvalues=1))
            else:
                raise NotImplementedError('plot_profile(): "{}" is not a valid component!'.format(c))
    
    elif cfg_datatype == 'custom':
        for ca in cfg['cellattr']:
            data.append(get_valuematrix_from_cellmatrix(cm, ca, outputs=1, nvalues=1))
    
    else:
        raise NotImplementedError('plot_profile(): "{}" is not a valid datatype!'.format(cfg['datatype']))
    
    # matrix to mask the profile
    goodmat = get_valuematrix_from_cellmatrix(cm, eval, outputs=2, nvalues=1)
    
    #
    # prepare data depending on plot style
    plot_data = np.zeros(shape=(data[0].shape[0], data[0].shape[1], 2))
    
    for i in range(data[0].shape[0]):
        for j in range(data[0].shape[1]):
            if cfg_style != 'vector':
                # merge data into a scalar by computing its magnitude
                plot_data[i,j] = np.sqrt((np.array([d[i,j,0] for d in data]) ** 2).sum())
            else:
                # re-order data for vector plotting
                temp =  np.array([d[i,j,0] for d in data])
                plot_data[i,j] = temp
                
    # mask unusable matrix data (eg. parts not inside the river)
    plot_data = np.ma.masked_array(plot_data, np.dstack((~goodmat, ~goodmat)))

        
    #
    # actually plot
    
#    rc('text', usetex=True)        # enable this to use TeX type setting
#    rc('font', family='serif')    # font closer to LaTeX documents
    
    # river depth
    # FIXME see [] on how to customize line plots
    if axes is None:
        fig, ax = plt.figure(figsize=(20, 5))
    else:
         ax = axes

    ax.plot(x, z_depths, '0.5')
    ax.plot(x, z_depth, 'k', linewidth=2)

    if ymin is not None:
        ax.set_ylim(ymin, ymax)
    else:
        ax.set_ylim(ymax=0)
    ax.set_xlim(0, max(x))
#    plt.tight_layout()
#     ax.set_xlabel('Profile station [m]')
    ax.set_ylabel('Profile depth [m]')
    ax.set_title(cfg['title'] if 'title' in cfg else 'section view')
    
    # default "jet" color map is bad: http://www.jwave.vt.edu/~rkriz/Projects/create_color_table/color_07.pdf
    cmap = plt.cm.turbo    # @UndefinedVariable
#    cmap = plt.cm.Spectral    # @UndefinedVariable
    
    
    if cfg_style == 'gradient':
        # re align grid and append one value, because pcolor() puts values in between the x/z grid
        # see http://matplotlib.org/api/pyplot_api.html?highlight=pcolor#matplotlib.pyplot.pcolor
        x_pcolor = np.append(np.array(x), x[-1] + p.cell_height) - p.cell_height / 2.0
        z_pcolor = np.append(np.array(z), z[-1] - p.cell_height) + p.cell_height / 2.0
        colorPlot = ax.pcolor(x_pcolor, z_pcolor, plot_data[:,:,0], cmap=cmap)
#        plt.colorbar(pad=0.01)
    
    elif cfg_style == 'contour':
        number_of_contours = 10
        colorPlot = ax.contourf(x, z, plot_data[:,:,0], 256, interpolation='none',  cmap=cmap) # @UndefinedVariable, dunno why eclipse doent find this cm..
        ax.contour(x, z, plot_data[:,:,0], number_of_contours, interpolation='none') # @UndefinedVariable
        #plt.colorbar()
        
    elif cfg_style == 'vector':
        ax.quiver(x, z, plot_data[:,:,0], plot_data[:,:,1], width=0.001)
    
#    elif cfg_style == 'stream':
#        plt.streamplot(np.array(x), np.array(z), plot_data[:,:,0], plot_data[:,:,1])
    
    #fig.subplots_adjust(bottom=0.10, left=0.04, right=0.98, top=0.92)
    # color bar if needed
    # if cfg_style in ['gradient', 'contour']:
    #     # fixes to remove unneeded whitespace
    #     plt.colorbar(colorPlot, fraction=.05, pad=0.01, ax=axes)
    
    ax.grid(True)

    if vmin is not None:
        colorPlot.set_clim(vmin=vmin)
    if vmax is not None:
        colorPlot.set_clim(vmax=vmax)

    # optionally save the produced image. otherwise it will be displayed
    # if not 'saveas' in cfg or cfg['saveas'] == False:
    #     plt.show()
    # else:
    #     fig.savefig(cfg['saveas'])
    
    # clear figure and axis.
    #plt.clf()
    #plt.cla()
    
def plot_profile_2d(p, cfg=[]):
    ''' plot the ground view 
    cfg vars:
    vscale        float        scale the the displayed velocity vectors
    saveas        string        path to output file, optional
    title        string        title of plot, optional
    '''
    
    
    import matplotlib.pyplot as plt
    import numpy as np
    
    # position of ensembles
    x = [e.position.x for e in p.ensembles]
    y = [e.position.y for e in p.ensembles]
    
    # velocity components. expects uv_rot to be False!!
    vx = [e.velocity.x if not e.void else 0 for e in p.ensembles]
    vy = [e.velocity.y if not e.void else 0 for e in p.ensembles]
    
    # average length for the quiverkey
    avglen = np.sqrt(np.array(vx) ** 2 + np.array(vy) ** 2).mean()
    displen = np.around(avglen, np.int(np.ceil(np.abs(np.log10(avglen)))))
    
    # scale the vectors
    vscale = cfg['vscale'] if 'vscale' in cfg else 1
    #print(vscale)
    # FIXME: scale doesnt have any effect!
    
    plt.plot(x, y, marker='x', color='black')
    plt.scatter(x[0], y[0], marker='o')
    plt.axis('equal')
    plt.grid()
    q = plt.quiver(x, y, vx * vscale, vy * vscale, width=0.003, color='grey')
    plt.quiverkey(q, 0.2, 0.05, displen * vscale, 'velocity: {} [m/s]'.format(displen))
    
    plt.title(cfg['title'] if 'title' in cfg else 'plan view')
    plt.ylabel('y coordinate [m]')
    plt.xlabel('x coordinate [m]')
    
    if not 'saveas' in cfg or cfg['saveas'] == False:    
        plt.show()
    else:
        plt.savefig(cfg['saveas'])
    
    pass

    # clear figure and axis.
    plt.clf()
    plt.cla()
    
if __name__ == '__main__':    
    print('importing')
    import pickle    
    print('unpickle')
    p, p4, p5 = pickle.load(open('../testfiles/debug.pickle','rb'))
    print('plot_profile')
    
#    cfg_plot = dict(datatype='velocity', components=['y', 'z'], style='contour', saveas='../testfiles/testplot.pdf')
    cfg_plot = dict(datatype='velocity', components=['x', 'y'], style='vector', saveas=False)
    
    #plot_profile_3d(p, cfg_plot)
    from averaging import thin_out 
    p1 = thin_out(p, dict(keep_ensemble=5, keep_cell=1))
    #plot_profile_2d(p, cfg_plot)
    
    plot_profile_3d(p1, dict(Xsaveas='../testfiles/aspecttest.pdf', style='vector', components=['y','z']))