#########################################################
#
#  Given the subdivision of the grid assigned to the
# current processor convert it into a form that is
# appropriate for the GA datastructure.
#
#  The main function of these modules is to change the
# node numbering. The GA datastructure assumes they
# are numbered consecutively from 0.
#
#  The module also changes the communication pattern
# datastructure into a form needed by parallel_advection
#
#  Authors: Linda Stals and Matthew Hardy, June 2005
#
#
#########################################################

from mesh import *
from Numeric import array, Int, Float, zeros

#########################################################
# Convert the format of the data to that used by
# pyvolution
#
#
# *) Change the nodes global ID's to an integer value,
#starting from 0.
#
# *) The triangles and boundary edges must also be
# updated accordingly.
#
# -------------------------------------------------------
#
# *) The nodes, triangles and boundary edges defined by
# the new numbering scheme are returned
#
#########################################################

def build_local_GA(nodes, triangles):

    Nnodes =len(nodes)
    Ntriangles = len(triangles)

    index={} #a dictionary mapping existing node ids to the new ids 0,1,2,...

    GAnodes = []

    nodes  = nodes # do not rearrange nodes in origianal grid
    for node_idnew in range(Nnodes): #move through the list of nodes, renumbering to 0,1,2, ...
        index[nodes[node_idnew][0]]=node_idnew #make the dictionary entry, for later use by triangles
        GAnodes.append(nodes[node_idnew][1:3])
        nodes[node_idnew][0]=node_idnew #renumber the node

    #Now loop over the triangles, changing the node ids and orientation.
    for t in range(Ntriangles):
        for i in range(3):
            n=index[triangles[t][i]] #use the dictionary to get the new node id
            triangles[t][i]=n #relabel the node

    del (index)

    return GAnodes


#########################################################
# Change the communication format to that needed by the
# parallel advection file.
#
# *) The index contains [global triangle no,
# local triangle no.]
#
# -------------------------------------------------------
#
# *) The ghost_recv and full_send dictionaries are
# returned.
#
# *) ghost_recv dictionary is local id, global id, value
#
# *) full_recv dictionary is local id, global id, value
#
# *) The information is ordered by the global id so the
# local and global id do not need to be compared when the
# information is sent/received because the order is
# predetermined.
#
#########################################################

def sort_tup(x, y):
    return cmp(x[0], y[0])

def build_local_commun(index, ghostc, fullc, nproc):

    # initialise

    full_send = {}
    ghost_recv = {}

    # build the ghost_recv dictionary (sort the
    # information by the global numbering)

    ghostc.sort(sort_tup)
    for c in ghostc:
        if not ghost_recv.has_key(c[1]):
            ghost_recv[c[1]] = [0, 0]
            ghost_recv[c[1]][0] = []
            ghost_recv[c[1]][1] = []
        ghost_recv[c[1]][0].append(index[c[0]])
        ghost_recv[c[1]][1].append(c[0])


    # build a temporary copy of the full_send dictionary
    # (this version allows the information to be stored
    # by the global numbering)

    tmp_send = {}
    for global_id in fullc:
        for i in range(len(fullc[global_id])):
            neigh = fullc[global_id][i]
            if not tmp_send.has_key(neigh):
                tmp_send[neigh] = []
            tmp_send[neigh].append([global_id, index[global_id]])

    # extract the full send information and put it in the form
    # required for the full_send dictionary

    for neigh in tmp_send:
        neigh_commun = tmp_send[neigh]
        neigh_commun.sort(sort_tup)
        full_send[neigh] = [0, 0]
        full_send[neigh][0] = []
        full_send[neigh][1] = []
        for t in neigh_commun:
            full_send[neigh][0].append(t[1])
            full_send[neigh][1].append(t[0])



    # parallel advection expects numeric arrays
    for key in full_send:
        full_send[key][0] = array(full_send[key][0],Int)
        full_send[key][1] = array(full_send[key][1],Int)


    for key in ghost_recv:
        ghost_recv[key][0] = array(ghost_recv[key][0],Int)
        ghost_recv[key][1] = array(ghost_recv[key][1],Int)



    return ghost_recv, full_send


#########################################################
# Convert the format of the data to that used by
# pyvolution
#
#
# *) Change the nodes global ID's to an integer value,
# starting from 0. The node numbering in the triangles
# must also be updated to take this into accound.
#
# *) The triangle number will also change, which affects
# the boundary tag information and the communication
# pattern.
#
# -------------------------------------------------------
#
# *) The nodes, triangles, boundary edges and communication
# pattern defined by the new numbering scheme are returned
#
#########################################################

def build_local_mesh(submesh, lower_t, upper_t, nproc):

    # combine the full nodes and ghost nodes

    nodes = submesh["full_nodes"]
    nodes.extend(submesh["ghost_nodes"])

    # combine the full triangles and ghost triangles

    triangles = submesh["full_triangles"]
    triangles.extend(map(lambda t: t[1], submesh["ghost_triangles"]))

    # renumber the boundary edges to correspond to the new
    # triangle numbering

    GAboundary = {}
    for b in submesh["full_boundary"]:
        GAboundary[b[0]-lower_t,b[1]]=submesh["full_boundary"][b]

    # make note of the new triangle numbers, including the ghost
    # triangles

    index = {}
    for i in range(lower_t, upper_t):
        index[i] = i-lower_t
    for i in range(len(submesh["ghost_triangles"])):
        index[submesh["ghost_triangles"][i][0]] = i+upper_t-lower_t

    # change the node numbering (and update the numbering in the
    # triangles)

    GAnodes = build_local_GA(nodes, triangles)

    # extract the local quantities
    
    quantities ={}
    for k in submesh["full_quan"]:
        Nf = len(submesh["full_quan"][k])
        Ng = len(submesh["ghost_quan"][k])
        quantities[k] = zeros((Nf+Ng, 3), Float)
        quantities[k][0:Nf] = submesh["full_quan"][k] 
        quantities[k][Nf:Nf+Ng] = submesh["ghost_quan"][k]
                             
    # change the communication pattern into a form needed by
    # the parallel_advection.py file

    gcommun = submesh["ghost_commun"]
    fcommun = submesh["full_commun"]
    [ghost_rec, full_send] = build_local_commun(index, gcommun, fcommun, nproc)

    # clean up before exiting

    del(index)

    return GAnodes, triangles, GAboundary, quantities, ghost_rec, full_send