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 precision as prec
from dynamico.meshes import Cartesian_mesh as Mesh
import math as math
import matplotlib.pyplot as plt
import numpy as np
import time

def thermal_bubble_3D(Lx,nx,Ly,ny,llm,ztop=1000., zc=350.,
                      rc=250, thetac=0.5, x0=0., y0=0.):
    Cpd, Rd, g, p0,theta0, T0 = 1004.5, 287.,9.81, 1e5, 300., 300.
    nqdyn = 1
    
    Phi = lambda eta : g*ztop*eta
    p=lambda Phi : p0*np.exp(-Phi/(Rd*T0))
    zz = lambda p: -(Rd*T0*np.log(p/p0))/g
    rr = lambda x,y,p: np.sqrt((x-x0)**2 + (y-y0)**2 + (zz(p)-zc)**2)
    sa = lambda x,y,p: rr(x,y,p) < rc
    deform = lambda x,y,p: (0.5*thetac*(1+np.cos(np.pi*rr(x,y,p)/rc)))*sa(x,y,p)
    temp =  lambda p: theta0*(p/p0)**(Rd/Cpd)
    T = lambda x,y,p: deform(x,y,p) + temp(p)

    mesh = Mesh(nx,ny,llm,nqdyn,Lx,Ly,0.)
    thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0)

    Phi_il = Phi(mesh.llp1/float(llm))
    Phi_ik = Phi((mesh.ll+.5)/llm)
    p_ik = p(Phi_ik)
    T_ik = T(mesh.xx, mesh.yy, p_ik)

    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 'ztop (m) = ', Phi_il[0,0,-1]/g, ztop
    ptop = p(g*ztop)
    print 'ptop (Pa) = ', gas.p[0,0,-1], ptop
    params=dyn.Struct()
    params.ptop=ptop
    params.dx=dx
    params.dx_g0=dx/g
    params.g = g
    pbot = p(Phi_il[:,:,0])
    gas_bot = thermo.set_pT(pbot, temp(pbot))
    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

#Lx, nx, llm, thetac, T, Nslice, courant = 2000., 100, 50, 30., 5., 10, 2.8
Lx, nx, llm, thetac, T, Nslice, courant = 2000., 20, 79, 30, 5., 10, 2.8
#Lx, nx, llm, thetac, T, Nslice, courant = 3000., 75, 25, -30, 5., 10, 2.8

nqdyn, dx = 1, Lx/nx
Ly,ny,dy = Lx,nx,dx
thermo, mesh, params, flow0, gas0 = thermal_bubble_3D(Lx,nx,Ly,ny,llm, thetac=thetac)

# compute hybrid coefs from initial distribution of mass
mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0])
print 'Type of mass_bl, mass_dak, mass_dbk : ', [x.dtype for x in mass_bl, mass_dak, mass_dbk]
unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk)

dz = flow0[3].max()/(params.g*llm)
#dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2+dz**-2))
dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2))
nt = int(math.ceil(T/dt))
dt = T/nt
print 'Time step : %d x %g s' % (nt,dt)

#caldyn_thermo, caldyn_eta = unst.thermo_theta, unst.eta_mass
caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass
#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)
    def next_flow(m,S,u,Phi,W):
        # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
        return scheme.advance((m,S,u,Phi,W),nt)

else: # time stepping in Fortran
    scheme = time_step.ARK2(None, dt)
    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):
        # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
        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())
    
m,S,u,Phi,W=flow0
if caldyn_thermo == unst.thermo_theta:
    s=S/m
    theta = thermo.T0*np.exp(s/thermo.Cpd)
    S=m*theta
    title_format = 'Potential temperature at t=%g s (K)'
else:
    title_format = 'Specific entropy at t=%g s (J/K/kg)'

w=mesh.field_mass()
z=mesh.field_mass()

for it in range(Nslice):
    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) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g')
    jj=ny/2
    xx,zz,ss,ww = mesh.xx[jj,:,:]/1000., z[jj,:,:]/1000., s[jj,:,:], w[jj,:,:]

    f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4))
    
    c=ax1.contourf(xx,zz,ss,20)
    ax1.set_xlim((-.5,.5)), ax1.set_xlabel('x (km)')
    ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)')        
    plt.colorbar(c,ax=ax1)
    ax1.set_title(title_format % (it*T,))

    #    plt.show()
    
    #    plt.figure(figsize=(12,5))
    c=ax2.contourf(xx,zz,ww,20)
    ax2.set_xlim((-.5,.5)), ax2.set_xlabel('x (km)')
    ax2.set_ylim((0.,1.)), ax2.set_ylabel('z (km)')        
    plt.colorbar(c,ax=ax2)
    ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,))
    #    plt.tight_layout()
    #    plt.show()
    plt.savefig('fig_NH_3D_bubble/%02d.png'%it)
    
    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*nx*ny*llm)
    print 'nanosec per gridpoint per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1)