from dynamico import unstructured as unst
from dynamico import dyn
from dynamico import time_step
from dynamico import DCMIP
from dynamico import meshes
import dynamico.xios as xios

import math as math
import matplotlib.pyplot as plt
import numpy as np
import time

from mpi4py import MPI
comm = MPI.COMM_WORLD
mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
print '%d/%d starting'%(mpi_rank,mpi_size)

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

# Parameters for the simulation
g              = 9.80616   # gravitational acceleration in meters per second squared
omega          = 7.29211e-5  # Earth's angular velocity in radians per second
f0             = 2.0*omega   # Coriolis parameter
u_0            = 20.0        # velocity in meters per second
T_0            = 288.0       # temperature in Kelvin
d2             = 1.5e6**2    # square of half width of Gaussian mountain profile in meters
h_0            = 2.0e3       # mountain height in meters
lon_c          = np.pi/2.0   # mountain peak longitudinal location in radians
lat_c          = np.pi/6.0   # mountain peak latitudinal location in radians
radius         = 6.371229e6  # mean radius of the Earth in meters
ref_press      = 100145.6    # reference pressure (mean surface pressure) in Pascals
ref_surf_press = 930.0e2     # South Pole surface pressure in Pascals
Rd             = 287.04      # ideal gas constant for dry air in joules per kilogram Kelvin
Cpd            = 1004.64     # specific heat at constant pressure in joules per kilogram Kelvin
kappa          = Rd/Cpd      # kappa=R_d/c_p
N_freq         = np.sqrt(g**2/(Cpd*T_0)) # Brunt-Vaisala buoyancy frequency
N2, g2kappa = N_freq**2, g*g*kappa

def DCMIP2008c5(grid,llm):
    def Phis(lon,lat): 
        rgrc = radius*np.arccos(np.sin(lat_c)*np.sin(lat)+np.cos(lat_c)*np.cos(lat)*np.cos(lon-lon_c))
        return g*h_0*np.exp(-rgrc**2/d2)
    def ps(lon, lat, Phis): 
        return ref_surf_press * np.exp(
            - radius*N2*u_0/(2.0*g2kappa)*(u_0/radius+f0)*(np.sin(lat)**2-1.)
            - N2/(g2kappa)*Phis )
    def ulon(lat): return u_0*np.cos(lat)
    def ulat(lat): return 0.*lat
    def f(lon,lat): return f0*np.sin(lat) # Coriolis parameter

    nqdyn, preff, Treff = 1, 1e5, 300.
    thermo = dyn.Ideal_perfect(Cpd, Rd, preff, Treff)

    meshfile = meshes.MPAS_Format('grids/x1.%d.grid.nc'%grid)
#    mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
    mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, f)
    mesh.pmesh=pmesh

    lev,levp1 = np.arange(llm),np.arange(llm+1)
    lon_i, lat_i, lon_e, lat_e =  mesh.lon_i, mesh.lat_i, mesh.lon_e, mesh.lat_e
    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')

    Phis_i = Phis(lon_i, lat_i)
    ps_i = ps(lon_i, lat_i, Phis_i)

    if llm==18:
        ap_l=[0.00251499, 0.00710361, 0.01904260, 0.04607560, 0.08181860,
              0.07869805, 0.07463175, 0.06955308, 0.06339061, 0.05621774, 0.04815296,
              0.03949230, 0.03058456, 0.02193336, 0.01403670, 0.007458598, 0.002646866,
              0.0, 0.0 ]
        bp_l=[0.0, 0.0, 0.0, 0.0, 0.0, 0.03756984, 0.08652625, 0.1476709, 0.221864,
              0.308222, 0.4053179, 0.509588, 0.6168328, 0.7209891, 0.816061, 0.8952581, 
              0.953189, 0.985056, 1.0 ]
    if llm==26:
        ap_l=[0.002194067, 0.004895209, 0.009882418, 0.01805201, 0.02983724, 0.04462334, 0.06160587, 
               0.07851243, 0.07731271, 0.07590131, 0.07424086, 0.07228744, 0.06998933, 0.06728574, 0.06410509, 
               0.06036322, 0.05596111, 0.05078225, 0.04468960, 0.03752191, 0.02908949, 0.02084739, 0.01334443, 
               0.00708499, 0.00252136, 0.0, 0.0 ]
        bp_l=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01505309, 0.03276228, 0.05359622, 
                0.07810627, 0.1069411, 0.1408637, 0.1807720, 0.2277220, 0.2829562, 0.3479364, 0.4243822, 
                0.5143168, 0.6201202, 0.7235355, 0.8176768, 0.8962153, 0.9534761, 0.9851122, 1.0 ]

    ap_l, bp_l = ref_press*np.asarray(ap_l[-1::-1]), bp_l[-1::-1] # reverse indices
    ptop = ap_l[-1] # pressure BC
    
    print ptop, ap_l, bp_l
    ps_il,ap_il = np.meshgrid(ps_i,ap_l, indexing='ij')
    ps_il,bp_il = np.meshgrid(ps_i,bp_l, indexing='ij')
    hybrid_coefs = meshes.mass_coefs_from_pressure_coefs(g, ap_il, bp_il)
    pb_il = ap_il + bp_il*ps_il
    mass_ik, pb_ik = mesh.field_mass(), mesh.field_mass()
    for l in range(llm):
        pb_ik[:,l]=.5*(pb_il[:,l]+pb_il[:,l+1])
        mass_ik[:,l]=(pb_il[:,l]-pb_il[:,l+1])/g
    Tb_ik = T_0 + 0.*pb_ik
    gas = thermo.set_pT(pb_ik,Tb_ik)
    Sik  = gas.s*mass_ik
# start at hydrostatic geopotential
    Phi_il = mesh.field_w()
    Phi_il[:,0]=Phis_i
    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
    print 'ptop (Pa) = ', gas.p[0,-1], ptop
    dx=mesh.de.min()
    params=dyn.Struct()
    params.g, params.ptop = g, 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, hybrid_coefs, params, (mass_ik,Sik,ujk,Phi_il,Wil), gas


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

#grid, llm = 40962, 26
nqtot, llm, grid = 1, 26, 2562
T, Nslice, courant = 14400, 24, 3.0
caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag
#caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass
thermo, mesh, hybrid_coefs, params, flow0, gas0 = DCMIP2008c5(grid,llm)
llm, dx = mesh.llm, params.dx
print 'llm, nb_hex, dx =', llm, mesh.Ai.size, dx
if caldyn_eta == unst.eta_lag:
    print 'Lagrangian coordinate.'
else:
    print 'Mass-based coordinate.'

unst.ker.dynamico_init_hybrid(*hybrid_coefs)
    
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_DCMIP2008c5/le_de.png'); plt.close()

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

scheme = time_step.ARK2(None, dt, a32=0.7)
caldyn_step = unst.caldyn_step_HPE(mesh,scheme,nt, caldyn_thermo,caldyn_eta, thermo,params,params.g)
def next_flow(m,S,u):
    caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u
    caldyn_step.remap()
    caldyn_step.next()
    return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), 
            caldyn_step.u.copy(), caldyn_step.geopot.copy())

def plots(it):
    s=S/m
    for l in range(llm):
        z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
        vol[:,l]=(Phi[:,l+1]-Phi[:,l])/params.g/m[:,l] # specific volume
    gas = thermo.set_vs(vol, s)
    s=.5*(s+abs(s))
    t = (it*T)/3600
    print( 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g') )
    mesh.plot_i(gas.T[:,llm/2])
    plt.title('T at t=%dh'%(t))
    plt.savefig('fig_DCMIP2008c5/T%02d.png'%it)
    plt.close()

    mesh.plot_i(m[:,llm/2])
    plt.title('mass at t=%dh'%(t))
    plt.savefig('fig_DCMIP2008c5/m%02d.png'%it)
    plt.close()

    mesh.plot_i(Phi[:,0])
    plt.title('Surface geopotential at t=%dh'%(t))
    plt.savefig('fig_DCMIP2008c5/Phis%02d.png'%it)
    plt.close()

z, vol = mesh.field_mass(), mesh.field_mass()
m,S,u,Phi,W = flow0 
caldyn_step.geopot[:,0]=Phi[:,0]
#plots(0)

context=xios.XIOS_Context(mesh.pmesh,mesh,nqtot, T)

for it in range(Nslice):
    context.update_calendar(it)
    unst.setvar('debug_hevi_solver',False)
    time1, elapsed1 =time.time(), unst.getvar('elapsed')
    m,S,u,Phi = next_flow(m,S,u)
    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*m.size)
    print 'nanosec per gridpoint per call to caldyn_hevi : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
    
    s=S/m
    s=.5*(s+abs(s))
    for l in range(llm):
        z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
        vol[:,l]=(Phi[:,l+1]-Phi[:,l])/params.g/m[:,l] # specific volume
    gas = thermo.set_vs(vol, s)
    ss = np.asarray(gas.T, dtype=np.double)
    context.send_field_primal('theta', ss)
    #plots(it+1)

print 'xios.context_finalize()'
context.finalize()
print 'xios.finalize()'
xios.finalize()
print 'Done'