from dynamico import meshes
from dynamico import unstructured as unst
from dynamico import time_step
from dynamico import precision as prec

import math as math
import matplotlib.pyplot as plt
import numpy as np
import netCDF4 as cdf

#--------------------------------- Main program -----------------------------

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

filename, llm, nqdyn, g, f, radius = 'cart_%03d_%03d.nc'%(nx,ny), 1, 1, 1., 1., None
unst.setvar('g',g)

print 'Reading Cartesian mesh ...'
meshfile = meshes.DYNAMICO_Format(filename)

def coriolis(lon,lat):
   return f+0.*lon

mesh=meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, coriolis)
caldyn = unst.Caldyn_RSW(mesh)

#xx = (mesh.lat_i-nx/2.)*dx
#yy = (mesh.lon_i-ny/2.)*dy
xx,yy = mesh.lat_i, mesh.lon_i

x1,x2,yy = xx-1., xx+1., yy
u0 = 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(-5.*yy*yy))

flow0=prec.asnum([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=unst.np_num)
scheme = time_step.RK4(caldyn.bwd_fast_slow, dt, precision=unst.np_num)

flow=flow0

def minmax(name, x): print('Min/max %s :'%name, x.min(), x.max() )
def reshape(data): return data.reshape((nx,ny))

x, y = map(reshape, (xx,yy) )

for i in range(10):
    h,u=flow
    flow, fast, slow = caldyn.bwd_fast_slow(flow, prec.zero)
    junk, du_fast = fast
    dh, du_slow = slow
#    minmax('lon',mesh.lon_i)
#    minmax('lat',mesh.lat_i)
#    minmax('x',xx)
#    minmax('y',yy)
    minmax('PV', caldyn.qv-1.)
#    minmax('geopot', caldyn.geopot)
#    minmax('du_fast', du_fast)
    plt.figure(); plt.pcolor(x,y,reshape(caldyn.qv))
    plt.colorbar(); plt.title('potential vorticity')
    plt.savefig('fig_RSW_2D_mesh/%03d.png'%i)
    plt.close()
#    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)
#        flow = scheme.advance(flow,5)
        # 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.)