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

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

class DYNAMICO_Format(meshes.Mesh_Format):
    """ Netcdf input file handling.
    Creates a dictionary of names corresponding to original data file."""

    start_index=1 # indexing starts at 1 as in Fortran
    translate = { 'down_up':'edge_down_up', 'left_right':'edge_left_right' } #new name is edge_left_right

    def __init__(self,gridfilename):
        try:
            self.nc = cdf.Dataset(gridfilename, "r")
            print('#NC4: Opening file:',gridfilename)
        except RuntimeError:
            print """ Unable to open grid file %s, maybe you forgot to download it ?
            To do so, go to the 'Python/' dir and execute './get_MPAS_grids.sh'.
            """ % gridfile
            raise
    def get_pdim(self,comm,name):
        if name in self.translate: name = self.translate[name]
        return parallel.PDim(self.nc.dimensions[name], comm)
    def getvar(self,name):
        if name in self.translate: name = self.translate[name]
        return self.nc.variables[name][:]
    def get_pvar(self,pdim,name):
        # provides parallel access to a NetCDF variable, with name translation
        # first deal with special cases
        if name == 'edge_deg': return parallel.CstPArray1D(pdim, np.int32, 2)
        # general case
        if name in self.translate: name = self.translate[name]
        return parallel.PArray(pdim, self.nc.variables[name])
    def normalize(self, mesh, radius): pass

#--------------------------------- 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 = 'cart128.nc', 1, 1, 1., 1., 1.
unst.setvar('g',g)

print 'Reading Cartesian mesh ...'
meshfile = DYNAMICO_Format(filename)
#print("#NC4: check point: 0")
#print(meshfile)

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

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

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

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

#print('h0 : ', h0[1234]-1.)
#print('h0 : ', h0[2345]-1.)

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

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

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, 0.)
    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/%02d.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)
        # 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.)