/*

 * Copyright © 2016-2017 Inria.  All rights reserved.

 *

 * $COPYRIGHT$

 *

 * Additional copyrights may follow

 * See COPYING in top-level directory.

 *

 * $HEADER$

 */

#define _GNU_SOURCE         /* See feature_test_macros(7) */

#include <stdio.h>

#include <stdlib.h>

#include <sys/types.h>

#include <dirent.h>

#include <scotch.h>

#include <netloc.h>

#include <netlocscotch.h>

#include <private/netloc.h>

#include <hwloc.h>

static int arch_tree_to_scotch_arch(netloc_arch_tree_t *tree, SCOTCH_Arch *scotch);

static int comm_matrix_to_scotch_graph(double **matrix, int n, SCOTCH_Graph *graph);

static int netlocscotch_get_mapping_from_graph(SCOTCH_Graph *graph,

        netlocscotch_core_t **pcores);

static int compareint(void const *a, void const *b)

{

   const int *int_a = (const int *)a;

   const int *int_b = (const int *)b;

   return *int_a-*int_b;

}

static int build_subarch(SCOTCH_Arch *scotch, NETLOC_int num_nodes, NETLOC_int *node_list,

        SCOTCH_Arch *subarch)

{

    int ret;

    /* Hack to avoid problem with unsorted node list in the subarch and scotch

     * FIXME TODO */

    qsort(node_list, num_nodes, sizeof(*node_list), compareint);

    ret = SCOTCH_archSub(subarch, scotch, num_nodes, node_list);

    if (ret != 0) {

        fprintf(stderr, "Error: SCOTCH_archSub failed\n");

    }

    return ret;

}

/* Convert a netloc tree to a scotch tleaf architecture */

int arch_tree_to_scotch_arch(netloc_arch_tree_t *tree, SCOTCH_Arch *scotch)

{

    int ret;

    ret = SCOTCH_archTleaf(scotch, tree->num_levels, tree->degrees, tree->cost);

    if (ret != 0) {

        fprintf(stderr, "Error: SCOTCH_archTleaf failed\n");

        return NETLOC_ERROR;

    }

    return NETLOC_SUCCESS;

}

static int build_subgraph(SCOTCH_Graph *graph, int *vertices, int num_vertices,

        SCOTCH_Graph *nodegraph)

{

    int ret;

    SCOTCH_Num base;       /* Base value               */

    SCOTCH_Num vert;       /* Number of vertices       */

    SCOTCH_Num *verttab;   /* Vertex array [vertnbr+1] */

    SCOTCH_Num *vendtab;   /* Vertex array [vertnbr]   */

    SCOTCH_Num *velotab;   /* Vertex load array        */

    SCOTCH_Num *vlbltab;   /* Vertex label array       */

    SCOTCH_Num edge;       /* Number of edges (arcs)   */

    SCOTCH_Num *edgetab;   /* Edge array [edgenbr]     */

    SCOTCH_Num *edlotab;   /* Edge load array          */

    SCOTCH_graphData(graph, &base, &vert, &verttab, &vendtab, &velotab,

            &vlbltab, &edge, &edgetab, &edlotab);

    int *vertex_is_present = (int *)malloc(vert*sizeof(int));

    for (int v = 0; v < vert; v++) {

        vertex_is_present[v] = -1;

    }

    for (int v = 0; v < num_vertices; v++) {

        vertex_is_present[vertices[v]] = v;

    }

    // TODO handle other cases 

    if (vendtab) {

        for (int i = 0; i < vert; i++) {

            assert(vendtab[i] == verttab[i+1]);

        }

    }

    SCOTCH_Num *new_verttab;   /* Vertex array [vertnbr+1] */

    SCOTCH_Num *new_vendtab;   /* Vertex array [vertnbr]   */

    SCOTCH_Num *new_velotab;   /* Vertex load array        */

    SCOTCH_Num *new_vlbltab;   /* Vertex label array       */

    SCOTCH_Num new_edge;       /* Number of edges (arcs)   */

    SCOTCH_Num *new_edgetab;   /* Edge array [edgenbr]     */

    SCOTCH_Num *new_edlotab;   /* Edge load array          */

    new_verttab = (SCOTCH_Num *)malloc((num_vertices+1)*sizeof(SCOTCH_Num));

    new_vendtab = NULL;

    if (velotab)

        new_velotab = (SCOTCH_Num *)malloc(num_vertices*sizeof(SCOTCH_Num));

    else

        new_velotab = NULL;

    if (vlbltab)

        new_vlbltab = (SCOTCH_Num *)malloc(num_vertices*sizeof(SCOTCH_Num));

    else

        new_vlbltab = NULL;

    new_edgetab = (SCOTCH_Num *)malloc(edge*sizeof(SCOTCH_Num));

    new_edlotab = (SCOTCH_Num *)malloc(edge*sizeof(SCOTCH_Num));

    int edge_idx = 0;

    new_verttab[0] = 0;

    for (int v = 0; v < num_vertices; v++) {

        if (velotab)

            new_velotab[v] = velotab[vertices[v]];

        if (vlbltab)

            new_vlbltab[v] = vlbltab[vertices[v]];

        for (int e = verttab[vertices[v]]; e < verttab[vertices[v]+1]; e++) {

            int dest_vertex = edgetab[e];

            int new_dest = vertex_is_present[dest_vertex];

            if (new_dest != -1) {

                new_edgetab[edge_idx] = new_dest;

                new_edlotab[edge_idx] = edlotab[e];

                edge_idx++;

            }

        }

        new_verttab[v+1] = edge_idx;

    }

    new_edge = edge_idx;

    SCOTCH_Num *old_edgetab = new_edgetab;

    new_edgetab = (SCOTCH_Num *)

        realloc(new_edgetab, new_edge*sizeof(SCOTCH_Num));

    if (!new_edgetab) {

        new_edgetab = old_edgetab;

    }

    SCOTCH_Num *old_edlotab = new_edlotab;

    new_edlotab = (SCOTCH_Num *)

        realloc(new_edlotab, new_edge*sizeof(SCOTCH_Num));

    if (!new_edlotab) {

        new_edlotab = old_edlotab;

    }

    ret = SCOTCH_graphBuild (nodegraph, base, num_vertices,

                new_verttab, new_vendtab, new_velotab, new_vlbltab,

                new_edge, new_edgetab, new_edlotab);

    free(vertex_is_present);

    return ret;

}

static int build_current_arch(SCOTCH_Arch *scotch_arch,

        SCOTCH_Arch *scotch_subarch, netloc_arch_t *arch)

{

    int ret;

    /* First we need to get the topology of the whole machine */

    ret = netloc_arch_build(arch, 1);

    if( NETLOC_SUCCESS != ret ) {

        return ret;

    }

    if (scotch_subarch) {

        /* Set the current nodes and slots in the arch */

        ret = netloc_arch_set_current_resources(arch);

    } else {

        ret = netloc_arch_set_global_resources(arch);

    }

    if( NETLOC_SUCCESS != ret ) {

        return ret;

    }

    SCOTCH_archInit(scotch_arch);

    ret = arch_tree_to_scotch_arch(arch->arch.global_tree, scotch_arch);

    if (NETLOC_SUCCESS != ret) {

        return ret;

    }

    if (scotch_subarch) {

        /* Now we can build the sub architecture */

        SCOTCH_archInit(scotch_subarch);

        ret = build_subarch(scotch_arch, arch->num_current_hosts,

                arch->current_hosts, scotch_subarch);

    }

    return ret;

}

int netlocscotch_build_global_arch(SCOTCH_Arch *arch)

{

    int ret;

    netloc_arch_t *netloc_arch = netloc_arch_construct();

    ret = build_current_arch(arch, NULL, netloc_arch);

    netloc_arch_destruct(netloc_arch);

    return ret;

}

int netlocscotch_build_current_arch(SCOTCH_Arch *arch, SCOTCH_Arch *subarch)

{

    int ret;

    netloc_arch_t *netloc_arch = netloc_arch_construct();

    ret = build_current_arch(arch, subarch, netloc_arch);

    if (ret == NETLOC_SUCCESS)

        netloc_arch_destruct(netloc_arch);

    return ret;

}

int netlocscotch_get_mapping_from_graph(SCOTCH_Graph *graph,

        netlocscotch_core_t **pcores)

{

    int ret;

    SCOTCH_Arch scotch_arch;

    SCOTCH_Arch scotch_subarch;

    netlocscotch_core_t *cores = NULL;

    netloc_arch_t *arch = netloc_arch_construct();

    ret = build_current_arch(&scotch_arch, &scotch_subarch, arch);

    if (NETLOC_SUCCESS != ret) {

        netloc_arch_destruct(arch);

        return ret;

    }

    NETLOC_int graph_size;

    SCOTCH_graphSize(graph, &graph_size, NULL);

    int num_hosts = arch->num_current_hosts;

    SCOTCH_Strat strategy;

    SCOTCH_stratInit(&strategy);

    /* We force Scotch to use all the processes

     * barat is 0.01 as in SCOTCH_STRATDEFAULT */

    SCOTCH_stratGraphMapBuild(&strategy, SCOTCH_STRATQUALITY, graph_size, 0.01);

    /* The ranks are the indices of the nodes in the complete graph */

    NETLOC_int *ranks = (NETLOC_int *)malloc(graph_size*sizeof(NETLOC_int));

    ret = SCOTCH_graphMap(graph, &scotch_subarch, &strategy, ranks);

    SCOTCH_stratExit(&strategy);

    SCOTCH_archExit(&scotch_subarch);

    SCOTCH_archExit(&scotch_arch);

    if (ret != 0) {

        fprintf(stderr, "Error: SCOTCH_graphMap failed\n");

        goto ERROR;

    }

    cores = (netlocscotch_core_t *)

        malloc(graph_size*sizeof(netlocscotch_core_t));

    if (!arch->has_slots) {

        /* We have the mapping but only for the nodes, not inside the nodes */

        UT_array *process_by_node[num_hosts];

        for (int n = 0; n < num_hosts; n++) {

            utarray_new(process_by_node[n], &ut_int_icd);

        }

        /* Find the processes mapped to the nodes */

        for (int p = 0; p < graph_size; p++) {

            int rank = ranks[p];

            if (rank >= num_hosts || rank < 0) {

                ret = NETLOC_ERROR;

                goto ERROR;

            }

            utarray_push_back(process_by_node[rank], &p);

        }

        /* Use the intranode topology */

        for (int n = 0; n < num_hosts; n++) {

            int *process_list = (int *)process_by_node[n]->d;

            int num_processes = utarray_len(process_by_node[n]);

            netloc_arch_node_t *node =

                arch->node_slot_by_idx[arch->current_hosts[n]].node;

            NETLOC_int node_ranks[num_processes];

            /* We need to extract the subgraph with only the vertices mapped to the

             * current node */

            SCOTCH_Graph nodegraph; /* graph with only elements for node n */

            build_subgraph(graph, process_list, num_processes, &nodegraph);

            /* Build the scotch arch of the all node */

            SCOTCH_Arch scotch_nodearch;

            ret = arch_tree_to_scotch_arch(node->slot_tree, &scotch_nodearch);

            if (NETLOC_SUCCESS != ret) {

                goto ERROR;

            }

            /* Restrict the scotch arch to the available cores */

            SCOTCH_Arch scotch_nodesubarch;

            ret = build_subarch(&scotch_nodearch, node->num_current_slots,

                    node->current_slots, &scotch_nodesubarch);

            if (NETLOC_SUCCESS != ret) {

                goto ERROR;

            }

            /* Find the mapping to the cores */

            ret = SCOTCH_graphMap(&nodegraph, &scotch_nodesubarch, &strategy, node_ranks);

            if (ret != 0) {

                fprintf(stderr, "Error: SCOTCH_graphMap failed\n");

                goto ERROR;

            }

            /* Report the node ranks in the global rank array */

            for (int p = 0; p < num_processes; p++) {

                int process = process_list[p];

                int arch_idx = node->current_slots[node_ranks[p]];

                cores[process].core = node->slot_os_idx[arch_idx];

                cores[process].nodename = strdup(node->node->hostname);

                cores[process].rank = node->slot_ranks[node_ranks[p]];

            }

        }

        for (int n = 0; n < num_hosts; n++) {

            utarray_free(process_by_node[n]);

        }

    } else {

        for (int p = 0; p < graph_size; p++) {

            int host_idx = arch->current_hosts[ranks[p]];

            netloc_arch_node_t *node = arch->node_slot_by_idx[host_idx].node;

            int slot_rank = arch->node_slot_by_idx[host_idx].slot;

            cores[p].nodename = strdup(node->node->hostname);

            cores[p].core = node->slot_os_idx[node->slot_idx[slot_rank]];

            cores[p].rank = node->slot_ranks[node->slot_idx[slot_rank]];

        }

    }

    *pcores = cores;

ERROR:

    free(ranks);

    netloc_arch_destruct(arch);

    if (ret == NETLOC_SUCCESS)

        return ret;

    free(cores);

    return ret;

}

int netlocscotch_get_mapping_from_comm_matrix(double **comm, int num_vertices,

        netlocscotch_core_t **pcores)

{

    int ret;

    SCOTCH_Graph graph;

    ret = comm_matrix_to_scotch_graph(comm, num_vertices, &graph);

    if (NETLOC_SUCCESS != ret) {

        return ret;

    }

    ret = netlocscotch_get_mapping_from_graph(&graph, pcores);

    /* Free arrays */

    {

        SCOTCH_Num base;       /* Base value               */

        SCOTCH_Num vert;       /* Number of vertices       */

        SCOTCH_Num *verttab;   /* Vertex array [vertnbr+1] */

        SCOTCH_Num *vendtab;   /* Vertex array [vertnbr]   */

        SCOTCH_Num *velotab;   /* Vertex load array        */

        SCOTCH_Num *vlbltab;   /* Vertex label array       */

        SCOTCH_Num edge;       /* Number of edges (arcs)   */

        SCOTCH_Num *edgetab;   /* Edge array [edgenbr]     */

        SCOTCH_Num *edlotab;   /* Edge load array          */

        SCOTCH_graphData(&graph, &base, &vert, &verttab, &vendtab, &velotab,

                &vlbltab, &edge, &edgetab, &edlotab);

        free(edlotab);

        free(edgetab);

        free(verttab);

        SCOTCH_graphExit(&graph);

    }

    return ret;

}

int netlocscotch_get_mapping_from_comm_file(char *filename, int *pnum_processes,

        netlocscotch_core_t **pcores)

{

    int ret;

    int n;

    double **mat;

    ret = netloc_build_comm_mat(filename, &n, &mat;;

    if (ret != NETLOC_SUCCESS) {

        return ret;

    }

    *pnum_processes = n;

    ret = netlocscotch_get_mapping_from_comm_matrix(mat, n, pcores);

    free(mat[0]);

    free(mat);

    return ret;

}

static int comm_matrix_to_scotch_graph(double **matrix, int n, SCOTCH_Graph *graph)

{

    int ret;

    SCOTCH_Num base;       /* Base value               */

    SCOTCH_Num vert;       /* Number of vertices       */

    SCOTCH_Num *verttab;   /* Vertex array [vertnbr+1] */

    SCOTCH_Num *vendtab;   /* Vertex array [vertnbr]   */

    SCOTCH_Num *velotab;   /* Vertex load array        */

    SCOTCH_Num *vlbltab;   /* Vertex label array       */

    SCOTCH_Num edge;       /* Number of edges (arcs)   */

    SCOTCH_Num *edgetab;   /* Edge array [edgenbr]     */

    SCOTCH_Num *edlotab;   /* Edge load array          */

    base = 0;

    vert = n;

    verttab = (SCOTCH_Num *)malloc((vert+1)*sizeof(SCOTCH_Num));

    for (int v = 0; v < vert+1; v++) {

        verttab[v] = v*(n-1);

    }

    vendtab = NULL;

    velotab = NULL;

    vlbltab = NULL;

    edge = n*(n-1);

    /* Compute the lowest load to reduce of the values of the load to avoid overflow */

    double min_load = -1;

    for (int v1 = 0; v1 < vert; v1++) {

        for (int v2 = 0; v2 < vert; v2++) {

            double load = matrix[v1][v2];

            if (load >= 0.01 && (load < min_load || min_load < 0)) /* TODO set an epsilon */

                min_load = load;

        }

    }

    edgetab = (SCOTCH_Num *)malloc(n*(n-1)*sizeof(SCOTCH_Num));

    edlotab = (SCOTCH_Num *)malloc(n*(n-1)*sizeof(SCOTCH_Num));

    for (int v1 = 0; v1 < vert; v1++) {

        for (int v2 = 0; v2 < vert; v2++) {

            if (v2 == v1)

                continue;

            int idx = v1*(n-1)+((v2 < v1) ? v2: v2-1);

            edgetab[idx] = v2;

            edlotab[idx] = (int)(matrix[v1][v2]/min_load);

        }

    }

    ret = SCOTCH_graphBuild(graph, base, vert,

            verttab, vendtab, velotab, vlbltab, edge, edgetab, edlotab);

    return ret;

}