/*
 *   -- micMAGMA (version 0.1) --
 *      Univ. of Tennessee, Knoxville
 *      Univ. of California, Berkeley
 *      Univ. of Colorado, Denver
 *      April 2012
 *
 * @author Mark Gates
 */

// defines stub functions, if no GPU implementation is available.
#include <stdio.h>
#include <stdlib.h>
#include <map>
#include <iostream>
#include <mkl.h>

#include "magma.h"
#include "magma_mic_bindings.h"


// ========================================
// initialization
magma_err_t
magma_init( void )
{
    uint16_t hostid, *self;
    int tries, i, num_nodes = -1;
    hostid = 0;
    self = (unsigned short *)&hostid;
    num_nodes = scif_get_nodeIDs(NULL, 0, self);

    //number of mic card in the system
    gNumEpd = num_nodes; //scif_get_nodeIDs return number of scif nodes including host    

    gEpd_array = (scif_epd_t *)malloc(sizeof(scif_epd_t) * (num_nodes));

    // Use node 1 by default
    int node = 1;

    int req_pn = 2049, con_pn;
    struct scif_portID portID;

    for(i = 1; i < num_nodes; i++) {
        tries = 20;
        if ((gEpd_array[i] = scif_open()) < 0) {
            printf("scif_open failed with error %d\n", errno);
            return errno;
        }

        if ((con_pn = scif_bind(gEpd_array[i], req_pn - i - 1)) < 0) {
            printf("scif_bind failed with error %d\n", errno);
            return errno;
        }

        portID.node = i;
        portID.port = 2050;

        retry:
        if ((scif_connect(gEpd_array[i], &portID)) < 0) {
            if ((ECONNREFUSED == errno) && (tries > 0)) {
                printf("scif_connect failed with error %d retrying\n", errno);
                tries--;
                sleep(1);
                goto retry;
            }
            printf("scif_connect failed for mic-%d with error %d\n", i, errno);
            exit(1);
        }

        //printf("connect to mic-%d success\n", i);
    }

    gCurrID = 1;
    gEpd = gEpd_array[gCurrID];

    //resize vector for data structure
    gHost_nbyte.resize(num_nodes);
    gHost_offset.resize(num_nodes);
    gDev_nbyte.resize(num_nodes);
    gDev_offset.resize(num_nodes);

    return MAGMA_SUCCESS;
}

// --------------------
magma_err_t
magma_finalize( void )
{
    int i, err = MAGMA_SUCCESS;
    int control_msg = CLOSE_MIC;

    for(i = 1; i < gNumEpd; i++) {
        gEpd = gEpd_array[i];
        if ((err = scif_send(gEpd, &control_msg, sizeof(control_msg), 1)) <= 0) {
            err = errno;
            printf("scif_send failed with err %d\n", errno);
            fflush(stdout);
        }

        close(gEpd);
    }

    free(gEpd_array);

    return err;
}


// ========================================
// memory allocation on the MIC
magma_err_t
magma_malloc( magma_ptr* ptrPtr, size_t nbyte )
{
    int err = 0;
    int control_msg = ALLOCATE_MEM;
    
    // make nbyte multiple of PAGE_SIZE
    if (nbyte % PAGE_SIZE != 0) {
        nbyte = (size_t)((nbyte / PAGE_SIZE) + 1) * PAGE_SIZE;
    }
    
    if (scif_send(gEpd, &control_msg, sizeof(control_msg), 1) <= 0) {
        err = errno;
        printf("magma_malloc: scif_send failed with err %d\n", errno);
        fflush(stdout);
    }
    
    magma_mic_mem_info mmm_info;
    mmm_info.nbyte = nbyte;
    if (scif_send(gEpd, &mmm_info, sizeof(mmm_info), 1) <= 0) {
        err = errno;
        printf("magma_malloc: scif_send failed with err %d\n", errno);
        fflush(stdout);
    }

    if (scif_recv(gEpd, &mmm_info, sizeof(mmm_info), 1) <= 0) {
        err = errno;
        printf("magma_malloc: scif_recv failed with err %d\n", errno);
        fflush(stdout);
    }
    
    (*ptrPtr) = (void *)(mmm_info.address);
    
    //save device pointer address, offset and number of bytes
    if((*ptrPtr) != 0) {
        (gDev_nbyte[gCurrID])[(off_t)(*ptrPtr)] = nbyte;
        (gDev_offset[gCurrID])[(off_t)(*ptrPtr)] = mmm_info.offset;

        //printf("mic_malloc:address->0x%lx offset: 0x%lx byte: %d\n", (off_t)(*ptr), gDev_offset[(off_t)(*ptr)], gDev_nbyte[(off_t)(*ptr)]);
        return MAGMA_SUCCESS;
    }

    return MIC_FAILURE;
}

// --------------------
magma_err_t
magma_free( magma_ptr ptr )
{
    int err = 0;
    int control_msg = FREE_MEM;
    magma_mic_mem_info mmm_info;
    
    //send message for unregister
    if (scif_send(gEpd, &control_msg, sizeof(control_msg), 1) <= 0) {
        err = errno;
        printf("scif_send failed with err %d\n", errno);
        fflush(stdout);
    }
    
    //send unregister information
    mmm_info.address = (off_t)ptr;
    mmm_info.offset  = (gDev_offset[gCurrID])[(off_t)ptr];
    mmm_info.nbyte   = (gDev_nbyte[gCurrID]) [(off_t)ptr];
    
    //printf("mic_free:address->0x%lx offset: 0x%lx byte: %d\n", (off_t)(ptr), gDev_offset[(off_t)(ptr)], gDev_nbyte[(off_t)(ptr)]);
    if (scif_send(gEpd, &mmm_info, sizeof(mmm_info), 1) <= 0) {
        err = errno;
        printf("scif_send failed with err %d\n", errno);
        fflush(stdout);
    }
    if (scif_recv(gEpd, &mmm_info, sizeof(mmm_info), 1) <= 0) {
        err = errno;
        printf("magma_free: scif_recv failed with err %d\n", errno);
        fflush(stdout);
    }
    
    //remove from the database
    std::map<off_t, size_t>::iterator devnbyte_it;
    std::map<off_t, off_t>::iterator devoff_it;
    
    devnbyte_it = (gDev_nbyte[gCurrID]).find((off_t)ptr);
    (gDev_nbyte[gCurrID]).erase(devnbyte_it);
    devoff_it = (gDev_offset[gCurrID]).find((off_t)ptr);
    (gDev_offset[gCurrID]).erase(devoff_it);
    
    return MAGMA_SUCCESS;
}

// --------------------
// memory allocation on the host CPU
magma_err_t
magma_malloc_pinned( void** ptrPtr, size_t size )
{
    int i, tmpCurrID, err = 0;
    off_t offset, suggested_offset;
    off_t address;
    scif_epd_t tmpEpd;

    //make size multiple of PAGE_SIZE                                                                             
    if (size % (2 * PAGE_SIZE) != 0) {
        size = ((size / (2 * PAGE_SIZE)) + 1) * (2 * PAGE_SIZE);
    }

    posix_memalign(ptrPtr, PAGE_SIZE, size);
    if((*ptrPtr) == NULL)
        return MIC_FAILURE;

    suggested_offset = (off_t)(*ptrPtr);//0x0;
    address = (off_t)*ptrPtr;

    tmpEpd = gEpd;
    tmpCurrID = gCurrID;

    //This might have problem, make sure it is correct later
    for(i = 1; i < gNumEpd; i++) {
        gCurrID = i;
        gEpd = gEpd_array[i];
        if((offset = scif_register(gEpd, *ptrPtr, size, suggested_offset,
                        SCIF_PROT_READ | SCIF_PROT_WRITE, 0)) < 0)
        {
            err = errno;
            printf("scif_register failed with err %d\n", errno);
            fflush(stdout);
            return MIC_FAILURE;
        }
        //save offset, base address size                                                                          
        (gHost_nbyte[gCurrID])[(off_t)(*ptrPtr)] = size;
        (gHost_offset[gCurrID])[(off_t)(*ptrPtr)] = offset;
    }

    gEpd = tmpEpd;
    gCurrID = tmpCurrID;

    //printf("Host side: address: 0x%lx offset: 0x%lx\n", address, offset);                                   
    suggested_offset += offset;

    //save offset, base address size                                                                          
    //gHost_nbyte[(off_t)(*ptrPtr)] = size;
    //gHost_offset[(off_t)(*ptrPtr)] = offset;

    return MAGMA_SUCCESS;
}

// --------------------
// this actually frees memory.
magma_err_t
magma_free_pinned( void* ptr )
{
    int i, tmpCurrID, err;
    off_t nbyte, offset;
    scif_epd_t tmpEpd;

    std::map<off_t, size_t  >::iterator hnbyte_it;
    std::map<off_t, off_t>::iterator hoff_it;
    
    tmpEpd = gEpd;
    tmpCurrID = gCurrID;

    //This might have problem, make sure it is correct later
    for(i = 1; i < gNumEpd; i++) {
        gCurrID = i;
        gEpd = gEpd_array[i];
        nbyte  = (gHost_nbyte[gCurrID])[(off_t)ptr];
        offset = (gHost_offset[gCurrID])[(off_t)ptr];

        //unregister the allocated memory                                                                         
        if ((err = scif_unregister(gEpd, offset, nbyte)) < 0) {
            err = errno;
            printf("scif_unregister failed with err %d\n", errno);
            fflush(stdout);
            return MIC_FAILURE;
        }

        //remove from the database                                                                                
        hnbyte_it = (gHost_nbyte[gCurrID]).find((off_t)ptr);
        (gHost_nbyte[gCurrID]).erase(hnbyte_it);
        hoff_it = (gHost_offset[gCurrID]).find((off_t)ptr);
        (gHost_offset[gCurrID]).erase(hoff_it);

    }

    gEpd = tmpEpd;
    gCurrID = tmpCurrID;

    free(ptr);
    return MAGMA_SUCCESS;
}

// ========================================
// device & queue support

magma_err_t
magma_setdevice( magma_int_t id )
{
    gCurrID = id + 1;
    gEpd = gEpd_array[gCurrID];
    return MAGMA_SUCCESS;
}

magma_err_t
magma_get_devices(
    magma_device_t* devices,
    magma_int_t     size,
    magma_int_t*    numPtr )
{
    *numPtr = 1;
    //printf( "%s( ptr, %d, numPtr )\n", __func__, size );
    return MAGMA_SUCCESS;
}

magma_err_t
magma_queue_create( magma_device_t device, magma_queue_t* queuePtr )
{
    //printf( "%s( %d, %d )\n", __func__, device, *queuePtr );
    return MAGMA_SUCCESS;
}

magma_err_t
magma_queue_destroy( magma_queue_t  queue )
{
    //printf( "%s( %d )\n", __func__, queue );
    return MAGMA_SUCCESS;
}

// waits for all kernels and all communication to complete.
magma_err_t
magma_queue_sync( magma_queue_t queue )
{
    magma_mic_synchronize_kernels();
    magma_event_t event;
    magma_mic_fence_mark( &event );
    magma_mic_fence_wait(  event );
    return MAGMA_SUCCESS;
}

// waits for all kernels to complete.
magma_err_t
magma_kernel_sync()
{
    magma_mic_synchronize_kernels();
    return MAGMA_SUCCESS;
}


// ========================================
// event support
magma_err_t
magma_event_create( magma_event_t* eventPtr )
{
    // nothing to do.
    return MAGMA_SUCCESS;
}

magma_err_t
magma_event_destroy( magma_event_t event )
{
    // nothing to do.
    return MAGMA_SUCCESS;
}

magma_err_t
magma_event_record( magma_event_t event, magma_queue_t queue )
{
    // not yet implemented
    printf( "%s( %d )\n", __func__, event );
    return MAGMA_SUCCESS;
}

magma_err_t
magma_event_query( magma_event_t event )
{
    // not yet implemented
    printf( "%s( %d )\n", __func__, event );
    return MAGMA_SUCCESS;
}

magma_err_t
magma_event_sync( magma_event_t event )
{
    magma_mic_fence_wait( event );
    return MAGMA_SUCCESS;
}
