from dynamico import unstructured as unst
from dynamico import dyn
from dynamico import time_step
from dynamico import DCMIP
from dynamico import meshes
from dynamico.meshes import MPAS_Mesh as Mesh
import math as math
import matplotlib.pyplot as plt
import numpy as np
import time

#------------------------ initial condition -------------------------

Cpd, Rd, g = 1004.5, 287., 9.81
N, Teq, peq = 0.01, 300., 1e5   # background 
lonc, d2 = 2.*math.pi/3., 25e6  # perturbation
N2, g2 = N*N, g*g
N2_g2, g2_N2 = N2/g2, g2/N2
G, kappa = g2/(N2*Cpd), Rd/Cpd

def psTs(lat) : 
    cos2latm1 = np.cos(2*lat)-1.
    Ts = G+(Teq-G)*np.exp(-.25*N2_g2*u0*u0*cos2latm1)
    ps = peq*np.exp(u0/(4.*G*Rd)*u0*cos2latm1)*((Ts/Teq)**(1./kappa))
    return ps, Ts
    
def Phi_top(lat) : 
    ps, Ts = psTs(lat)
    return -g2_N2*np.log( 1+(Ts/G)*( ((ptop/ps)**kappa) - 1 ) )
    
def T0_p0(lon,lat,Phi):
    ps, Ts = psTs(lat)
    pb = ps*( (G/Ts)*(np.exp(-N2_g2*Phi)-1.)+1. )**(1./kappa)
    Tb = G + (Ts-G)*np.exp(N2_g2*Phi)
    r  = radius*np.arccos(np.cos(lat)*np.cos(lon-lonc))
    Theta = dT*np.sin(math.pi*Phi/(g*ztop))*d2/(d2+r*r)
    T = Tb + Theta*(pb/peq)**kappa
    return T, pb

def DCMIP31_3D(grid):
    def Phi(eta, lat): return eta*Phi_top(lat)
    def ulon(lat): return u0*np.cos(lat)
    def ulat(lat): return 0.*lat
    def f(lon,lat): return 0.*np.sin(lat) # Coriolis parameter


    nqdyn, preff, Treff = 1, 1e5, 300.
    thermo = dyn.Ideal_perfect(Cpd, Rd, preff, Treff)
    mesh = Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f)
    lev,levp1 = np.arange(llm),np.arange(llm+1)
    lat_ik,k_i = np.meshgrid(mesh.lat_i,lev, indexing='ij')
    lon_ik,k_i = np.meshgrid(mesh.lon_i,lev, indexing='ij')
    lat_il,l_i = np.meshgrid(mesh.lat_i,levp1, indexing='ij')            
    lon_il,l_i = np.meshgrid(mesh.lon_i,levp1, indexing='ij')
    lat_ek,k_e = np.meshgrid(mesh.lat_e,lev, indexing='ij')

    Phi_ik, Phi_il = Phi(k_i/float(llm),lat_ik), Phi(l_i/float(llm),lat_il)
    Tb_ik, pb_ik = T0_p0(lon_ik, lat_ik, Phi_ik)
    Tb_il, pb_il = T0_p0(lon_il, lat_il, Phi_il)
    gas = thermo.set_pT(pb_ik,Tb_ik)
    mass_ik = mesh.field_mass()
    for l in range(llm):
        mass_ik[:,l]=(pb_il[:,l]-pb_il[:,l+1])/g
    Sik  = gas.s*mass_ik    

# start at hydrostatic geopotential
    pb_ik = .5*(pb_il[:,1:]+pb_il[:,0:-1]) # pressure at full layers
    gas = thermo.set_ps(pb_ik,gas.s)
    for l in range(llm):
        Phi_il[:,l+1] = Phi_il[:,l] + g*mass_ik[:,l]*gas.v[:,l]

    ujk, Wil = mesh.ucov3D(ulon(lat_ek),ulat(lat_ek)), mesh.field_w()
    
    print 'ztop (m) = ', Phi_il[0,-1]/g, ztop
    print 'ptop (Pa) = ', gas.p[0,-1], ptop
    dx=mesh.de.min()
    params=dyn.Struct()
    params.g, params.ztop, params.ptop = g, ztop, ptop
    params.dx, params.dx_g0 = dx, dx/g
    params.pbot, params.rho_bot = 1e5+0.*mesh.lat_i, 1e6+0.*mesh.lat_i
    return thermo, mesh, params, (mass_ik,Sik,ujk,Phi_il,Wil), gas

#------------------------ main program -------------------------

ztop, u0, dT, Xfactor =  1e4, 20., 1., 125.
radius, grid, llm = 6.37122e6/Xfactor, 10242, 20
T, Nslice, courant = 360., 10, 3.0

junk, ptop = T0_p0(0.,0.,g*ztop)

thermo, mesh, params, flow0, gas0 = DCMIP31_3D(grid)
llm, dx = mesh.llm, params.dx
dz = params.ztop/llm
print 'llm, nb_hex, dx, dz =', llm, mesh.Ai.size, dx, dz

#dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dz**-2))
dt = courant*.5*dx/np.sqrt(gas0.c2.max())

nt = int(math.ceil(T/dt))
dt = T/nt
print 'Time step : %g s' % dt

mesh.plot_e(mesh.le/mesh.de) ; plt.title('le/de')
plt.savefig('fig_NH_3D_DCMIP31/le_de.png') ;plt.close()

mesh.plot_i(mesh.Ai) ; plt.title('Ai')
plt.savefig('fig_NH_3D_DCMIP31/Ai.png'); plt.close()

caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag

if False: # time stepping in Python
    caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g)
    scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt, a32=0.7)
    def next_flow(m,S,u,Phi,W): return scheme.advance((m,S,u,Phi,W),nt)
else: # time stepping in Fortran
    scheme = time_step.ARK2(None, dt, a32=0.7)
    caldyn_step = unst.caldyn_step_NH(mesh,scheme,nt, caldyn_thermo,caldyn_eta, thermo,params,params.g)
    def next_flow(m,S,u,Phi,W):
        caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u
        caldyn_step.geopot[:,:], caldyn_step.W[:,:] = Phi,W
        caldyn_step.next()
        return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(),
                caldyn_step.geopot.copy(), caldyn_step.W.copy())

def plots(it):
    s=S/m ; s=.5*(s+abs(s))
    for l in range(llm):
        w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l]
        z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
        
    print( 'ptop, model top (m), max(w), max(W) (m/s) :', 
          unst.getvar('ptop'), Phi.max()/unst.getvar('g'), w.max(), W.max() )

    mesh.plot_i(w[:,llm/2])
    plt.title('vertical velocity at t=%d'%(it*T))
    plt.savefig('fig_NH_3D_DCMIP31/w%02d.png'%it)
    plt.close()

w=mesh.field_mass()
z=mesh.field_mass()
m,S,u,Phi,W=flow0 
plots(0)

for it in range(Nslice):
    unst.setvar('debug_hevi_solver',False)
    time1, elapsed1 =time.time(), unst.getvar('elapsed')
    m,S,u,Phi,W = next_flow(m,S,u,Phi,W)
    time2, elapsed2 =time.time(), unst.getvar('elapsed')
    factor = 1000./nt
    print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
    factor = 1e9/(4*nt*w.size)
    print 'nanosec per gridpoint per call to caldyn_hevi : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
    plots(it+1)