/*
 * adios_transform_mgard_write.c
 *
 * 	Author: Jong Choi
 * 	Contact: choij@ornl.gov
 */

#include <stdint.h>
#include <assert.h>
#include <limits.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>

#include "adios_logger.h"
#include "adios_transforms_common.h"
#include "adios_transforms_write.h"
#include "adios_transforms_hooks_write.h"
#include "adios_transforms_util.h"

#include "core/common_adios.h"

#ifdef HAVE_MGARD
#include "mgard_capi.h"

#ifdef HAVE_ZCHECKER
#  ifndef _NOMPI
#    define USE_ZCHECKER 1
#  else
#    undef USE_ZCHECKER
#  endif
#endif

#ifdef USE_ZCHECKER
#include <ZC_rw.h>
#include <zc.h>
#include "zcheck_comm.h"
#endif

typedef unsigned int uint;

// Variables need to be defined as static variables
static double mgard_tolD = 0.0;
static float mgard_tolF = 0.0;
static int use_zchecker = 0;
static char *zc_configfile = "zc.config";

static int check_file(const char* filename){
    struct stat buffer;
    int exist = stat(filename,&buffer);
    if(exist == 0)
        return 1;
    else // -1
        return 0;
}

uint16_t adios_transform_mgard_get_metadata_size(struct adios_transform_spec *transform_spec)
{
    //log_debug("function: %s\n", __FUNCTION__);
    return 0; // Set amount of transform-internal metadata space to allocate
}

void adios_transform_mgard_transformed_size_growth(const struct adios_var_struct *var, const struct adios_transform_spec *transform_spec,
                                                uint64_t *constant_factor, double *linear_factor, double *capped_linear_factor, uint64_t *capped_linear_cap)
{
    //log_debug("function: %s\n", __FUNCTION__);
}

int adios_transform_mgard_apply(struct adios_file_struct *fd,
                             struct adios_var_struct *var,
                             uint64_t *transformed_len,
                             int use_shared_buffer,
                             int *wrote_to_shared_buffer)
{
    //log_debug("function: %s\n", __FUNCTION__);
    //log_debug("use_shared_buffer: %d\n", use_shared_buffer);

    int iflag = 1; //0 -> float, 1 -> double
    int nrow, ncol;
    int out_size;
    void *tol;

    // Get type info
    switch (var->pre_transform_type)
    {
        case adios_double:
            iflag = 1;
            tol = &mgard_tolD;
            break;
        case adios_real:
            iflag = 0;
            tol = &mgard_tolF;
            break;
        default:
            adios_error(err_transform_failure, "No supported data type\n");
            return -1;
            break;
    }

    // Get the input data and data length
    const uint64_t input_size = adios_transform_get_pre_transform_var_size(var);
    const void *input_buff = var->data;


    // Get dimension info
    struct adios_dimension_struct* d = var->pre_transform_dimensions;
    int ndims = (uint) count_dimensions(d);
    //log_debug("ndims: %d\n", ndims);
    if (ndims != 2)
    {
        log_warn("MGARD compression only supports 2D arrays.\n");
        ncol = nrow = 0;
    }
    else if (fd->group->adios_host_language_fortran == adios_flag_yes)
    {
        //log_debug("%s: Fortran dimension enabled\n", "MGARD");
        ncol = (int) adios_get_dim_value(&d->dimension);
        nrow = (int) adios_get_dim_value(&d->next->dimension);
    }
    else
    {
        //log_debug("%s: C dimension enabled\n", "MGARD");
        nrow = (int) adios_get_dim_value(&d->dimension);
        ncol = (int) adios_get_dim_value(&d->next->dimension);
    }

    struct adios_transform_spec_kv_pair* param;
    int i = 0;
    if (adios_verbose_level>7) log_debug("param_count: %d\n", var->transform_spec->param_count);
    for (i=0; i<var->transform_spec->param_count; i++)
    {
        param = &(var->transform_spec->params[i]);
        if (adios_verbose_level>7) log_debug("param: %s %s\n", param->key, param->value);

        if (!strcmp(param->key, "tol") || !strcmp(param->key, "accuracy"))
        {
            mgard_tolD = atof(param->value);
            mgard_tolF = (float) mgard_tolD;
        }
        else if (!strcmp(param->key, "zchecker") || !strcmp(param->key, "zcheck") || !strcmp(param->key, "z-checker") || !strcmp(param->key, "z-check"))
        {
            use_zchecker = (param->value == NULL)? 1 : atof(param->value);
        }
        else if (strcmp(param->key, "zc_init") == 0)
        {
            zc_configfile = strdup(param->value);
        }
        else
        {
            log_warn("An unknown MGARD parameter: %s\n", param->key);
        }
    }
    if (mgard_tolD < 1e-8)
    {
        log_warn("MGARD compression supports accurary 1e-8 or higher. Will use 1e-8 now\n");
        mgard_tolD = 1e-8;
        mgard_tolF = 1.1e-8; // 1e-8 is not exactly represented in float so we need something bit higher
    }

    unsigned char* mgard_comp_buff;
#ifdef USE_ZCHECKER
    log_debug("%s: %s\n", "Z-checker", "Enabled"); 
    ZC_DataProperty* dataProperty = NULL;
    ZC_CompareData* compareResult = NULL;
#endif

    int rtn = -1;
    // zero sized data will not be compressed
    // MGARD only works with 2D data
    if(input_size > 0u && ndims == 2)
    {

#ifdef USE_ZCHECKER
        int zc_type = zcheck_type(var->pre_transform_type);
        size_t r[5] = {0,0,0,0,0};
        zcheck_dim(fd, var, r);
        if (use_zchecker)
        {
            if (check_file(zc_configfile))
            {
                ZC_Init(zc_configfile);
                //ZC_DataProperty* ZC_startCmpr(char* varName, int dataType, void* oriData, size_t r5, size_t r4, size_t r3, size_t r2, size_t r1);
                dataProperty = ZC_startCmpr(var->name, zc_type, (void *) input_buff, r[4], r[3], r[2], r[1], r[0]);
            }
            else
            {
                log_warn("Failed to access Z-Check config file (%s). Disabled. \n", zc_configfile);
                use_zchecker = 0;
            }
        }
#endif
        //unsigned char *mgard_compress(int itype_flag, void *data, int *out_size, int nrow, int ncol, void* tol);
        mgard_comp_buff = mgard_compress(iflag, (void *) input_buff, &out_size,  nrow,  ncol, tol );
#ifdef USE_ZCHECKER
        if (use_zchecker)
        {
            //ZC_CompareData* ZC_endCmpr(ZC_DataProperty* dataProperty, int cmprSize);
            compareResult = ZC_endCmpr(dataProperty, (int)out_size);
            // For entropy
            ZC_DataProperty* property = ZC_genProperties(var->name, zc_type, (void *) input_buff, r[4], r[3], r[2], r[1], r[0]);
            dataProperty->entropy = property->entropy;
            freeDataProperty(property);

            ZC_startDec();
            void* hat = mgard_decompress(iflag, mgard_comp_buff, out_size,  nrow,  ncol);         
            ZC_endDec(compareResult, "SZ", hat);
            free(hat);
            log_debug("Z-Checker done.\n");
        }
#endif
        rtn = 0;
    }

    if(0 != rtn)  // compression avoided, then just copy the buffer
    {
        mgard_comp_buff = (unsigned char *) malloc (input_size);
        memcpy(mgard_comp_buff, input_buff, input_size);
        out_size = input_size;
    }


    log_debug("%s: %d,%d\n", "MGARD now,ncol", nrow, ncol);
    if (iflag)
    {
        log_debug("%s: %g\n", "MGARD tol", mgard_tolD);
    }
    else
    {
        log_debug("%s: %f\n", "MGARD tol", mgard_tolF);
    }
    log_debug("%s: %d\n", "MGARD out_size", out_size);
    log_debug("%s: %p\n", "MGARD output buffer", mgard_comp_buff);

    /*
    FILE *qfile;
    char fname[80];
    sprintf(fname, "input_type%d_%dx%d.dat", iflag, nrow, ncol);
    log_debug("%s: %s\n", "MGARD save", fname);
    qfile = fopen (fname, "wb");
    fwrite (input_buff, 1, nrow*ncol*8, qfile);
    fclose(qfile);
    */

    double *v = (double*) input_buff;
    log_debug("%s: %g %g %g %g %g ... %g %g %g %g %g\n", "MGARD input_buff",
              v[0], v[1], v[2], v[3], v[4],
              v[nrow*ncol-5], v[nrow*ncol-4], v[nrow*ncol-3], v[nrow*ncol-2], v[nrow*ncol-1]);

    // Output
    uint64_t output_size = (uint64_t) out_size/* Compute how much output size we need */;
    void* output_buff;

    log_debug("%s: %" PRIu64 "\n", "MGARD output_size", output_size);
    log_debug("%s: %d\n", "MGARD use_shared_buffer", use_shared_buffer);

    if (use_shared_buffer) {
        // If shared buffer is permitted, serialize to there
        assert(shared_buffer_reserve(fd, output_size));

        // Write directly to the shared buffer
        output_buff = fd->buffer + fd->offset;
        memcpy(output_buff, mgard_comp_buff, (size_t)output_size);
        // No more need
        free(mgard_comp_buff);
    } else { // Else, fall back to var->adata memory allocation
        output_buff = mgard_comp_buff;
        //assert(output_buff);
    }
    *wrote_to_shared_buffer = use_shared_buffer;

    // Do transform from input_buff into output_buff, and update output_size to the true output size

    // Wrap up, depending on buffer mode
    if (*wrote_to_shared_buffer) {
        shared_buffer_mark_written(fd, output_size);
    } else {
        var->adata = output_buff;
        var->data_size = output_size;
        var->free_data = adios_flag_yes;
    }

    *transformed_len = output_size; // Return the size of the data buffer

#ifdef USE_ZCHECKER
    // Have to do this after setting buffer size for adios
    if (use_zchecker && !rtn)
    {
        zcheck_write(dataProperty, compareResult, fd, var);
        log_debug("Z-Checker written.\n");
        freeDataProperty(dataProperty);
        freeCompareResult(compareResult);
        ZC_Finalize();
    }
#endif
    return 1;
}

#else

DECLARE_TRANSFORM_WRITE_METHOD_UNIMPL(mgard)

#endif
