/*
 * Copyright (C) by Argonne National Laboratory
 *     See COPYRIGHT in top-level directory
 */

#include "mpidimpl.h"

/* Generates a bitarray based on orig_comm where all procs in group are marked with 1 */
static void group_to_bitarray(MPIR_Group *group, MPIR_Comm *orig_comm, int **bitarray, int *bitarray_size) {
    int mask;
    int *group_ranks, *comm_ranks, i;

    /* Calculate the bitarray size in ints and allocate space */
    *bitarray_size = (orig_comm->local_size / (8 * sizeof(int)) + (orig_comm->local_size % (8 * sizeof(int)) ? 1 : 0));
    *bitarray = (int *) MPL_malloc(sizeof(int) * *bitarray_size, MPL_MEM_OTHER);

    /* If the group is empty, return an empty bitarray. */
    if (group == MPIR_Group_empty) {
        for (i = 0; i < *bitarray_size; i++) (*bitarray)[i] = 0;
        return;
    }

    /* Get the ranks of group in orig_comm */
    group_ranks = (int *) MPL_malloc(sizeof(int) * group->size, MPL_MEM_OTHER);
    comm_ranks = (int *) MPL_malloc(sizeof(int) * group->size, MPL_MEM_OTHER);

    for (i = 0; i < group->size; i++) group_ranks[i] = i;
    for (i = 0; i < *bitarray_size; i++) (*bitarray)[i] = 0;

    MPIR_Group_translate_ranks_impl(group, group->size, group_ranks,
                                    orig_comm->local_group, comm_ranks);

    /* For each process in the group, shift a bit to the correct location and
     * add it to the bitarray. */
    for (i = 0; i < group->size ; i++) {
        if (comm_ranks[i] == MPI_UNDEFINED) continue;
        int idx = comm_ranks[i] / (sizeof(int) * 8);
        mask = 0x1 << comm_ranks[i] % (sizeof(int) * 8);
        (*bitarray)[idx] |= mask;
    }

    MPL_free(group_ranks);
    MPL_free(comm_ranks);
}

/* Generates an MPIR_Group from a bitarray */
static MPIR_Group *bitarray_to_group(MPIR_Comm *comm_ptr, int *bitarray)
{
    MPIR_Group *ret_group;
    MPIR_Group *comm_group;
    UT_array *ranks_array;
    int i, found = 0;

    /* Create a utarray to make storing the ranks easier */
    utarray_new(ranks_array, &ut_int_icd, MPL_MEM_OTHER);

    MPIR_Comm_group_impl(comm_ptr, &comm_group);

    /* Converts the bitarray into a utarray */
    for (i = 0; i < comm_ptr->local_size; i++) {
        if (bitarray[i / (sizeof(int) * 8)] & (0x1 << (i % (sizeof(int) * 8)))) {
            utarray_push_back(ranks_array, &i, MPL_MEM_OTHER);
            found++;
        }
    }

    if (found)
        /* Converts the utarray into a group */
        MPIR_Group_incl_impl(comm_group, found, ut_int_array(ranks_array), &ret_group);
    else
        ret_group = MPIR_Group_empty;

    utarray_free(ranks_array);
    MPIR_Group_release(comm_group);

    return ret_group;
}

int MPID_Comm_get_all_failed_procs(MPIR_Comm *comm_ptr, MPIR_Group **failed_group, int tag)
{
    int mpi_errno = MPI_SUCCESS, ret_errno;
    int i, j, bitarray_size;
    int *bitarray, *remote_bitarray;
    MPIR_Group *local_fail;

    MPIR_FUNC_ENTER;

    /* Kick the progress engine in case it's been a while so we get all the
     * latest updates about failures */
    MPID_Progress_poke();

    /* Generate the list of failed processes */
    MPIDI_CH3U_Check_for_failed_procs();

    mpi_errno = MPIDI_CH3U_Get_failed_group(-2, &local_fail);
    MPIR_ERR_CHECK(mpi_errno);

    /* Generate a bitarray based on the list of failed procs */
    group_to_bitarray(local_fail, comm_ptr, &bitarray, &bitarray_size);
    remote_bitarray = MPL_malloc(sizeof(int) * bitarray_size, MPL_MEM_OTHER);
    if (local_fail != MPIR_Group_empty)
        MPIR_Group_release(local_fail);

    /* For now, this will be implemented as a star with rank 0 serving as
     * the source */
    if (comm_ptr->rank == 0) {
        for (i = 1; i < comm_ptr->local_size; i++) {
            /* Get everyone's list of failed processes to aggregate */
            ret_errno = MPIC_Recv(remote_bitarray, bitarray_size, MPI_INT,
                i, tag, comm_ptr, MPI_STATUS_IGNORE);
            if (ret_errno) continue;

            /* Combine the received bitarray with my own */
            for (j = 0; j < bitarray_size; j++) {
                if (remote_bitarray[j] != 0) {
                    bitarray[j] |= remote_bitarray[j];
                }
            }
        }

        for (i = 1; i < comm_ptr->local_size; i++) {
            /* Send the list to each rank to be processed locally */
            ret_errno = MPIC_Send(bitarray, bitarray_size, MPI_INT, i,
                tag, comm_ptr, MPIR_ERR_NONE);
            if (ret_errno) continue;
        }

        /* Convert the bitarray into a group */
        *failed_group = bitarray_to_group(comm_ptr, bitarray);
    } else {
        /* Send my bitarray to rank 0 */
        mpi_errno = MPIC_Send(bitarray, bitarray_size, MPI_INT, 0,
            tag, comm_ptr, MPIR_ERR_NONE);

        /* Get the resulting bitarray back from rank 0 */
        mpi_errno = MPIC_Recv(remote_bitarray, bitarray_size, MPI_INT, 0,
            tag, comm_ptr, MPI_STATUS_IGNORE);

        /* Convert the bitarray into a group */
        *failed_group = bitarray_to_group(comm_ptr, remote_bitarray);
    }

    MPL_free(bitarray);
    MPL_free(remote_bitarray);

  fn_exit:
    MPIR_FUNC_EXIT;
    return mpi_errno;
  fn_fail:
    goto fn_exit;
}
