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 import xios
from dynamico import precision as prec
from dynamico.meshes import Cartesian_mesh as Mesh

import math as math
import numpy as np
import time
from numpy import pi, log, exp, sin, cos

# Baroclinic instability test based on Ullrich et al. 2015, QJRMS

def baroclinic_3D(Lx,nx,Ly,ny,llm,ztop=25000.):
    Rd    = 287.0      # Gas constant for dryy air (j kg^-1 K^-1)
    T0    = 288.0      # Reference temperature (K) 
    lap   = 0.005      # Lapse rate (K m^-1)
    b     = 2.         # Non dimensional vertical width parameter
    u0    = 35.        # Reference zonal wind speed (m s^-1)
    a     = 6.371229e6 # Radius of the Earth (m)
    ptop  = 2000.
    y0    = Ly*0.5
    Cpd   = 1004.5
    p0    = 1e5

    omega  = 7.292e-5  # Angular velocity of the Earth (s^-1)    
    phi0   = 45.       # Reference latitude North pi/4 (deg)
    f0     = 2*omega*np.sin(phi0) 
    beta0  = 2*omega*np.cos(phi0)/a
    fb     = 2*omega*np.sin(phi0) - y0*2*omega*np.cos(phi0)/a

    def Phi_xy(x,y):
        fc = y*y - (Ly*y/pi)*sin(2*pi*y/Ly)
        fd = Ly*Ly/(2*pi*pi)*cos(2*pi*y/Ly)
        return .5*u0*( fb*(y-y0-Ly/(2*pi)*sin(2*pi*y/Ly)) + .5*beta0*(fc-fd-(Ly*Ly/3.)- Ly*Ly/(2*pi*pi)) )

    def Phi_xyeta(x,y,eta): return T0*g/lap*(1-eta**(Rd*lap/g)) + Phi_xy(x,y)*log(eta)*exp(-((log(eta)/b)**2))
    def ulon(x,y,eta): return -u0*(sin(pi*y/Ly)**2)*log(eta)*(eta**(-log(eta)/b/b))
    def tmean(eta) : return T0*eta**(Rd*lap/g)
    def T(x,y,eta) : return tmean(eta)+(Phi_xy(x,y)/Rd)*(((2/(b*b))*(log(eta))**2)-1)*exp(-((0.5*log(eta))**2)) 
    def p(eta): return p0*eta     # eta  = p/p_s

    def eta(alpha) : return (1-(lap*ztop*alpha/(T0)))**(g/(Rd*lap)) # roughly equispaced levels

    filename = 'cart_%03d_%03d.nc'%(nx,ny)
    print 'Reading Cartesian mesh ...'
    def coriolis(lon,lat): return f0+0.*lon
    meshfile = meshes.DYNAMICO_Format(filename)
    nqdyn, radius = 1, None
    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
    mesh  = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)

    alpha_k = (np.arange(llm)  +.5)/llm
    alpha_l = (np.arange(llm+1)+ 0.)/llm
    x_ik, alpha_ik = np.meshgrid(mesh.lon_i, alpha_k, indexing='ij')
    y_ik, alpha_ik = np.meshgrid(mesh.lat_i, alpha_k, indexing='ij')
    x_il, alpha_il = np.meshgrid(mesh.lon_i, alpha_l, indexing='ij')
    y_il, alpha_il = np.meshgrid(mesh.lat_i, alpha_l, indexing='ij')
    x_ek, alpha_ek = np.meshgrid(mesh.lon_e, alpha_k, indexing='ij')
    y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij')

    print('----------------')
    print 'ztop(ptop) according to Eq. 7:', T0/lap*(1.-(ptop/p0)**(Rd*lap/g)) 
    print(np.shape(alpha_k),np.shape(alpha_l))
    print(mesh.__dict__.keys())
    thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0)

    eta_il = eta(alpha_il)
    eta_ik = eta(alpha_ik)
    eta_ek = eta(alpha_ek)
    print('min max eta_il', np.min(eta_il),np.max(eta_il))

    Phi_il = Phi_xyeta(x_il, y_il, eta_il)
    Phi_ik = Phi_xyeta(x_ik, y_ik, eta_ik)
    p_ik = p(eta_ik)
    T_ik = T(x_ik, y_ik, eta_ik)  #ik full level(40), il(41)

    gas = thermo.set_pT(p_ik,T_ik)
    mass_ik = mesh.field_mass()
    for l in range(llm): 
        mass_ik[:,l]=(Phi_il[:,l+1]-Phi_il[:,l])/(g*gas.v[:,l])
    Sik, ujk, Wil  = gas.s*mass_ik, mesh.field_u(), mesh.field_w()
    
    print(np.shape(ujk))
    print('P_ik',p_ik[0,:])

    u_ek = mesh.ucov3D(ulon(x_ek, y_ek, eta_ek), 0.*eta_ek)

    print(np.shape(u_ek))

    print 'ztop (m) = ', Phi_il[0,-1]/g, ztop
    ptop = p(eta(1.))
    print 'ptop (Pa) = ', gas.p[0,-1], ptop

    params=dyn.Struct()
    params.ptop=ptop
    params.dx=dx
    params.dx_g0=dx/g
    params.g = g

    # define parameters for lower BC
    pbot = p(eta_il[:,0])
    print 'min p, T :', pbot.min(), tmean(pbot/p0)
    gas_bot = thermo.set_pT(pbot, tmean(pbot/p0))
    params.pbot = gas_bot.p
    params.rho_bot = 1e6/gas_bot.v
    
    return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas

with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator
    comm = client.comm
    mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
    print '%d/%d starting'%(mpi_rank,mpi_size)

    g, Lx, Ly = 9.81, 4e7, 6e6
    nx, ny, llm = 200, 30, 22
    dx,dy=Lx/nx,Ly/ny

    unst.setvar('g',g)

    thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm)

    T, nslice, dt = 3600., 1, 3600.

    with xios.Context_Curvilinear(mesh,1, dt) as context:
        # now XIOS knows about the mesh and we can write to disk
        for i in range(48):
            context.update_calendar(i)
            print 'send_field', i, gas0.T.shape
    #    context.send_field_primal('ps', lat_i)
            context.send_field_primal('temp', gas0.T)
            context.send_field_primal('p', gas0.p)