#########################################################
#
#
#  Read in a data file and subdivide the triangle list
#
#
#  The final routine, pmesh_divide_metis, does automatic
# grid partitioning. Once testing has finished on this
# routine the others should be removed.
#
#  Authors: Linda Stals and Matthew Hardy, June 2005
#  Modified: Linda Stals, Nov 2005
#            Jack Kelly, Nov 2005
#
#
#########################################################

from os import sep
from sys import path
from math import floor

from Numeric import zeros, Float, Int, reshape, argsort


#########################################################
#
# If the triangles list is reordered, the quantities
# assigned to the triangles must also be reorded.
#
# *) quantities contain the quantites in the old ordering
# *) proc_sum[i] contains the number of triangles in
# processor i
# *) tri_index is a map from the old triangle ordering to
# the new ordering, where the new number for triangle
# i is proc_sum[tri_index[i][0]]+tri_index[i][1]
#
# -------------------------------------------------------
#
# *) The quantaties are returned in the new ordering
#
#########################################################

def reorder(quantities, tri_index, proc_sum):

    # Find the number triangles

    N = len(tri_index)

    # Temporary storage area

    index = zeros(N, Int)
    q_reord = {}

    # Find the new ordering of the triangles

    for i in range(N):
        bin = tri_index[i][0]
        bin_off_set = tri_index[i][1]
        index[i] = proc_sum[bin]+bin_off_set

    # Reorder each quantity according to the new ordering

    for k in quantities:
        q_reord[k] = zeros((N, 3), Float)
        for i in range(N):
            q_reord[k][index[i]]=quantities[k].vertex_values[i]
    del index

    return q_reord


#########################################################
#
# Do a simple coordinate based division of the grid
#
# *) n_x is the number of subdivisions in the x direction
# *) n_y is the number of subdivisions in the y direction
#
# -------------------------------------------------------
#
# *)  The nodes, triangles, boundary, and quantities are
# returned. triangles_per_proc defines the subdivision.
# The first triangles_per_proc[0] triangles are assigned
# to processor 0, the next triangles_per_proc[1] are
# assigned to processor 1 etc. The boundary and quantites
# are ordered the same way as the triangles
#
#########################################################

def pmesh_divide(domain, n_x = 1, n_y = 1):

    # Find the bounding box
    
    x_coord_min = domain.xy_extent[0]
    x_coord_max = domain.xy_extent[2]
    y_coord_min = domain.xy_extent[1]
    y_coord_max = domain.xy_extent[3]

    # Find the size of each sub-box

    x_div = (x_coord_max-x_coord_min)/n_x
    y_div = (y_coord_max-y_coord_min)/n_y

    # Initialise the lists
    
    tri_list = []
    triangles_per_proc = []
    proc_sum = []
    for i in range(n_x*n_y):
        tri_list.append([])
        triangles_per_proc.append([])
        proc_sum.append([])
        tri_list[i] = []

    # Subdivide the triangles depending on which sub-box they sit
    # in (a triangle sits in the sub-box if its first vectex sits
    # in that sub-box)

    tri_index = {}
    N = domain.number_of_elements
    for i in range(N):
        t = domain.triangles[i]

        x_coord = domain.centroid_coordinates[i][0]
        bin_x = int(floor((x_coord-x_coord_min)/x_div))
        if (bin_x == n_x):
            bin_x = n_x-1

        y_coord = domain.centroid_coordinates[i][1]
        bin_y = int(floor((y_coord-y_coord_min)/y_div))
        if (bin_y == n_y):
            bin_y = n_y-1

        bin = bin_y*n_x + bin_x
        tri_list[bin].append(t)
        tri_index[i] = ([bin, len(tri_list[bin])-1])

    # Find the number of triangles per processor and order the
    # triangle list so that all of the triangles belonging to
    # processor i are listed before those belonging to processor
    # i+1

    triangles = []
    for i in range(n_x*n_y):
        triangles_per_proc[i] = len(tri_list[i])
        for t in tri_list[i]:
            triangles.append(t)

    # The boundary labels have to changed in accoradance with the
    # new triangle ordering, proc_sum and tri_index help with this

    proc_sum[0] = 0
    for i in range(n_x*n_y-1):
        proc_sum[i+1]=proc_sum[i]+triangles_per_proc[i]

    # Relabel the boundary elements to fit in with the new triangle
    # ordering

    boundary = {}
    for b in domain.boundary:
        t =  tri_index[b[0]]
        boundary[proc_sum[t[0]]+t[1], b[1]] = domain.boundary[b]

    quantities = reorder(domain.quantities, tri_index, proc_sum)

    # Extract the node list
    
    nodes = domain.coordinates.copy()
    ttriangles = zeros((len(triangles), 3), Int)
    for i in range(len(triangles)):
        ttriangles[i] = triangles[i]
        
    return nodes, ttriangles, boundary, triangles_per_proc, quantities


#########################################################
#
# Do a coordinate based division of the grid using the
# centroid coordinate
#
# *) n_x is the number of subdivisions in the x direction
# *) n_y is the number of subdivisions in the y direction
#
# -------------------------------------------------------
#
# *)  The nodes, triangles, boundary, and quantities are
# returned. triangles_per_proc defines the subdivision.
# The first triangles_per_proc[0] triangles are assigned
# to processor 0, the next triangles_per_proc[1] are
# assigned to processor 1 etc. The boundary and quantites
# are ordered the same way as the triangles
#
#########################################################
def pmesh_divide_steve(domain, n_x = 1, n_y = 1):
    
    # Find the bounding box
    x_coord_min = domain.xy_extent[0]
    x_coord_max = domain.xy_extent[2]
    y_coord_min = domain.xy_extent[1]
    y_coord_max = domain.xy_extent[3]


    # Find the size of each sub-box

    x_div = (x_coord_max-x_coord_min)/n_x
    y_div = (y_coord_max-y_coord_min)/n_y


    # Initialise the lists
    
    tri_list = []
    triangles_per_proc = []
    proc_sum = []
    for i in range(n_x*n_y):
        tri_list.append([])
        triangles_per_proc.append([])
        proc_sum.append([])
        tri_list[i] = []

    # Subdivide the triangles depending on which sub-box they sit
    # in (a triangle sits in the sub-box if its first vectex sits
    # in that sub-box)

    tri_index = {}
    N = domain.number_of_elements
 
    # Sort by x coordinate of centroid

    sort_order = argsort(argsort(domain.centroid_coordinates[:,0]))

    x_div = float(N)/n_x

    for i in range(N):
        t = domain.triangles[i]
        
        bin = int(floor(sort_order[i]/x_div))
        if (bin == n_x):
            bin = n_x-1

        tri_list[bin].append(t)
        tri_index[i] = ([bin, len(tri_list[bin])-1])

    # Find the number of triangles per processor and order the
    # triangle list so that all of the triangles belonging to
    # processor i are listed before those belonging to processor
    # i+1

    triangles = []
    for i in range(n_x*n_y):
        triangles_per_proc[i] = len(tri_list[i])
        for t in tri_list[i]:
           triangles.append(t)
        

    # The boundary labels have to changed in accoradance with the
    # new triangle ordering, proc_sum and tri_index help with this

    proc_sum[0] = 0
    for i in range(n_x*n_y-1):
        proc_sum[i+1]=proc_sum[i]+triangles_per_proc[i]

    # Relabel the boundary elements to fit in with the new triangle
    # ordering

    boundary = {}
    for b in domain.boundary:
        t =  tri_index[b[0]]
        boundary[proc_sum[t[0]]+t[1], b[1]] = domain.boundary[b]

    quantities = reorder(domain.quantities, tri_index, proc_sum)

    # Extract the node list
    
    nodes = domain.coordinates.copy()

    # Convert the triangle datastructure to be an array type,
    # this helps with the communication

    ttriangles = zeros((len(triangles), 3), Int)
    for i in range(len(triangles)):
        ttriangles[i] = triangles[i]
        
    return nodes, ttriangles, boundary, triangles_per_proc, quantities




###
#
# Divide the mesh using a call to metis, through pymetis.


path.append('..' + sep + 'pymetis')

from metis import partMeshNodal

def pmesh_divide_metis(domain, n_procs):
    
    # Initialise the lists
    # List, indexed by processor of # triangles.
    
    triangles_per_proc = []
    
    # List of lists, indexed by processor of vertex numbers
    
    tri_list = []
    
    # List indexed by processor of cumulative total of triangles allocated.
    
    proc_sum = []
    for i in range(n_procs):
        tri_list.append([])
        triangles_per_proc.append(0)
        proc_sum.append([])

    # Prepare variables for the metis call
    
    n_tri = len(domain.triangles)
    if n_procs != 1: #Because metis chokes on it...
        n_vert = len(domain.coordinates)
        t_list = domain.triangles.copy()
        t_list = reshape(t_list, (-1,))
    
        # The 1 here is for triangular mesh elements.
        
        edgecut, epart, npart = partMeshNodal(n_tri, n_vert, t_list, 1, n_procs)
        # print edgecut
        # print npart
        # print epart
        del edgecut
        del npart
        # Assign triangles to processes, according to what metis told us.
        
        # tri_index maps triangle number -> processor, new triangle number
        # (local to the processor)
        
        tri_index = {}
        triangles = []        
        for i in range(n_tri):
            triangles_per_proc[epart[i]] = triangles_per_proc[epart[i]] + 1
            tri_list[epart[i]].append(domain.triangles[i])
            tri_index[i] = ([epart[i], len(tri_list[epart[i]]) - 1])

        # print tri_list
        # print triangles_per_proc
        
        # Order the triangle list so that all of the triangles belonging
        # to processor i are listed before those belonging to processor
        # i+1

        for i in range(n_procs):
            for t in tri_list[i]:
                triangles.append(t)
            
        # The boundary labels have to changed in accoradance with the
        # new triangle ordering, proc_sum and tri_index help with this

        proc_sum[0] = 0
        for i in range(n_procs - 1):
            proc_sum[i+1]=proc_sum[i]+triangles_per_proc[i]

        # Relabel the boundary elements to fit in with the new triangle
        # ordering

        boundary = {}
        for b in domain.boundary:
            t =  tri_index[b[0]]
            boundary[proc_sum[t[0]]+t[1], b[1]] = domain.boundary[b]

        quantities = reorder(domain.quantities, tri_index, proc_sum)
    else:
        boundary = domain.boundary.copy()
        triangles_per_proc[0] = n_tri
        triangles = domain.triangles.copy()
        
        # This is essentially the same as a chunk of code from reorder.
        
        quantities = {}
        for k in domain.quantities:
            quantities[k] = zeros((n_tri, 3), Float)
            for i in range(n_tri):
                quantities[k][i] = domain.quantities[k].vertex_values[i]
        
    # Extract the node list
    
    nodes = domain.coordinates.copy()
    
    # Convert the triangle datastructure to be an array type,
    # this helps with the communication

    ttriangles = zeros((len(triangles), 3), Int)
    for i in range(len(triangles)):
        ttriangles[i] = triangles[i]
        
    return nodes, ttriangles, boundary, triangles_per_proc, quantities