import numpy as np
import scipy as sp
from scipy.integrate import cumulative_trapezoid, trapezoid
from scipy.interpolate import interp1d
from tqdm import tqdm

from .gradient import gradient

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

def _generic_bl_integral(y, integrand, y_edge=None, j_edge=None, **kwargs):
    '''
    Compute a definite boundary-layer integral from y=0 to 'edge'
    -----
    The integration limit is given by j_edge (on-grid) or y_edge
      (off-grid, using interpolation of the cumulative integral)
    '''
    
    interp_kind = kwargs.get('interp_kind','cubic')
    
    ## Check coordinate [y] and integrand
    if y.ndim != 1 or integrand.ndim != 1:
        raise ValueError
    if y.shape[0] != integrand.shape[0]:
        raise ValueError
    if not np.all(np.diff(y) > 0): ## check (+) monotonicity
        raise ValueError("y must be monotonically increasing")
    
    ## Assert y[0]~0.
    if y[0] < -1e-4:
        raise ValueError("y[0] < 0")
    if not np.isclose(y[0], 0., atol=1e-4):
        raise ValueError("y[0] not sufficiently close to 0")
    
    ## Check distance 'y_edge' and/or index 'j_edge'
    if j_edge is not None and not isinstance(j_edge,(int,np.integer)):
        raise ValueError("j_edge must be int")
    if y_edge is not None:
        y_edge = float(y_edge)
    if (y_edge is not None) and (j_edge is not None):
        if not np.isclose(y[j_edge], y_edge, rtol=1e-4):
            raise ValueError("y_edge and j_edge inconsistent")
    if y_edge is None and j_edge is None: ## No edge defined; integrate full profile
        j_edge = -1
    
    if j_edge is not None: ## Index-defined 'edge'
        integral_val = trapezoid(
            integrand[:j_edge+1],
            x=y[:j_edge+1],
            )
    else: ## 'edge' distance supplied, perform 1D interpolation
        integral_cum = cumulative_trapezoid(
            integrand,
            x=y,
            initial=0.0,
            )
        integral_val = interp1d(
            y,
            integral_cum,
            kind=interp_kind,
            bounds_error=True,
            )(y_edge)
    
    return integral_val

def calc_d1(y, u, rho=None, y_edge=None, j_edge=None, **kwargs):
    '''
    Displacement (mass-flux deficit) thickness δ1 = δ*
    -----
    δ1 = ∫ ( 1 - ρu / (ρe·ue) ) dy
    '''
    interp_kind = kwargs.get('interp_kind','cubic')
    
    if y_edge is None and j_edge is None:
        j_edge = -1 ## integrate over complete profile
    
    if rho is None:
        rho = np.ones_like(u)
    
    if j_edge is None:
        u_e   = interp1d(y, u,   kind=interp_kind, bounds_error=True)(y_edge)
        rho_e = interp1d(y, rho, kind=interp_kind, bounds_error=True)(y_edge)
    else:
        u_e   = u[j_edge]
        rho_e = rho[j_edge]
    
    integrand = 1. - (rho*u)/(rho_e*u_e) ## δ1
    
    return _generic_bl_integral(
        y,
        integrand,
        y_edge=y_edge,
        j_edge=j_edge,
        interp_kind=interp_kind,
    )

def calc_d2(y, u, rho=None, y_edge=None, j_edge=None, **kwargs):
    '''
    Momentum deficit thickness δ2 = θ
    -----
    δ2 = ∫ ( ρu / (ρe·ue) ) ( 1 - u / ue ) dy
    '''
    interp_kind = kwargs.get('interp_kind','cubic')
    
    if y_edge is None and j_edge is None:
        j_edge = -1 ## integrate over complete profile
    
    if rho is None:
        rho = np.ones_like(u)
    
    if j_edge is None:
        u_e   = interp1d(y, u,   kind=interp_kind, bounds_error=True)(y_edge)
        rho_e = interp1d(y, rho, kind=interp_kind, bounds_error=True)(y_edge)
    else:
        u_e   = u[j_edge]
        rho_e = rho[j_edge]
    
    integrand = (rho*u)/(rho_e*u_e) * (1. - u/u_e) ## δ2
    
    return _generic_bl_integral(
        y,
        integrand,
        y_edge=y_edge,
        j_edge=j_edge,
        interp_kind=interp_kind,
    )

def calc_d3(y, u, rho=None, y_edge=None, j_edge=None, **kwargs):
    '''
    Kinetic energy deficit thickness δ3
    -----
    δ3 = ∫ ( ρu / (ρe·ue) ) ( 1 - (u / ue)^2 ) dy
    '''
    interp_kind = kwargs.get('interp_kind','cubic')
    
    if y_edge is None and j_edge is None:
        j_edge = -1 ## integrate over complete profile
    
    if rho is None:
        rho = np.ones_like(u)
    
    if j_edge is None:
        u_e   = interp1d(y, u,   kind=interp_kind, bounds_error=True)(y_edge)
        rho_e = interp1d(y, rho, kind=interp_kind, bounds_error=True)(y_edge)
    else:
        u_e   = u[j_edge]
        rho_e = rho[j_edge]
    
    integrand = (rho*u)/(rho_e*u_e) * (1. - (u/u_e)**2) ## δ3
    
    return _generic_bl_integral(
        y,
        integrand,
        y_edge=y_edge,
        j_edge=j_edge,
        interp_kind=interp_kind,
    )

def calc_dRC(y, u, u_tau, rho=None, rho_wall=None, y_edge=None, j_edge=None, **kwargs):
    '''
    Rotta-Clauser Δ = ∫W+dy = ∫(ρe+ue+ - ρ+u+) dy
    -----
    W  = (ρe·ue - ρ·u)/(ρe·ue) = 1 - (ρu/(ρe·ue))
    W+ = (ρe·ue - ρ·u)/(ρw·uτ) = ρe+ue+ - ρ+u+
    '''
    interp_kind = kwargs.get('interp_kind','cubic')
    
    if y_edge is None and j_edge is None:
        j_edge = -1 ## integrate over complete profile
    
    if rho is None:
        if rho_wall is not None:
            raise ValueError("rho_wall provided but rho is None")
        rho = np.ones_like(u)
    
    if rho_wall is None:
        rho_wall = rho[0]
    else:
        if not np.isclose(rho[0], rho_wall, rtol=1e-4):
            raise ValueError("rho_wall inconsistent with rho[0]")
    
    if j_edge is None:
        u_e   = interp1d(y, u,   bounds_error=True)(y_edge)
        rho_e = interp1d(y, rho, bounds_error=True)(y_edge)
    else:
        u_e   = u[j_edge]
        rho_e = rho[j_edge]
    
    integrand = (rho_e*u_e - rho*u) / (rho_wall*u_tau) ## W+
    
    ## Returned value is units of input [y]
    return _generic_bl_integral(
        y,
        integrand,
        y_edge=y_edge,
        j_edge=j_edge,
        interp_kind=interp_kind,
    )

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

def calc_profile_edge_1d(y, u=None, ddy_u=None, **kwargs):
    '''
    Determine the edge location of a 1D boundary-layer profile
    
    The edge is defined as the first location where |du+/dy+|<ϵ
    
    Parameters
    ----------
    y : np.ndarray
        1D wall-normal coordinate (units arbitrary)
    u : np.ndarray, optional
        Profile used to compute du/dy internally
    ddy_u : np.ndarray, optional
        Precomputed du/dy. Exactly one of `u` or `ddy_u` must be provided
    
    Keyword Arguments
    -----------------
    epsilon : float, default=5e-5
        Threshold ϵ for |du+/dy+|
    acc : int, default=6
        Accuracy order for numerical differentiation
    edge_stencil : str, default='half'
        Gradient stencil near boundaries
    interp_kind : {'linear','cubic'}, default='cubic'
        Interpolation method when `ongrid=False`
    ongrid : bool, default=True
        Snap to first grid point if True; otherwise interpolate
    
    Returns
    -------
    y_edge : float
        Wall-normal location of the boundary-layer edge
    
    Notes
    -----
    - Assumes |du+/dy+| at y=0 is O(1)
    - ϵ is adaptively increased if no crossing is found
    - Deprecated kwargs (`lchar`, `deriv`) are rejected
    '''
    
    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 OR 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) ## 5% rtol
    
    ## 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): ## Got to end, recalibrate & keep going
            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 = 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 = 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,
    rho_wall,
    nu_wall,
    **kwargs,
    ):
    '''
    For 1D profile, get [δ1=δ*, δ2=θ, Reθ, Reτ] etc.
    '''
    
    rtol        = kwargs.get('rtol',1e-6)
    interp_kind = kwargs.get('interp_kind','cubic')
    
    ## Checks
    if y.ndim != 1:
        raise ValueError
    if u.ndim != 1 or u.shape[0] != y.shape[0]:
        raise ValueError
    if rho.ndim != 1 or rho.shape[0] != y.shape[0]:
        raise ValueError
    
    for val in (u_tau, d99, y_edge, rho_edge, nu_edge, u_edge, rho_wall, nu_wall):
        if not isinstance(val, (float, np.floating)):
            raise ValueError
    
    if (y_edge > y.max()) or (y_edge < y.min()):
        raise ValueError
    
    ## Consistency checks for edge quantities
    rho_edge_ = interp1d(y, rho, kind=interp_kind, bounds_error=True)(y_edge)
    np.testing.assert_allclose(rho_edge, rho_edge_, rtol=rtol)
    
    u_edge_ = interp1d(y, u, kind=interp_kind, bounds_error=True)(y_edge)
    np.testing.assert_allclose(u_edge, u_edge_, rtol=rtol)
    
    # === Compressible integrals (root names)
    
    d1 = calc_d1(
        y, u,
        rho=rho,
        y_edge=y_edge,
        interp_kind=interp_kind,
        )
    dstar = d1
    
    d2 = calc_d2(
        y, u,
        rho=rho,
        y_edge=y_edge,
        interp_kind=interp_kind,
        )
    theta = d2
    
    d3 = calc_d3(
        y, u,
        rho=rho,
        y_edge=y_edge,
        interp_kind=interp_kind,
        )
    
    dRC = calc_dRC(
        y, u,
        u_tau=u_tau,
        rho=rho,
        rho_wall=rho_wall,
        y_edge=y_edge,
        interp_kind=interp_kind,
        )
    
    # === Kinetic (u-only) integrals
    
    d1_k = calc_d1(
        y, u,
        rho=None,
        y_edge=y_edge,
        interp_kind=interp_kind,
        )
    dstar_k = d1_k
    
    d2_k = calc_d2(
        y, u,
        rho=None,
        y_edge=y_edge,
        interp_kind=interp_kind,
        )
    theta_k = d2_k
    
    d3_k = calc_d3(
        y, u,
        rho=None,
        y_edge=y_edge,
        interp_kind=interp_kind,
        )
    
    dRC_k = calc_dRC(
        y, u,
        u_tau=u_tau,
        rho=None,
        rho_wall=None,
        y_edge=y_edge,
        interp_kind=interp_kind,
        )
    
    ## Shape factors
    H12    = d1   / d2
    H12_k  = d1_k / d2_k
    H32    = d3   / d2
    H32_k  = d3_k / d2_k
    
    ## Reynolds numbers
    Re_tau   = d99   * u_tau  / nu_wall
    Re_theta = theta * u_edge / nu_edge
    Re_d99   = d99   * u_edge / nu_edge
    
    dd = {
        'd1': d1,
        'd1_k': d1_k,
        'dstar': dstar,
        'dstar_k': dstar_k,
        
        'd2': d2,
        'd2_k': d2_k,
        'theta': theta,
        'theta_k': theta_k,
        
        'd3': d3,
        'd3_k': d3_k,
        
        'dRC': dRC,
        'dRC_k': dRC_k,
        
        'H12': H12,
        'H12_k': H12_k,
        'H32': H32,
        'H32_k': H32_k,
        
        'Re_tau': Re_tau,
        'Re_theta': Re_theta,
        'Re_d99': Re_d99,
        }
    
    return dd
