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

import math as math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as manimation


def reldiff(a,b):
    return (a-b)/(np.absolute(a).max()+np.absolute(b).max()+1e-20)

def go(model,courant,scheme,start):
    def plotter(metric, flow, scalars, t):
        mik, gas_ik, ujk, Phi_il, Wil = model.diagnose(*flow)
        Q1, Q2 = scalars
        
        mil = metric.ml_ext(mik)
        wil = Wil/mil*metric.g0
        wik = metric.mk_ext(Wil)/mik*metric.g0
        xik = metric.xik/1000
        xjl = metric.xjl/1000
        xil = metric.xil/1000
        zil = Phi_il/metric.g0/1000
        zik = metric.mk_int(zil)
        zjl = metric.mj(zil)
        print 's ',np.max(gas_ik.s),np.min(gas_ik.s)
        print 'ujk ',np.max(ujk),np.min(ujk)
        print 'wil ',np.max(wil),np.min(wil)
        sik = gas_ik.s
        sik = .5*(abs(sik)+sik)

        plt.clf()
        plt.contourf(xik, zik, sik, 20 )
#        plt.contourf(xik, zik, wik, 20 )
        plt.xlim((-1,1)), plt.xlabel('x (km)')
        plt.ylim((0,ztop/1000.)), plt.ylabel('z (km)')
        plt.colorbar()
        plt.title("t = %f s" % (t))
        plt.savefig('fig_bubble/bubble.png',bbox_inches='tight')
        writer.grab_frame()
        
    transport = advect.SplitTransport(model.metric, advect.HorizVanLeer, advect.VertVanLeer)
    replace = lambda fd,fv : fv
    return DCMIP.run_with_scalars(model,scheme,transport,replace, plotter, courant, start,scalars, T1, N1, N2)

Lx, ztop, nx, llm = 2000.,1000., 100, 50
nqdyn, ny, dx = 1, 1, Lx/nx
T1, N1, N2 = 5., 1, 200
#T1, N1, N2 = 5., 2, 5

metric, thermo, BC, m0ik, gas0, Phi0_il, u0_jk = DCMIP.thermal_bubble(Lx,nx,llm, ztop=ztop, zc=350.)
Sik = m0ik*gas0.s
start = (m0ik, Sik, u0_jk, Phi0_il, 0*Phi0_il)
scalars = (Sik, m0ik*metric.mk_int(Phi0_il))

model = dyn.Nonhydro(thermo,metric,BC)
solver, caldyn_eta = dyn.MassBasedNH(model), unst.eta_mass
#solver, caldyn_eta = dyn.LagrangianNH(model), unst.eta_lag

rho_bot=BC.rho_bot
if False: # explicit time-stepping with rigid lid
    def rigid(flow,tau): # impose rigid BCs at top and bottom
        flow, fast, slow = solver.bwd_fast_slow(flow,0.)
        fast[4][:,0]=0.
        fast[4][:,-1]=0.
        return flow,fast,slow
    BC.rho_bot = 100*rho_bot # moderate stiffness
    courant, time_scheme = 2.0, lambda dt : time_step.RK4(rigid, dt)
else: # HEVI time-stepping with pressure BC
    # Note : parameters BC.rigid_top and BC.rigid_bottom are True by default, but are ignored by the Python code...
    BC.rigid_top=False
    BC.rigid_bottom=False
    BC.rho_bot = 1e6*rho_bot # large stiffness
#    BC.rho_bot = 100*rho_bot # moderate stiffness
    courant, time_scheme = 3.0, lambda dt : time_step.ARK2(solver.bwd_fast_slow, dt, a32=0.7)

FFMpegWriter = manimation.writers['ffmpeg']
metadata = dict(title='Aquaplanet', artist='DYNAMICO_LMDZ5',
                comment='Movie support!')
writer = FFMpegWriter(fps=15, metadata=metadata, bitrate=4000, codec="h263p")
fig = plt.figure(figsize=(12,5))
with writer.saving(fig, "fig_bubble/bubble.avi", 100):
    m0ik, gas0_ik, u0jk, Phi0_il, W0il = go(model, courant, time_scheme, start)

#-------------------------------------------------------------------------------------------------------
# Now use final state to check that Fortran and Python implementations produce the exact same tendencies
#-------------------------------------------------------------------------------------------------------

#W0il[:,0]=0.
#W0il[:,-1]=0.

unst.setvar('nb_threads', 1)
mesh = Mesh(nx,ny,llm,nqdyn,Lx,ny*dx,0.)
xx_ik, xx_il, ll = mesh.xx[:,0,:]/1000, mesh.xxp1[:,0,:]/1000, mesh.ll[:,0,:]

mass_bl,mass_dak,mass_dbk=unst.compute_hybrid_coefs(m0ik)
unst.init_hybrid(mass_bl,mass_dak,mass_dbk)

caldyn_thermo = unst.thermo_entropy
#caldyn_thermo = unst.thermo_theta
g = mesh.dx/metric.dx_g0
caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,metric,thermo,BC,g)
unst.setvar('rigid',False)

s0ik = gas0_ik.s
if caldyn_thermo == unst.thermo_theta:
    s0ik = thermo.T0*np.exp(s0ik/thermo.Cpd)

# 2 momentum components in DYNAMICO vs 1 component in slice
u0jk = u0jk + 100*dx # shift horizontal wind by 100 m/s
u0_ker = mesh.field_u()
mesh.set_ucomp(u0_ker,u0jk)

tau=0.1
# impose hybdrid distribution of mass (Fortran)
flow = (m0ik/dx,s0ik*m0ik/dx,u0_ker,Phi0_il,W0il/dx) # NB : different definitions of mass field in DYNAMICO / Python
flow,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
# shift Phi from target value Phi_bot
Phi0_il = Phi0_il+g*100.*np.exp(-10.*xx_il*xx_il)
m,S,u,Phi,W = flow
flow = m,S,u,Phi0_il,W
# compute tendencies (0=Python)
flow0 = (m*dx,S*dx,u0jk,Phi0_il,W0il, 0., 0.) # NB : different definitions of mass field in DYNAMICO / Python
flow0, fast0, slow0 = solver.bwd_fast_slow(flow0,tau,True) 
# compute tendencies (Fortran)
flow,fast,slow = caldyn.bwd_fast_slow(flow, tau) 


zz_il = Phi0_il/metric.g0/1000.
zz_ik=.5*(zz_il[:,0:-1]+zz_il[:,1:])
print
print 'ptop, model top (km) :', unst.getvar('ptop'), zz_il.max()
        
for data0,data,dname in ((flow0,flow,'flow'),(slow0,slow,'slow'),(fast0,fast,'fast')):
    def plot(xx,zz,var,vname):
        title = "%s_%s" % (vname,dname)
        filename = "bubble/bubble_%s.png" % title
        var = 1.*var
        if var.size>1:
            plt.figure(figsize=(12,5))
            plt.contourf(xx,zz,var,20)
#                    plt.pcolor(xx,zz,var)
            plt.xlim((-1.,1.)), plt.xlabel('x (km)')
            plt.ylim((0.,1.)), plt.ylabel('z (km)')        
            plt.colorbar()
            plt.title(title)
            plt.savefig(filename,bbox_inches='tight')
            print title, var.min(), var.max()
            if var.max()>1e100:
                raise OverflowError
                    
    dm,dS,du,dPhi,dW = data    
    du=mesh.ucomp(du)/dx
    du=du[0,:,:]

    dm0, dS0, du0, dPhi0, dW0, junk, junk = data0
    dm0, dS0, du0, dW0 = dm0/dx, dS0/dx, du0/dx, dW0/dx

    if dname == 'flow':
        hflux=mesh.ucomp(caldyn.hflux)
        plot(xx_ik, zz_ik, dm,'m')
        plot(xx_ik, zz_ik, du,'u')
        plot(xx_ik, zz_ik, du-du0,'uu')
        plot(xx_ik, zz_ik, dS/m,'s')
        plot(xx_il, zz_il, dW,'w')
        plot(xx_il, zz_il, dW-dW0,'ww')
    else:
        dS=reldiff(dS,dS0)
        dW=dW-dW0
        dPhi=dPhi-dPhi0
        du=reldiff(du,du0)
        plot(xx_ik, zz_ik, dS,'dS')
        plot(xx_il, zz_il, dPhi,'dPhi')
        plot(xx_il, zz_il, dW,'dW')
        plot(xx_ik, zz_ik, du,'du')

if unst.getvar('rigid'):
    print 'Fortran BCs : rigid'
else:
    print 'Fortran BCs : pressure/soft', unst.getvars('ptop','rho_bot','Phi_bot','pbot')

if BC.rigid_top:
    print 'Top BC : rigid'
else:
    print 'Top BC : pressure', BC.ptop

if BC.rigid_bottom:
    print 'Bottom BC : rigid'
else:
    print 'Bottom BC : soft', BC.rho_bot[0], BC.Phi_bot[0], BC.pbot[0]