print 'Starting'

from mpi4py import MPI
comm = MPI.COMM_WORLD
mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
print '%d/%d starting'%(mpi_rank,mpi_size)

#from dynamico import partition
from dynamico import parallel, meshes
from dynamico import unstructured as unst

import math as math
import numpy as np
import netCDF4 as cdf

import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection

print 'Done loading modules'

#----------------- partition hand-written 15-cell mesh ------------------#

if mpi_size<15:
    send=np.random.randn(mpi_size)
    recv=np.zeros(mpi_size)
    comm.Alltoall(send,recv)
    #time.sleep(mpi_rank)
#    print mpi_rank, send, recv
    
    adjncy=[1, 5, 0, 2, 6, 1, 3, 7, 2, 4, 8, 3, 9, 
              0, 6, 10, 1, 5, 7, 11, 2, 6, 8, 12, 3, 7, 9, 13, 4, 8, 14,
              5, 11, 6, 10, 12, 7, 11, 13, 8, 12, 14, 9, 13 ]
    xadj=[0, 2, 5, 8, 11, 13, 16, 20, 24, 28, 31, 33, 36, 39, 42, 44]
    
    nb_vert = len(xadj)-1
    vtxdist = [i*nb_vert/mpi_size for i in range(mpi_size+1)]
    xadj, adjncy, vtxdist = [np.asarray(x,np.int32) for x in xadj,adjncy,vtxdist]
    
    idx_start = vtxdist[mpi_rank]
    idx_end = vtxdist[mpi_rank+1]
    nb_vert = idx_end - idx_start
    
    xadj_loc = xadj[idx_start:idx_end+1]-xadj[idx_start]
    adjncy_loc = adjncy[ xadj[idx_start]:xadj[idx_end] ]
    part = 0*xadj_loc[0:-1];

    unst.ker.dynamico_partition_graph(mpi_rank, mpi_size, vtxdist, xadj_loc, adjncy_loc, 4, part)

#    for i in range(len(part)):
#        print 'vertex', i+idx_start, 'proc', part[i]

#-----------------------------------------------------------------------------#
#---------------         partition and plot MPAS mesh       ------------------#
#-----------------------------------------------------------------------------#

# Helper functions to plot unstructured graph

def local_mesh(get_mycells):
    #    mydegree, mybounds = [get_mycells(x) for x in nEdgesOnCell, verticesOnCell]
    mydegree, mybounds = [get_mycells(x) for x in primal_deg, primal_vertex]
    print '%d : len(mydegree)=%d'%(mpi_rank, len(mydegree))
    vertex_list = sorted(set(unst.list_stencil(mydegree,mybounds))) 
    print '%d : len(vertex_list))=%d'%(mpi_rank, len(vertex_list))
    get_myvertices = parallel.Get_Indices(dim_vertex, vertex_list)
    mylon, mylat = [get_myvertices(x)*180./math.pi for x in lonVertex, latVertex]
    vertex_dict = parallel.inverse_list(vertex_list)
    meshes.reindex(vertex_dict, mydegree, mybounds)
    return vertex_list, mydegree, mybounds, mylon, mylat

def members(struct, *names): return [struct.__dict__ [name] for name in names]

#--------------- read MPAS grid file ---------------#

grid = 'x1.2562'
#grid = 'x1.10242'
#grid = 'x4.163842'
print 'Reading MPAS file %s ...'%grid

meshfile = meshes.MPAS_Format('grids/%s.grid.nc'%grid)
pmesh = meshes.Unstructured_PMesh(comm, meshfile)
pmesh.partition_metis()

def coriolis(lon,lat): return 0.*lat
llm, nqdyn, radius = 1,1,1.
lmesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)

(primal_deg, primal_vertex, dim_vertex, dim_cell, cell_owner, 
 lonVertex, latVertex, lonCell, latCell) = members(
    pmesh, 'primal_deg', 'primal_vertex', 'dim_dual', 'dim_primal', 'primal_owner', 
    'lon_v', 'lat_v', 'lon_i', 'lat_i')

local_num, total_num, com_cells = np.zeros(1), np.zeros(1), lmesh.com_primal
local_num[0]=com_cells.own_len
comm.Reduce(local_num, total_num, op=MPI.SUM, root=0)
if(mpi_rank==0): print 'total num :', total_num[0], dim_cell.n

#---------------------------- plot -----------------------------#

print 'Plotting ...'

halo_vertex_list, mydegree, mybounds, mylon, mylat = local_mesh(com_cells.get_all)
buf = parallel.LocalArray1(com_cells)

fig, ax = plt.subplots()
buf.read_own(latCell) # reads only own values
buf.data = np.cos(10.*buf.data)
buf.update() # updates halo
lmesh.plot_patches(ax,[-math.pi/2,math.pi/2], mydegree, mybounds, mylon, mylat, buf.data)
plt.xlim(-190.,190.)
plt.ylim(-90.,90.)
plt.savefig('fig_partition/A%03d.pdf'%mpi_rank, dpi=1600)

fig, ax = plt.subplots()
buf.read_own(cell_owner)
buf.update()
lmesh.plot_patches(ax,[0,mpi_rank+1], mydegree, mybounds, mylon, mylat, buf.data)
plt.xlim(-190.,190.)
plt.ylim(-90.,90.)
plt.savefig('fig_partition/B%03d.pdf'%mpi_rank, dpi=1600)