import numpy as np
import scipy as sp
from tqdm import tqdm

from .gradient import gradient

# ======================================================================

def calc_profile_edge_1d(y, u=None, ddy_u=None, **kwargs):
    '''
    determine the edge location of 1D BL profile
    - [y] is a 1D coordinate vector
    - if 'ddy_u' is passed: is an externally calculated du/dy (passing it skips gradient calc in this func)
    - if 'u' is passed: 'u' still needs to be d/dy()'ed (can be less efficient)
    -----
    - 'edge' is where |du/dy|<ϵ
    - this func is units-agnostic. it is assumed that |du/dy|_y=0 is O(1)
    '''
    
    epsilon      = kwargs.get('epsilon',5e-5)
    acc          = kwargs.get('acc',6)
    edge_stencil = kwargs.get('edge_stencil','half')
    interp_kind  = kwargs.get('interp_kind','cubic')
    ongrid       = kwargs.get('ongrid',True) ## snap to grid point
    
    ## these options were deprecated
    ## check that they are not passed
    lchar = kwargs.get('lchar',None)
    if (lchar is not None):
        raise ValueError('lchar was deprecated as kwarg for turbx.calc_profile_edge_1d()')
    deriv = kwargs.get('deriv',None)
    if (deriv is not None):
        raise ValueError('deriv was deprecated as kwarg for turbx.calc_profile_edge_1d()')
    
    # do_debug_plot = kwargs.get('do_debug_plot',False) ## show a debug plot for tuning ϵ
    
    ## checks
    if not isinstance(y, np.ndarray):
        raise ValueError('y should be a numpy array')
    if (y.ndim!=1):
        raise ValueError
    
    ## make sure exactly one of u | ddy_u was passed
    if (u is None) and (ddy_u is None):
        raise ValueError('pass one of either u or ddy_u')
    if (u is not None) and (ddy_u is not None):
        raise ValueError('either pass u or ddy_u (not both)')
    
    ## make sure dims of passed u | ddy_u are correct
    if (u is not None):
        if not isinstance(u, np.ndarray):
            raise ValueError
        if (u.ndim!=1):
            raise ValueError
        if (u.shape[0]!=y.shape[0]):
            raise ValueError
    if (ddy_u is not None):
        if not isinstance(ddy_u, np.ndarray):
            raise ValueError
        if (ddy_u.ndim!=1):
            raise ValueError
        if (ddy_u.shape[0]!=y.shape[0]):
            raise ValueError
    
    if not any([(interp_kind=='linear'),(interp_kind=='cubic')]):
        raise ValueError("'interp_kind' should be one of: 'linear','cubic'")
    
    ny = y.shape[0]
    
    ## ddy_u was NOT passed, need to calculate du/dy here
    if (ddy_u is None):
        ddy_u = gradient(u, y, axis=0, d=1, acc=acc, edge_stencil=edge_stencil)
    
    ## confirm that |d□/dy|_w is ~O(1)
    np.testing.assert_allclose( ddy_u[0], 1., rtol=0.05)
    
    ## static ϵ
    # ## get [y] of first intersection
    # jtop = ny-1
    # for j in range(ny):
    #     if ( np.abs(ddy_u[j]) < epsilon ): ## |d□/dy|<ϵ
    #         jtop = j
    #         break
    
    ## variable ϵ
    ## get [y] of first intersection
    keep_going = True
    while keep_going:
        jtop = ny-1
        for j in range(ny):
            if ( np.abs(ddy_u[j]) < epsilon ): ## |d□/dy|<ϵ
                jtop = j
                break
        if (jtop == ny-1):
            epsilon *= 1.05
            tqdm.write(f'recalibrating: ϵ={epsilon:0.3e}')
        else:
            keep_going = False
    
    if (jtop == ny-1):
        raise ValueError('jtop == ny-1')
    if (jtop<3):
        raise ValueError('jtop<3')
    
    if ongrid: ## grid-snapped, i.e. first point satisfying |d□/dy|<ϵ
        y_edge = float( y[jtop] )
    
    else: ## find interpolated |d□/dy|==ϵ
        
        #intrp_func = sp.interpolate.interp1d(y, ddy_u, kind=interp_kind, bounds_error=True) ## deprecated!
        
        if (interp_kind=='cubic'):
            
            #intrp_func = sp.interpolate.CubicSpline(y, ddy_u, axis=0, bc_type='natural', extrapolate=False)
            intrp_func = sp.interpolate.PchipInterpolator(y, ddy_u, axis=0, extrapolate=False)
            #intrp_func = sp.interpolate.Akima1DInterpolator(y, ddy_u, axis=0, method='akima', extrapolate=False)
        
        elif (interp_kind=='linear'):
            raise NotImplementedError
            ## need to use np.interp() --> no longer retuns a function! needs to be wrapped.
        else:
            raise NotImplementedError
        
        def __f_opt_edge_locator(y_test, intrp_func, epsilon):
            ddy_u_test = intrp_func(y_test)
            root = np.abs( ddy_u_test ) - epsilon
            return root
        
        sol = sp.optimize.least_squares(
                    fun=__f_opt_edge_locator,
                    args=(intrp_func,epsilon,),
                    x0 = y[0] + 0.99*(y[jtop]-y[0]),
                    #xtol=1e-15,
                    #ftol=1e-15,
                    #gtol=1e-15,
                    method='dogbox',
                    #method='trf',
                    #bounds=(y[jtop-2], y[jtop]),
                    bounds=(y[0], y[jtop]),
                    )
        if not sol.success:
            raise ValueError
        if ( sol.x.shape[0] != 1 ):
            raise ValueError
        
        y_edge = float(sol.x[0])
        
        if ( y_edge > y[jtop] ):
            raise ValueError
    
    # if do_debug_plot:
    #     plt.close('all')
    #     fig1 = plt.figure(figsize=(2*(16/9),2), dpi=300)
    #     ax1 = plt.gca()
    #     ax1.set_xscale('log', base=10)
    #     ax1.plot(
    #         ddy_u,
    #         y,
    #         lw=0.8,
    #         )
    #     #ax1.set_xlim(1e-30, 1e1)
    #     ax1.axhline(y=y_edge, linestyle='dashed', c='gray', zorder=1, lw=0.5)
    #     ax1.axvline(x=epsilon, linestyle='dashed', c='gray', zorder=1, lw=0.5)
    #     fig1.tight_layout(pad=0.25)
    #     fig1.tight_layout(pad=0.25)
    #     plt.show()
    
    return y_edge

def calc_d99_1d(y, u, y_edge, u_edge=None, **kwargs):
    '''
    determine δ99 location of 1D profile
    - [y] is a 1D coordinate vector
    - [u] is some profile variable (streamwise velocity, pseudovelocity, etc.)
    - [y_edge] is the edge location
    - [u_edge] (optional) is the variable value at the edge location
    '''
    
    rtol        = kwargs.get('rtol',1e-3) ## rtol for asserting u[y_edge]==u_edge
    interp_kind = kwargs.get('interp_kind','cubic')
    d95         = kwargs.get('d95',False) ## δ95 instead of δ99
    
    if (y.ndim!=1):
        raise ValueError
    if (u.ndim!=1):
        raise ValueError
    if (u.shape[0]!=y.shape[0]):
        raise ValueError
    if not isinstance(y_edge, (float, np.floating)):
        raise ValueError
    if (u_edge is not None):
        if not isinstance(u_edge, (float, np.floating)):
            raise ValueError
    
    if (y_edge>y.max()):
        raise ValueError
    if (y_edge<y.min()):
        raise ValueError
    
    ny = y.shape[0]
    
    j_edge = np.abs(y-y_edge).argmin()
    if ( y[j_edge] < y_edge ):
        j_edge += 1
    je = min(j_edge,ny)
    je += 1 ## for numpy [:n]
    
    intrp_func = sp.interpolate.interp1d(x=y[:je], y=u[:je], kind=interp_kind, bounds_error=True)
    
    if (u_edge is not None): ## assert passed u_edge is == f(y_edge)
        np.testing.assert_allclose(u_edge, intrp_func(y_edge), rtol=rtol)
    else:
        u_edge = intrp_func(y_edge)
    
    def __f_opt_d99_locator(y_test, intrp_func, u_edge):
        root = np.abs( 0.99*u_edge - intrp_func(y_test) )
        return root
    
    def __f_opt_d95_locator(y_test, intrp_func, u_edge):
        root = np.abs( 0.95*u_edge - intrp_func(y_test) )
        return root
    
    if d95:
        __f_opt = __f_opt_d95_locator
    else:
        __f_opt = __f_opt_d99_locator
    
    sol = sp.optimize.least_squares(
                fun=__f_opt,
                args=(intrp_func,u_edge),
                x0=0.1*y_edge,
                xtol=1e-15,
                ftol=1e-15,
                gtol=1e-15,
                method='dogbox',
                bounds=(y.min(), y_edge),
                )
    if not sol.success:
        raise ValueError
    if ( sol.x.shape[0] != 1 ):
        raise ValueError
    
    delta = float(sol.x[0]) ## δ99 or δ95
    
    if ( delta > y_edge ):
        raise ValueError
    
    return delta

def calc_bl_integral_quantities_1d( y, u, rho, u_tau, d99, y_edge, rho_edge, nu_edge, u_edge, nu_wall, **kwargs):
    '''
    for 1D profile get [θ, δ*, Re_θ, Re_τ]
    '''
    
    rtol        = kwargs.get('rtol',1e-6)
    interp_kind = kwargs.get('interp_kind','cubic')
    
    if (y.ndim!=1):
        raise ValueError
    
    if (u.ndim!=1):
        raise ValueError
    if (u.shape[0]!=y.shape[0]):
        raise ValueError
    if (rho.ndim!=1):
        raise ValueError
    if (rho.shape[0]!=y.shape[0]):
        raise ValueError
    
    ## type check
    if not isinstance(u_tau, (float, np.floating)):
        raise ValueError
    if not isinstance(d99, (float, np.floating)):
        raise ValueError
    if not isinstance(y_edge, (float, np.floating)):
        raise ValueError
    if not isinstance(rho_edge, (float, np.floating)):
        print(type(rho_edge))
        raise ValueError
    if not isinstance(nu_edge, (float, np.floating)):
        raise ValueError
    if not isinstance(u_edge, (float, np.floating)):
        raise ValueError
    if not isinstance(nu_wall, (float, np.floating)):
        raise ValueError
    
    if (y_edge>y.max()):
        raise ValueError
    if (y_edge<y.min()):
        raise ValueError
    
    #ny = y.shape[0]
    
    rho_edge_ = sp.interpolate.interp1d(y, rho, kind=interp_kind, bounds_error=True)(y_edge)
    np.testing.assert_allclose(rho_edge, rho_edge_, rtol=rtol)
    
    u_edge_ = sp.interpolate.interp1d(y, u, kind=interp_kind, bounds_error=True)(y_edge)
    np.testing.assert_allclose(u_edge, u_edge_, rtol=rtol)
    
    # ===
    
    integrand_theta_cmp = (u*rho)/(u_edge*rho_edge)*(1-(u/u_edge))
    integrand_dstar_cmp = (1-((u*rho)/(u_edge*rho_edge)))
    
    theta_cmp_     = sp.integrate.cumulative_trapezoid(y=integrand_theta_cmp, x=y, initial=0.)
    theta_cmp_func = sp.interpolate.interp1d(y, theta_cmp_, kind=interp_kind)
    theta_cmp      = theta_cmp_func(y_edge)
    
    dstar_cmp_     = sp.integrate.cumulative_trapezoid(y=integrand_dstar_cmp, x=y, initial=0.)
    dstar_cmp_func = sp.interpolate.interp1d(y, dstar_cmp_, kind=interp_kind)
    dstar_cmp      = dstar_cmp_func(y_edge)
    
    integrand_theta_inc = (u/u_edge)*(1-(u/u_edge))
    integrand_dstar_inc = 1-(u/u_edge)
    
    theta_inc_     = sp.integrate.cumulative_trapezoid(y=integrand_theta_inc, x=y, initial=0.)
    theta_inc_func = sp.interpolate.interp1d(y, theta_inc_, kind=interp_kind)
    theta_inc      = theta_inc_func(y_edge)
    
    dstar_inc_     = sp.integrate.cumulative_trapezoid(y=integrand_dstar_inc, x=y, initial=0.)
    dstar_inc_func = sp.interpolate.interp1d(y, dstar_inc_, kind=interp_kind)
    dstar_inc      = dstar_inc_func(y_edge)
    
    # ===
    
    H12     = dstar_cmp/theta_cmp
    H12_inc = dstar_inc/theta_inc
    
    Re_tau       = d99*u_tau/nu_wall
    Re_theta     = theta_cmp*u_edge/nu_edge
    Re_theta_inc = theta_inc*u_edge/nu_edge
    Re_d99       = d99*u_edge/nu_edge
    
    dd = { 'theta_cmp':theta_cmp,
           'dstar_cmp':dstar_cmp,
           'theta_inc':theta_inc,
           'dstar_inc':dstar_inc,
           'H12':H12,
           'H12_inc':H12_inc,
           'Re_tau':Re_tau,
           'Re_theta':Re_theta,
           'Re_theta_inc':Re_theta_inc,
           'Re_d99':Re_d99,
         }
    
    return dd
