from dynamico.meshes import Cartesian_mesh as Mesh
from dynamico import meshes
from dynamico import unstructured as unst
from dynamico import dyn
from dynamico import time_step
from dynamico import DCMIP

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

def run(mesh,metric,thermo,BC,caldyn_eta,u, m0ik,gas0,Phi0_il, T=300., courant=2.9):
    caldyn_thermo = unst.thermo_entropy
    g = mesh.dx/metric.dx_g0

    dt = courant*.5*dx/np.sqrt(gas0.c2.max())
    nt=int(T/dt)+1
    dt=T/nt
    print 'nt,dt,dx', nt,dt,dx

    Sik = m0ik*gas0.s
    m0ik = m0ik/dx # NB : different definitions of mass field in DYNAMICO / Slice
    if caldyn_thermo == unst.thermo_theta:
        theta0 = thermo.T0*np.exp(gas0.s/thermo.Cpd)
        S0ik = m0ik*theta0
    else:
        S0ik = m0ik*gas0.s

    u0=mesh.field_u() 
    u0[range(0,2*nx,2),:] = u*mesh.dx # Doppler shift by u
    flow0=(m0ik,S0ik,u0,Phi0_il,0*Phi0_il)

    if False: # time stepping in Python
        caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,BC,g)
        scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt, a32=0.7)
        #scheme = time_step.ARK3(caldyn.bwd_fast_slow, dt)
        def next_flow(m,S,u,Phi,W):
            m,S,u,Phi,W = scheme.advance((m,S,u,Phi,W),nt)
            return m,S,u,Phi,W
    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,BC,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())

    m,S,u,Phi,W=flow0

    for i in range(12):
        time1=time.time()
        m,S,u,Phi,W = next_flow(m,S,u,Phi,W)
        time2=time.time()
        print 'ms per time step : ', 1000*(time2-time1)/nt

        w=.5*(W[:,0:llm]+W[:,1:llm+1])/m
#        junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
#        zz=caldyn.geopot[:,0:llm]/metric.g0/1000
        zz=Phi[:,0:llm]/metric.g0/1000
        print 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g')

        def plt_axes():
            plt.plot(xx,zz[:,llm/4])
            plt.plot(xx,zz[:,llm/2])
            plt.plot(xx,zz[:,3*llm/4])
            plt.ylim(0,ztop/1000.)
            plt.colorbar()

#        plt.figure(figsize=(12,3))
#        ucomp = mesh.ucomp(u)
#        plt.pcolor(xx,zz,caldyn.pk)
#        plt_axes()
#        plt.title("T(t = %f s)" % ((i+1)*dt*nt))
#        plt.savefig('fig_slice_GW_NH/T%02d.png'%i)
#        plt.close()
#
        plt.figure(figsize=(12,3))
        ucomp = mesh.ucomp(u)
        plt.pcolor(xx,zz,ucomp[:,:]/dx)
        plt_axes()
        plt.title("u(t = %f s)" % ((i+1)*dt*nt))
        plt.savefig('fig_slice_GW_NH/U%02d.png'%i)
        plt.close()

        plt.figure(figsize=(12,3))
        plt.pcolor(xx,zz,w)
        plt_axes()
        plt.title("w(t = %f s)" % ((i+1)*dt*nt))
        plt.savefig('fig_slice_GW_NH/W%02d.png'%i)
        plt.close()

def deform(x,eta): # x and eta go from 0 to 1
    h0, dd, xi = 500., 5000., 4000.
    cos2=lambda x: np.cos(np.pi*x)**2
    gauss=lambda x: np.exp(-x*x)
    xx = x*Lx
    phis = h0*gauss(xx/dd)*cos2(xx/xi)
    return (phis/ztop)*np.sin(np.pi*eta)

#Lx, nx, llm = 3e5, 100, 1e4, 30
Lx, nx, ztop, llm = 3e5, 300, 1e4, 80
nqdyn, ny, dx = 1, 1, Lx/nx
mesh = Mesh(nx,ny,llm,nqdyn,nx*dx,ny*dx,0.)
xx,ll = mesh.xx[:,0,:]/1000, mesh.ll[:,0,:]

metric, thermo, BC, m0ik, gas0, Phi0_il, u0_jk = DCMIP.DCMIP31(Lx, nx, llm, deform=deform, dTheta=1.)
#metric, thermo, BC, m0ik, gas0, Phi0_il, u0_jk = DCMIP.DCMIP21(Lx, nx, llm, deform=deform, h0=0.)

BC.rho_bot = 1e6*BC.rho_bot # large stiffness

mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(m0ik)
unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk)

run(mesh,metric,thermo,BC,unst.eta_mass,20.,  m0ik,gas0,Phi0_il, courant=2.) # lower courant number required if h0=2000m
#run(mesh,metric,thermo,BC,unst.eta_lag,20.,  m0ik,gas0,Phi0_il)