from dynamico.meshes import Cartesian_mesh as Mesh
from dynamico import unstructured as unst
from dynamico import dyn
from dynamico import time_step
from dynamico import DCMIP
from dynamico.precision import zero, np_num

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

nx,ny,llm,nqdyn=128,128,1,1
Lx,Ly,g,f = 8.,8.,1.,1.
dx,dy=Lx/nx,Ly/ny

unst.setvar('g',g)
mesh = Mesh(nx,ny,llm,nqdyn,Lx,Ly,f)
caldyn = unst.Caldyn_RSW(mesh)

x1,x2,yy = mesh.xx[:,:,0]-1., mesh.xx[:,:,0]+1., mesh.yy[:,:,0]
h0, u0 = mesh.field_mass(), mesh.field_u()  # flow initally at rest
h0[:] = 1+0.1*(np.exp(-2.*(x1*x1+yy*yy))+
            np.exp(-2.*(x2*x2+yy*yy)))
#h0[:] = 1+0.1*(np.exp(-2.*yy*yy))

flow0=(h0,u0)

cfl = .8
dt = cfl/math.sqrt((nx/Lx)**2+(ny/Ly)**2)

T=10.
N=int(T/dt)+1
dt=T/N
print N,dt,Lx/nx
scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt, precision=np_num)

print 'types : yy, u0, h0, u0', yy.dtype, u0.dtype, h0.dtype, u0.dtype

flow=flow0
for i in range(10):
    h,u=flow
    print 'types : h,u', h.dtype, u.dtype
    caldyn.bwd_fast_slow(flow, zero)
    plt.figure(); plt.pcolor(mesh.x,mesh.y,caldyn.qv)
    plt.colorbar(); plt.title('potential vorticity')
    plt.savefig('fig_RSW_2D/%03d.png'%i)
#    plt.figure(); plt.pcolor(mesh.x,mesh.y,h)
#    plt.colorbar(); plt.title('h')
#    plt.show()
#    plt.pcolor(x,y,vcomp(u)/dx)
#    plt.colorbar(); plt.title('v')
#    plt.show()
    for j in range(5):
        unst.elapsed=0.
        flow = scheme.advance(flow,N)
        h,u=flow
        print 'types : h,u', h.dtype, u.dtype
        # flops
        # mass flux : 1add+1mul per edge          =>  4
        # div U     : 4 add per cell              =>  4
        # KE        : 4*(2add+1mul) per cell      => 12
        # grad KE   : 1 add per edge              =>  2
        # grad h    : 1 add per edge              =>  2
        # qv    : 4+4+1 add +4mul + 1div per cell => 10
        # qu    : 1add+1mul per edge              =>  4
        # TrisK : 4add+4mul+4add+1add per edge    => 26
        # Total :                                    64 FLOPS/cell
        print i,j, unst.elapsed, 100.*(1.-unst.getvar('elapsed')/unst.elapsed)
        print 'GFlops', 64*(N*nx*ny)/unst.elapsed/1e9
        unst.setvar('elapsed',0.)