# accelarations!
import numpy as np
from scipy import interpolate

def Bx0_func(y,V_A,Lm):
    return y/Lm* V_A* np.exp(-y**2/(2*Lm**2)) * np.exp(0.5)

def term1_func(y,u_bar,v,h,h_bar):
    Ny = len(y)
    tmp1 = np.zeros(Ny)
    tmp2 = np.zeros(Ny)
    for i in range(Ny):
        tmp1[i] = np.mean(v[i,:]*h[i,:])/h_bar[i]

    tck = interpolate.splrep(y,u_bar, s=0)
    dubardy = interpolate.splev(y,tck, der=1)
    for i in range(0,Ny):
        tmp2[i] = y[i]-dubardy[i]
    return tmp1*tmp2

def term2_func(h_bar,hv_dev,u_dev,y):
    Ny = len(y)
    tmp = np.zeros(Ny)
    dtmp = np.zeros(Ny)
    for i in range(Ny):
        tmp[i] = np.mean(hv_dev[i,:]*u_dev[i,:])
    
    tck = interpolate.splrep(y,tmp, s=0)
    dtmpdy = interpolate.splev(y,tck, der=1)

    for i in range(0,Ny):
        dtmp[i] = -dtmpdy[i]/h_bar[i]
    return dtmp

def term3_func(u_dev,Q_dev,Ru,h_bar):
    Ny = len(u_dev)
    tmp = np.zeros(Ny)
    for i in range(Ny):
        tmp[i] = np.mean(u_dev[i,:]*Q_dev[i,:])/h_bar[i] + np.mean(Ru[i,:]) /h_bar[i]
    return tmp

def term4_func(u,h,h_bar,tao_d):
    Ny = len(h_bar)
    tmp = np.zeros(Ny)
    for i in range(Ny):
        tmp[i] = -np.mean(h[i,:]*u[i,:])/h_bar[i]/tao_d
    return tmp

def term5_func(y,x,h,Bx,By,h_bar):
    Ny = len(y)
    Nx = len(x)
    tmp = np.zeros(Ny)
    dBxdx = np.zeros((Ny,Nx))
    dBxdy = np.zeros((Ny,Nx))


    for j in range(Nx):
        tck = interpolate.splrep(y,Bx[:,j], s=0)
        dBxdy[:,j]  = interpolate.splev(y,tck , der=1)


    for i in range(0,Ny):
        tck = interpolate.splrep(x,Bx[i,:], s=0)
        dBxdx[i,:] = interpolate.splev(x,tck, der=1)

    for i in range(Ny):
        tmp[i] = np.mean(h[i,:]*Bx[i,:]*dBxdx[i,:] + h[i,:]*By[i,:]*dBxdy[i,:]) / h_bar[i]
    return tmp

def acc_term(u,v,eta,bx,by,y,x,V_A,Lm,tao_r,tao_d,k_x):
    Ny = len(y)
    Nx = len(x)
    Bx0 = np.zeros((Ny,Nx))
    Q = np.zeros((Ny,Nx))
    Ru = np.zeros((Ny,Nx))

    # zonal mean value
    u_bar = np.zeros(Ny)
    v_bar = np.zeros(Ny)
    h_bar = np.zeros(Ny)
    hv_bar = np.zeros(Ny)
    Q_bar = np.zeros(Ny)
    # Ru_bar = np.zeros(Ny)

    # zonal mean deviations
    u_dev = np.zeros((Ny,Nx))
    # v_dev = np.zeros((Ny,Nx))
    # h_dev = np.zeros((Ny,Nx))
    hv_dev=np.zeros((Ny,Nx))
    Q_dev=np.zeros((Ny,Nx))

    for i in range(Ny):
        Bx0[i,:] = Bx0_func(y[i],V_A,Lm)

    u = np.real(u)
    v = np.real(v)
    h = np.real(eta) + 1
    Bx = np.real(Bx0+bx)
    By = np.real(by)

    for i in range(Ny):
        for j in range(Nx):
            Q[i,j] = (np.exp(-0.5*y[i]**2)*np.cos(k_x*x[j]) +1 -h[i,j])/ tao_r
            if(Q[i,j]>0):
                Ru[i,j] = -Q[i,j]*u[i,j]
            else:
                Ru[i,j] = 0

    for i in range(Ny):
        u_bar[i] = np.mean(u[i,:])
        v_bar[i] = np.mean(v[i,:])
        h_bar[i] = np.mean(h[i,:])
        hv_bar[i] = np.mean(h[i,:]*v[i,:])
        Q_bar[i] = np.mean(Q[i,:])


    for i in range(Ny):
        for j in range(Nx):
            u_dev[i,j] = u[i,j]-u_bar[i]
            hv_dev[i,j] = h[i,j]*v[i,j]-hv_bar[i]
            Q_dev[i,j] = Q[i,j]-Q_bar[i]
            
    term1 = term1_func(y,u_bar,v,h,h_bar)
    term2 = term2_func(h_bar,hv_dev,u_dev,y)
    term3 = term3_func(u_dev,Q_dev,Ru,h_bar)
    term4 = term4_func(u,h,h_bar,tao_d)
    term5 = term5_func(y,x,h,Bx,By,h_bar)

    return term1,term2,term3,term4,term5







