import numpy as np
import pseudo_spectral_v3_revised as ps3
import pseudo_spectral_v4_revised as ps4

# print(ps4.assign.__doc__)

#! basic parameters of HD 1897733b

a = 8.2e7 #m HJ Radius
omega = 3.2e-5 # s^-1
gwspeed = 2e3 #m/s
beta = 2*omega/a

mu_dim =3e7 # magentic diffusion coefficient

# nondimensionalizing scale:
L_scale = np.sqrt(gwspeed/beta)
U_scale = gwspeed
T_scale = (gwspeed*beta)**(-0.5)




# x_nondim = np.pi*a/L_scale
# y_nondim = np.pi*a/L_scale/2

# initial mean flow U0
R = 3700 # specific gas constant for hydrogen
Hp = 400e3
T = Hp*10/R
print(T)
# U0 = mu_dim*Rm/Hp

x_nondim = 2*np.pi
y_nondim = np.pi

# print(0.01*gwspeed/(10**(0.7)/Hp))


#! calculation set up

n_order = 500# number of basis functions  j
n_p = n_order
n_eqt = 5 # number of equations  k
n_var = 5 # number of varibles  l
k_x = 0.5 # wave number

mu =  mu_dim/(L_scale**2/T_scale)# magnetic diffusion coefficient
B_0 = 0 # external vertical magnetic field
# V_A = 0.04
Bx_mean = 0
By0 =0


tao_r = 5# radiation time 
tao_d = 5# drag time
p = (tao_r/tao_d)**(1/4) # Prandtl number


Nx = 500
Ny = 250

x = np.linspace(-x_nondim,x_nondim,Nx)
y = np.linspace(-y_nondim,y_nondim,Ny)

mean_flow = np.zeros((Ny,Nx),dtype=complex)
Q0 = 0.0
for i in range(Nx):
    mean_flow[:,i] = Q0*np.exp(-y**2/2)


# by = np.load('../data/mean_flow/Lm1.571V_A0.040_by.npy')
# print(by)
