# -*- coding: iso-8859-15 -*- # stupid python2 fix
'''
Created on 23.09.2012

@author: Jakob
@note: 
'''

import numpy as np







def logfit_ensemble(e, cfg):
    '''
    fits the log law to one measured velocity profile.
    kwargs:
    e ...    the ensemble with the velocity profile
    
    cfg:
    component ...    specify the velocity component, that is being processed (0=u, 1=v, 2=w, 3=magnitude)
    logheight ...    percentage (0 ... 1.0) that is used for fitting
    
    '''
    
    # this is just optional for debugging
    from scipy.stats.stats import pearsonr
    
    # prepare data in lists
    depth = e.depth
    z_max = depth * cfg['logheight']
    z = []
    v = []
    
    # NOTE: z is starting from bottom to top.
    
    if not e.void:
        for c in e.cells:
            if depth + c.z_position <= z_max:
                z.append(depth + c.z_position)
                v.append(c.velocity.magn2d() if cfg['component'] == 3 else  c.velocity[cfg['component']] )
                
        if len(z) > 1:
            # (np.log() is actually ln!)
            logz= np.log(z)
            
            # actual linear regression. basic formula: y = a*x + b 
            a, b = np.linalg.lstsq(np.vstack([v, np.ones(len(v))]).T, logz)[0]
            
            # correlation
            r,p = pearsonr(logz, v)
            
            # shear stress
            v_shear = 0.4 / a
            tau_shear = v_shear **2 * 1000 # assumption: density of water: 1000 kg/m^3
            
            # roughness height
            ks = 30 * np.exp(b)
            
        
            # write back results
            e.ks = ks
            e.tau_shear = tau_shear
            e.v_shear = v_shear
            e.logfit_debug = (a, b, r, p, len(z))
    
    return e


def logfit_profile(p, cfg):
    
    from copy import deepcopy
    p = deepcopy(p)
    
    for i in range(len(p.ensembles)):
        p.ensembles[i] = logfit_ensemble(p.ensembles[i], cfg)
        
    return p 


    
    
    
if __name__ == '__main__':
    pass    
#    import pickle    
#    from quickviz import *
#    
#    pre_profile = pickle.load(open('../testfiles/profile_stage0.pickle','rb'))
##    cfg_outlier = dict(radius_h=15, radius_v=0, limit=0.05)
#    
#    cfg_logfit = dict(component=3, logheight=0.25)
##    logfit_ensemble(pre_profile.ensembles[88], cfg_logfit)
##    
##    plot_logfit_ensemble(pre_profile.ensembles[88], cfg_logfit)
#    
#    p2 = logfit_profile(pre_profile, cfg_logfit)
#    plot_logfit_profile(p2, cfg_logfit)
#    #plot_logfit_ensemble(p2.ensembles[494], cfg_logfit)
##    plot_logfit_ensembles(p2, cfg_logfit)
#    
##    p3 = logfit_outliers(p2, cfg_outlier)