import math as math
import matplotlib.pyplot as plt
import numpy as np
import dynamico.dyn as dyn
import dynamico.time_step as time_step
import dynamico.DCMIP as DCMIP
import dynamico.advect as advect

dx, Nx, llm, tau_relax = 500., 400, 60, 25.
#dx, Nx, llm, tau_relax = 250., 800, 120, 25.
T1, N1, N2 = 30., 6, 20 # check energy every 60s, plot every 10min, run 10h
#dx, Nx, llm, tau_relax = 20000., 200, 200, 1e3
#T1, N1, N2 = 1800., 4, 6 # check energy every 60s, plot every 10min, run 2h

Lx=Nx*dx
metric, thermo, BC, m0ik, gas0, Phi0_il, u0_jk = DCMIP.DCMIP21(
    Lx, Nx, llm, h0=250., ztop=3e4, dd=4000, xi=3200)
#    Lx, Nx, llm, h0=250., ztop=3e4, dd=10*dx, xi=8*dx)
Sik = m0ik*gas0.s
start =(m0ik, Sik, u0_jk, Phi0_il, 0*Phi0_il)
scalars = (Sik, m0ik*metric.mk_int(Phi0_il))

relax = lambda Phi : (1./tau_relax)*(1+np.tanh((Phi/9.81-23e3)/3e3))

def sponge(Phi_il, vjk, dujk_dt, tau):
    tau_rel = 1./relax(Phi_il[:,0:-1])
    dujk = (vjk-u0_jk + tau*dujk_dt)/(tau+tau_rel)
    ujk = u0_jk + tau_rel*dujk
    dujk_dt = dujk_dt - dujk
    return ujk, dujk_dt

BC.sponge=sponge # will be called by fast_bwd
BC.rigid_top = False
BC.rigid_bottom = False
rho_bot = BC.rho_bot
BC.rho_bot = 1e6*rho_bot

def go(model,courant,scheme,start):
    def plotter(metric, flow, scalars, t):
        mik, gas_ik, ujk, Phi_il, Wil = model.diagnose(*flow)
        mil = metric.ml_ext(mik)
        wil = metric.g0*Wil/mil
        xil = metric.xil/1000
        xjl = metric.xjl/1000
        zil = Phi_il/metric.g0/1000
        zjl = metric.mj(zil)
        plt.figure(figsize=(12,5))
#        plt.figure(figsize=(6,5))
#        plt.pcolormesh(xjl, zjl, gas_ik.T-gas0.T)
#        plt.contourf(xil, zil, wil, np.arange(-.2,.2,.02))
        plt.contourf(xil, zil, wil, np.arange(-1.,1.,.1))
        plt.plot(xil,zil[:,llm/10])
        plt.plot(xil,zil[:,llm/5])
        plt.plot(xil,zil[:,llm/4])
        plt.plot(xil,zil[:,llm/2])
        plt.xlabel('x (km)') 
        plt.ylabel('z (km)')
#        plt.ylim((-1,30.)), plt.xlim((-dx/20.,dx/10.))
        plt.ylim((0.,10.)), plt.xlim((-10.,10.))
        plt.colorbar()
        plt.title("t = %f s" % (t))
        plt.savefig("fig_slice_orographic/%04d"%t)
        plt.close()
    transport = advect.SplitTransport(model.metric, advect.HorizVanLeer, advect.VertVanLeer)
    replace = lambda fd,fv : fv
    DCMIP.run_with_scalars(model,scheme,transport,replace, plotter, courant, start,scalars, T1, N1, N2)

model = dyn.Nonhydro(thermo,metric,BC)
solver = dyn.MassBasedNH(model)
#solver = dyn.LagrangianNH(model)

def bwd_fast_slow(flow,tau):
    flow,fast,slow=solver.bwd_fast_slow(flow,tau)
    m,S,u,Phi,W,X,Y = flow
    dm,dS,du,dPhi,dW,dX,dY = slow
    dW = dW - relax(Phi)*W # damp vertical velocity
    slow = dm,dS,du,dPhi,dW,dX,dY
    return flow,fast,slow

time_scheme = lambda dt : time_step.ARK2(bwd_fast_slow, dt, a32=0.7)
go(model, 3.0, time_scheme, start )