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)

import sys
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

#from dynamico import partition
from dynamico import parallel
from dynamico import unstructured as unst
from dynamico.unstructured import list_stencil, reindex

print 'Done loading modules'

sys.stdout.flush()

#----------------- 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 patches(degree, bounds, lon, lat):
    for i in range(degree.size):
        nb_edge=degree[i]
        bounds_cell = bounds[i,0:nb_edge]
        lat_cell    = lat[bounds_cell]
        lon_cell    = lon[bounds_cell]
        orig=lon_cell[0]
        lon_cell    = lon_cell-orig+180.
        lon_cell    = np.mod(lon_cell,360.)
        lon_cell    = lon_cell+orig-180.
#        if np.abs(lon_cell-orig).max()>10. :
#            print '%d patches :'%mpi_rank, lon_cell
        lonlat_cell = np.zeros((nb_edge,2))
        lonlat_cell[:,0],lonlat_cell[:,1] = lon_cell,lat_cell
        polygon = Polygon(lonlat_cell, True)
        yield polygon

def plot_mesh(ax, clim, degree, bounds, lon, lat, data):
    nb_vertex = lon.size # global
    p = list(patches(degree, bounds, lon, lat))
    print '%d : plot_mesh %d %d %d'%( mpi_rank, degree.size, len(p), len(data) ) 
    p = PatchCollection(p, linewidth=0.01)
    p.set_array(data) # set values at each polygon (cell)
    p.set_clim(clim)
    ax.add_collection(p)

def local_mesh(get_mycells):
    mydegree, mybounds = [get_mycells(x) for x in nEdgesOnCell, verticesOnCell]
    print '%d : len(mydegree)=%d'%(mpi_rank, len(mydegree))
    vertex_list = sorted(set(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)
    reindex(vertex_dict, mydegree, mybounds)
    return vertex_list, mydegree, mybounds, mylon, mylat

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

#grid = 'x1.2562'
grid = 'x1.10242'
#grid = 'x4.163842'
print 'Reading MPAS file %s ...'%grid
sys.stdout.flush()

nc = cdf.Dataset('grids/%s.grid.nc'%grid, "r")
dim_cell, dim_edge, dim_vertex = [
    parallel.PDim(nc.dimensions[name], comm) 
    for name in 'nCells','nEdges','nVertices']
edge_degree   = parallel.CstPArray1D(dim_edge, np.int32, 2)
vertex_degree = parallel.CstPArray1D(dim_vertex, np.int32, 3)
nEdgesOnCell, verticesOnCell, edgesOnCell, cellsOnCell, latCell = [
    parallel.PArray(dim_cell, nc.variables[var])
    for var in 'nEdgesOnCell', 'verticesOnCell', 'edgesOnCell', 'cellsOnCell', 'latCell' ]
cellsOnVertex, edgesOnVertex, kiteAreasOnVertex, lonVertex, latVertex = [
    parallel.PArray(dim_vertex, nc.variables[var])
    for var in 'cellsOnVertex', 'edgesOnVertex', 'kiteAreasOnVertex', 'lonVertex', 'latVertex']
nEdgesOnEdge, cellsOnEdge, edgesOnEdge, verticesOnEdge, weightsOnEdge = [
    parallel.PArray(dim_edge, nc.variables[var])
    for var in 'nEdgesOnEdge', 'cellsOnEdge', 'edgesOnEdge', 'verticesOnEdge', 'weightsOnEdge']

# Indices start at 0 on the C/Python side and at 1 on the Fortran/MPAS side
# hence an offset of 1 is added/substracted where needed.
for x in (verticesOnCell, edgesOnCell, cellsOnCell, cellsOnVertex, edgesOnVertex,
          cellsOnEdge, edgesOnEdge, verticesOnEdge) : x.data = x.data-1
edge2cell, cell2edge, edge2vertex, vertex2edge, cell2cell, edge2edge = [
    unst.Stencil_glob(a,b) for a,b in 
    (edge_degree, cellsOnEdge), (nEdgesOnCell, edgesOnCell),
    (edge_degree, verticesOnEdge), (vertex_degree, edgesOnVertex),
    (nEdgesOnCell, cellsOnCell), (nEdgesOnEdge, edgesOnEdge) ]

#---------------- partition edges and cells ------------------#

print 'Partitioning ...'
sys.stdout.flush()

edge_owner = unst.partition_mesh(nEdgesOnEdge, edgesOnEdge, mpi_size)
edge_owner = parallel.LocPArray1D(dim_edge, edge_owner)
cell_owner = unst.partition_from_stencil(edge_owner, nEdgesOnCell, edgesOnCell)
cell_owner = parallel.LocPArray1D(dim_cell, cell_owner)

#--------------------- construct halos  -----------------------#

print 'Constructing halos ...'
sys.stdout.flush()

def chain(start, links):
    for link in links:
        start = link(start).neigh_set
        yield start

edges_E0 = unst.find_my_cells(edge_owner)
cells_C0, edges_E1, vertices_V1, edges_E2, cells_C1 = chain(
    edges_E0, ( edge2cell, cell2edge, edge2vertex, vertex2edge, edge2cell) )

edges_E0, edges_E1, edges_E2 = unst.progressive_list(edges_E0, edges_E1, edges_E2)
cells_C0, cells_C1 = unst.progressive_list(cells_C0, cells_C1)

print 'E2,E1,E0 ; C1,C0 : ', map(len, (edges_E2, edges_E1, edges_E0, cells_C1, cells_C0))
sys.stdout.flush()

#com_edges = parallel.Halo_Xchange(24, dim_edge, edges_E2, dim_edge.get(edges_E2, edge_owner))

mycells, halo_cells = cells_C0, cells_C1
get_mycells, get_halo_cells = dim_cell.getter(mycells), dim_cell.getter(halo_cells)
com_cells = parallel.Halo_Xchange(42, dim_cell, halo_cells, get_halo_cells(cell_owner))

local_num, total_num = np.zeros(1), np.zeros(1)
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
sys.stdout.flush()

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

if True:
    print 'Plotting ...'
    sys.stdout.flush()

    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
    plot_mesh(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()
    plot_mesh(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)