'''
module to fill the unmeasured areas near surface and near riverbed with values
'''

def split_velocity(v_magn, v_avg):
    ''' split the computed velocity magnitude into a vector pointing
    in the average direction of the other vectors in the ensemble
    kwargs: 
    v_magn    float    the new velocity magnitude
    v_avg    Vector     the vector containing the average direction in spherical cooridnates
    returns:
    Vector with magnitude  v_magn and direction from v_avg in cartesian coords
    '''
    from vector import Vector 
    return Vector(v_magn, v_avg[1], v_avg[2]).tocartesian()
    

def    extrapolate_profile(p, cfg,  debug=False):
    ''' call extrapolate_ensemble() for every ensemble in profile p '''
    
    from copy import deepcopy 
    p = deepcopy(p)
    
    p.ensembles = [extrapolate_ensemble(e, p.cell_height, cfg, debug) for e in p.ensembles]
    p.update_velocities()

    return p


def    extrapolate_ensemble(e, cell_height, cfg, debug=False):
    '''
    extrapolate missing velocities in the blanking area below water surface
    and above river bed (but not on the river banks)
    
    cfg vars:
    topcells    int        number of cells counting from the top, that will be
                        used for fitting an extrapolation line below water surface
    method        int        method used for velocity extrapolation
                        1:    auto: linear on top, roughness based (log-law), on bottom
                            (if no roughness parameters exist, fallback to modified power law, see manual)
                        2:    powerlaw only: power law for both, based on equal current in measured zone
                        
    forcepowerlaw        bool        use powerlaw for both upper and lower zone!
    
    '''
    from copy import deepcopy
    import numpy as np
    from adcpprocess import ArtificialCellObj
    from vector import Vector
    e = deepcopy(e)
    
    # if whole ensemble is void: don't do anything
    if not e.void:
        # 
        # stuff thats needed in any case
        cells_top = []
        cells_bed = []
        n_cells_top = int(np.floor(-e.cells[0].z_position / cell_height)) # number of new cells on top
        n_cells_bottom = int(np.floor((e.depth + e.cells[-1].z_position) / cell_height)) # number of new cells on bottom
        
        '''
        for extrapolation on top:
        the height of the new cells will be calculated like this:
        z(n) = n * cell_height + d, 
        where    n    number of new cell, counting from 0 with the one on water surface
                d    z_position of first cell   '''
        d = -e.cells[0].z_position - n_cells_top * cell_height
            
        # compute averaged direction angles in ensemble
        v_spherical = np.array([c.velocity.tospherical() for c in e.cells])
        v_sph_avg = np.mean(v_spherical, axis=0)
        
        
        if debug:
            print('ensemble info:')
            print(( " total depth: {} (cells: {}), cell height: {} \n" 
                   " unmeasured top: {} (cells: {}) \n"
                   " unmeasured bottom: {} (cells: {}) \n".format(
                        e.depth, len(e.cells), cell_height, 
                        0 - e.cells[0].z_position, n_cells_top,
                        e.depth + e.cells[-1].z_position, n_cells_bottom
                         )))
        
        
        if not cfg['forcepowerlaw']:
            
            #
            # cells near water surface
            
            topcells = cfg['topcells']
            x = np.array([c.z_position for c in e.cells[:topcells]])
            v = []
            v.append([c.velocity[0] for c in e.cells[:topcells]])
            v.append([c.velocity[1] for c in e.cells[:topcells]])
            v.append([c.velocity[2] for c in e.cells[:topcells]])
            
            # linear regression
            a, b, r2 = [], [], []
            for i in range(3):
                res = np.linalg.lstsq(np.vstack([x, np.ones(len(x))]).T, v[i])
                a.append(res[0][0])
                b.append(res[0][1])
                #r2.append(res[1][0])    # correlation coefficient 
            
            
            cells_top = []
            for i in range(1, n_cells_top):
                position = i * -cell_height + d if i != 0 else 0
                velocity = Vector((np.array(a) * position + np.array(b)).tolist())
                cells_top.append(ArtificialCellObj(position, velocity))
            
                
            
            #
            # cells on bottom
            
            
            if hasattr(e, 'ks'):
                # roughness based estimation
#                print('ks based')
                
                
                z = e.cells[-1].z_position
                while True:
                    z -= cell_height
                    if z < -e.depth:
                        break
                    v_bed = (2.5 * np.log(30 / e.ks * (e.depth  + z))) * e.v_shear
                    #v_bed_vec = Vector(v_bed, v_sph_avg[1], v_sph_avg[2]).tocartesian()    # split the "length" of vector into its components
                    v_bed_vec = split_velocity(v_bed, v_sph_avg)
                    cells_bed.append(ArtificialCellObj(z, v_bed_vec))
                    
#                    print z, v_bed, v_bed_vec
            else:
#                print('power law based')
                # power law based extrapolation: 
                # find nearest cells to y = y_0.2 = 0.2 * h
                y_02 = 0.2 * e.depth
                z_08 = 0.8 * -e.depth
                            
                
                
                # check if this simple approach is possible
#                print(('depth: {}, lowest cell: {}, z08:{}'.format(e.depth, e.cells[-1].z_position, z_08)))
                if  e.cells[-1].z_position > z_08:
                    need_integral_powerlaw = True
                else:
                    # use simple approach
                    need_integral_powerlaw = False
                
                    v = []
                    v.append([c.velocity[0] for c in e.cells])
                    v.append([c.velocity[1] for c in e.cells])
                    v.append([c.velocity[2] for c in e.cells])
                    # interpolate to get velocity at y=0.2h
                    v_02 = []
                    z = np.array([c.z_position for c in e.cells])
                    v_02.append(np.interp(z_08, z, v[0]))
                    v_02.append(np.interp(z_08, z, v[1]))
                    v_02.append(np.interp(z_08, z, v[2]))
                        
                    # compute a (one for each component)
                    a = np.array(v_02) / (y_02 ** (1/6.0))
                    
                    # actually create new cells
                    z_new =  e.cells[-1].z_position
                    for i in range(int(np.floor((e.cells[-1].z_position + e.depth) / cell_height))-1):
                        z_new -= cell_height
                        y_new = e.depth + z_new
                        v_new = a * y_new ** (1/6.0)
                        cells_bed.append(ArtificialCellObj(z_new, Vector(v_new.tolist())))
                
                
                    
            if debug == True:
                z_top = np.array([c.z_position for c in cells_top])
                z_bed = np.array([c.z_position for c in cells_bed])
                z_orig = np.array([c.z_position for c in e.cells])
                
                v_top = np.array([c.velocity.magn2d() for c in cells_top])
                v_orig = np.array([c.velocity.magn2d() for c in e.cells])
                v_bed = np.array([c.velocity.magn2d() for c in cells_bed])
                
                import matplotlib.pyplot as plt
                i = -1
                plt.plot(z_top, v_top, marker='x')
                plt.plot(z_orig, v_orig, marker='x')
                plt.plot(z_bed, v_bed, marker='x')
                plt.ylim(ymin=0)
                plt.xlabel('depth z [m]')
                plt.ylabel('velocity magnitude [m/s]')
                plt.grid(True)
            
#                plt.show()
            
            

            
        elif cfg['forcepowerlaw'] or need_integral_powerlaw:
            # measured current Q_M
            v_sum = np.sum([c.velocity.magn2d() for c in e.cells])
            q_m = v_sum * cell_height
            
            # to be fitted current
            # u = a * y^b, Q_PL = a * zeta, where zeta = y^(b+1)/ (b+1)
            b = 1 / 6.0
            y1 = e.cells[-1].z_position + e.depth - cell_height/2
            y2 = e.cells[0].z_position + e.depth + cell_height/2
            zeta = (y2 ** (b+1) - y1 ** (b+1)) / (b+1)
            
            a = q_m / zeta
            
            #print(('Q_m: {}, Q_pl: {} [m3/s/ensemble]'.format(q_m, a*zeta)))
            
            # time for plotting
            if debug == True:
                import matplotlib.pyplot as plt
                z_orig = np.array([c.z_position for c in e.cells])
                v_orig = np.array([c.velocity.magn2d() for c in e.cells])
                
                z_new = np.linspace(0,-e.depth, 50)
                v_new = a * (z_new + e.depth) ** b
                
                plt.plot(v_orig, z_orig)
                plt.plot(v_new, z_new)
                plt.ylim(ymin=-e.depth,ymax=0)
                plt.grid()
#                plt.show()            

            
            # add cells to top if required
            if cfg['forcepowerlaw']:
                # then we are here because we also want the top cells to be extrapolated with this
                for i in range(n_cells_top):
                    #_pos = i * -cell_height + d if i != 0 else 0
                    y_pos = e.depth - d - i * cell_height
                    z_pos = y_pos - e.depth
                    vmagn = a * y_pos ** b
                    v = split_velocity(vmagn, v_sph_avg)
                    cells_top.append(ArtificialCellObj(z_pos, v))
            
            # add cells to bottom
            for i in range(n_cells_bottom):
                z_pos = e.cells[-1].z_position - (i+1) * cell_height 
                y_pos = z_pos + e.depth
                vmagn = a * y_pos ** b
                v = split_velocity(vmagn, v_sph_avg)
                cells_bed.append(ArtificialCellObj(z_pos, v))

    
        # concatenate the new and the existing cells
        e.cells = cells_top + e.cells + cells_bed
        
        if debug:
            print('ensemble summary:\n top:')
            for c in cells_top:
                print('{c.x} {c.y} {c.z} {m}'.format(c=c.velocity, m=c.velocity.magn2d()))
            print(' bottom:')
            for c in cells_bed:
                print('{c.x} {c.y} {c.z} {m}'.format(c=c.velocity, m=c.velocity.magn2d()))
        
            print('average:'.format(v_sph_avg))
    return e
    
    


if __name__ == '__main__':    
    import pickle
    from quickviz import *
        
    p0, p4, p5 = pickle.load(open('../testfiles/debug.pickle','rb'))
    

    p1 = extrapolate_profile(p4, dict(topcells=5, forcepowerlaw=1), debug=0)
#    print [c.velocity.magn() for c in p0.ensembles[240].cells]
#    
#    e1 = extrapolate_ensemble(p0.ensembles[240], .25, dict(topcells=5, forcepowerlaw=0), debug=False)
#    
#    print [c.velocity.magn() for c in p0.ensembles[240].cells]
#    print [c.velocity.magn() for c in e1.cells]

    
    plot_profile_3d(p1)