print 'Loading DYNAMICO modules ...'
from dynamico.meshes import morton_index, key_to_perm, reorder_indices, reorder_values, MPAS_Format, Unstructured_Mesh
print '...Done'
import numpy as np
import netCDF4 as cdf
import matplotlib.pyplot as plt
import time

#------------ Read command-line arguments ----------#

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("-n", type=int, default=2562,
                    help="number of primal cells in MPAS mesh")
args = parser.parse_args()
grid = args.n

#-------------- read MPAS mesh ----------------#

meshfile='grids_MPAS/x1.%d.grid.nc'%grid
print 'Reading MPAS mesh %s ...'%meshfile

def f(lon,lat): return 0.*lon
mesh=MPAS_Format(meshfile)
mesh=Unstructured_Mesh(mesh, 1, 1, 1., f)
primal_ne, dual_ne = mesh.primal_ne, mesh.dual_ne

mesh=cdf.Dataset(meshfile, "r")

def getdims(*names): return [len(mesh.dimensions[name]) for name in names]
def getvars(*names): 
    for name in names : 
        print "getvar %s ..."%name
        time1=time.time()
        ret=[mesh.variables[name][:] for name in names]
        print "... Done in %f seconds"%(time.time()-time1)
    return ret

primal_num, edge_num, dual_num = getdims('nCells','nEdges','nVertices')                                   
print 'Number of primal cells, dual cells and edges : %d, %d, %d '%(primal_num,dual_num,edge_num)
# get degree of primal, dual cells
primal_deg, trisk_deg = getvars('nEdgesOnCell','nEdgesOnEdge')
# get indices for stencils 
primal_edge, left_right = getvars('edgesOnCell','cellsOnEdge')
dual_edge, down_up = getvars('edgesOnVertex','verticesOnEdge')
primal_vertex, dual_vertex, trisk = getvars('verticesOnCell','cellsOnVertex','edgesOnEdge')
# get positions, lengths, surfaces and weights
le,de,Ai,Av = getvars('dvEdge','dcEdge','areaCell','areaTriangle')
lat_i,lon_i = getvars('latCell','lonCell')
lat_v,lon_v = getvars('latVertex','lonVertex')
lat_e,lon_e,angle_e = getvars('latEdge','lonEdge','angleEdge')
wee,Riv2 = getvars('weightsOnEdge','kiteAreasOnVertex')
mesh.close()

#------------ reorder mesh according to Morton index --------------#

def morton_order(lon,lat):
    x,y,z = np.cos(lat)*np.cos(lon), np.cos(lat)*np.sin(lon), np.sin(lat)
    keys = morton_index(x,y,z)
    return key_to_perm(keys)

perm_i = morton_order(lon_i,lat_i)
perm_v = morton_order(lon_v,lat_v)
perm_e = morton_order(lon_e,lat_e)

# MPAS => Python
for idx in primal_edge, left_right, dual_edge, down_up, primal_vertex, dual_vertex, trisk : idx[:]=idx[:]-1

reorder_indices(perm_i, primal_deg, primal_vertex, primal_edge, primal_ne, lon_i, lat_i, Ai)
reorder_indices(perm_v, dual_vertex, dual_edge, dual_ne, lon_v, lat_v, Av, Riv2)
reorder_indices(perm_e, trisk_deg, left_right, down_up, trisk, lon_e, lat_e, angle_e, le, de, wee)
reorder_values(perm_i, dual_vertex, left_right)
reorder_values(perm_v, primal_vertex, down_up)
reorder_values(perm_e, primal_edge, dual_edge, trisk)

# Python => MPAS
for idx in primal_edge, left_right, dual_edge, down_up, primal_vertex, dual_vertex, trisk : idx[:]=idx[:]+1

#--- Open copy of MPAS mesh file ; keep file structure but replace contents of variables ---#

mesh=cdf.Dataset('grids/x1.%d.grid.nc'%grid, "r+")

# create new variables dual_ne, primal_ne
mesh.createVariable('signOnCell','i4',('nCells','maxEdges'))
mesh.createVariable('signOnVertex','i4',('nVertices','vertexDegree'))

def putvars(datas,names): 
    for name,data in zip(names,datas) :
        print "putvar %s ..."%name
        time1=time.time()
        mesh.variables[name][:] = data[:]
        print "... Done in %f seconds"%(time.time()-time1)
    
putvars( (primal_deg, trisk_deg), ('nEdgesOnCell','nEdgesOnEdge') )
putvars( (primal_edge, left_right), ('edgesOnCell','cellsOnEdge') )
putvars( (dual_edge, down_up), ('edgesOnVertex','verticesOnEdge') )
putvars( (primal_vertex, dual_vertex, trisk), ('verticesOnCell','cellsOnVertex','edgesOnEdge') )
putvars( (le,de,Ai,Av), ('dvEdge','dcEdge','areaCell','areaTriangle') )
putvars( (lat_i,lon_i), ('latCell','lonCell') )
putvars( (lat_v,lon_v), ('latVertex','lonVertex') )
putvars( (lat_e,lon_e,angle_e), ('latEdge','lonEdge','angleEdge') )
putvars( (wee,Riv2), ('weightsOnEdge','kiteAreasOnVertex') )
putvars( (primal_ne,dual_ne), ('signOnCell','signOnVertex') )

mesh.close()