#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include "configure.h"
#include "mkl_dnn.h"

#define CHECK_ERR(f, err)                                          \
    do                                                             \
    {                                                              \
        (err) = (f);                                               \
        if ((err) != E_SUCCESS)                                    \
        {                                                          \
            printf("[%s:%d] err (%d)\n", __FILE__, __LINE__, err); \
            goto bail_out;                                         \
        }                                                          \
    } while (0)

#define dimension (4)

static dnnError_t init_conversion(dnnPrimitive_t *cv, double **ptr_out,
                                  dnnLayout_t lt_pr, dnnLayout_t lt_us, double *ptr_us)
{
    dnnError_t err;
    *ptr_out = NULL;
    if (!dnnLayoutCompare_F64(lt_pr, lt_us))
    {
        CHECK_ERR(dnnConversionCreate_F64(cv, lt_us, lt_pr), err);
        CHECK_ERR(dnnAllocateBuffer_F64((void **)ptr_out, lt_pr), err);
    }
    else
    {
        *ptr_out = ptr_us;
    }
    return E_SUCCESS;

bail_out:
    if (*ptr_out)
        dnnReleaseBuffer_F64(*ptr_out);
    return err;
}

static dnnError_t simple_net(int want_groups_conv)
{
    dnnError_t err;

    size_t outputSize[dimension] = {N, N, 64, BATCH_SIZE};
    size_t outputStrides[dimension] = {1, N, N * N, N * N * 64};

    size_t inputSize[dimension] = {N, N, 3, BATCH_SIZE};
    size_t inputStrides[dimension] = {1, N, N * N, N * N * 3};

    size_t filterSize[dimension] = {3, 3, 3, 64};
    size_t filterStrides[dimension] = {1, 3, 3 * 3, 3 * 3 * 3};

    size_t convolutionStride[dimension - 2] = {1, 1};
    int inputOffset[dimension - 2] = {-1, -1};

    size_t biasSize[1] = {outputSize[2]};
    size_t biasStrides[1] = {outputStrides[2]};

    dnnLayout_t lt_user_input = NULL,
                lt_user_filt = NULL,
                lt_user_bias = NULL,
                lt_user_output = NULL;
    dnnPrimitive_t conv1 = NULL;
    dnnLayout_t lt_conv1_input = NULL,
                lt_conv1_filt = NULL,
                lt_conv1_bias = NULL,
                lt_conv1_output = NULL;
    double *resConv1[dnnResourceNumber] = {0};
    dnnPrimitive_t cv_user_to_conv1_input = NULL,
                   cv_user_to_conv1_filt = NULL,
                   cv_user_to_conv1_bias = NULL;
    dnnPrimitive_t relu1 = NULL;
    double *resRelu1[dnnResourceNumber] = {0};
    dnnLayout_t lt_relu1_output = NULL;
    dnnPrimitive_t cv_relu1_to_user_output = NULL;

    dnnPrimitive_t bn1 = NULL;
    double *resBn1[dnnResourceNumber] = {0};
    dnnLayout_t lt_bn1_output = NULL;
    dnnPrimitiveAttributes_t attributes = NULL;

    double *user_i = NULL,
           *user_c1_f = NULL,
           *user_c1_b = NULL,
           *user_o = NULL;

    /*** data allocation ***/
    user_i = (double *)malloc(sizeof(double) * (inputSize[0] * inputSize[1] * inputSize[2] * inputSize[3]));
    user_c1_f = (double *)malloc(sizeof(double) * (filterSize[0] * filterSize[1] * filterSize[2] * filterSize[3]));
    user_c1_b = (double *)malloc(sizeof(double) * (8 * 1024 * 1024));
    if (user_i == NULL || user_c1_f == NULL || user_c1_b == NULL)
    {
        err = E_MEMORY_ERROR;
        goto bail_out;
    }

    /*** User's data description ***/
    CHECK_ERR(dnnLayoutCreate_F64(&lt_user_input, dimension, inputSize, inputStrides), err);
    CHECK_ERR(dnnLayoutCreate_F64(&lt_user_filt, dimension, filterSize, filterStrides), err);
    CHECK_ERR(dnnLayoutCreate_F64(&lt_user_bias, 1, biasSize, biasStrides), err);
    CHECK_ERR(dnnLayoutCreate_F64(&lt_user_output, dimension, outputSize, outputStrides), err);

    /* Initialize attributes */
    CHECK_ERR(dnnPrimitiveAttributesCreate_F64(&attributes), err);

    /*** convolution section ***/
    if (!want_groups_conv)
    {
        CHECK_ERR(dnnConvolutionCreateForwardBias_F64(&conv1, attributes,
                                                      dnnAlgorithmConvolutionDirect, dimension, inputSize,
                                                      outputSize, filterSize, convolutionStride, inputOffset,
                                                      dnnBorderZeros),
                  err);
    }
    else
    {
        CHECK_ERR(dnnGroupsConvolutionCreateForwardBias_F64(&conv1, attributes,
                                                            dnnAlgorithmConvolutionDirect, 1, dimension, inputSize,
                                                            outputSize, filterSize, convolutionStride, inputOffset,
                                                            dnnBorderZeros),
                  err);
    }
    // Convolution describes what layout it expects
    CHECK_ERR(dnnLayoutCreateFromPrimitive_F64(&lt_conv1_input, conv1, dnnResourceSrc), err);
    CHECK_ERR(dnnLayoutCreateFromPrimitive_F64(&lt_conv1_filt, conv1, dnnResourceFilter), err);
    CHECK_ERR(dnnLayoutCreateFromPrimitive_F64(&lt_conv1_bias, conv1, dnnResourceBias), err);
    CHECK_ERR(dnnLayoutCreateFromPrimitive_F64(&lt_conv1_output, conv1, dnnResourceDst), err);

    CHECK_ERR(init_conversion(&cv_user_to_conv1_input, &resConv1[dnnResourceSrc], lt_conv1_input, lt_user_input, user_i), err);
    CHECK_ERR(init_conversion(&cv_user_to_conv1_filt, &resConv1[dnnResourceFilter], lt_conv1_filt, lt_user_filt, user_c1_f), err);
    CHECK_ERR(init_conversion(&cv_user_to_conv1_bias, &resConv1[dnnResourceBias], lt_conv1_bias, lt_user_bias, user_c1_b), err);
    CHECK_ERR(dnnAllocateBuffer_F64((void **)&resConv1[dnnResourceDst], lt_conv1_output), err);

    /*** BN1 section ***/
    CHECK_ERR(dnnBatchNormalizationCreateForward_F64(&bn1, attributes, lt_conv1_output, 0), err);
    resBn1[dnnResourceSrc] = resConv1[dnnResourceDst];
    CHECK_ERR(dnnLayoutCreateFromPrimitive_F64(&lt_bn1_output, bn1, dnnResourceDst), err);
    CHECK_ERR(dnnAllocateBuffer_F64((void **)&resBn1[dnnResourceDst], lt_bn1_output), err);
    CHECK_ERR(init_conversion(&cv_relu1_to_user_output, &user_o, lt_user_output, lt_bn1_output, resBn1[dnnResourceDst]), err);

    /*** ReLU section ***/
    CHECK_ERR(dnnReLUCreateForward_F64(&relu1, attributes, lt_bn1_output, 0.0f), err);
    resRelu1[dnnResourceSrc] = resBn1[dnnResourceDst];
    CHECK_ERR(dnnLayoutCreateFromPrimitive_F64(&lt_relu1_output, relu1, dnnResourceDst), err);
    CHECK_ERR(dnnAllocateBuffer_F64((void **)&resRelu1[dnnResourceDst], lt_relu1_output), err);
    CHECK_ERR(init_conversion(&cv_relu1_to_user_output, &user_o, lt_user_output, lt_relu1_output, resRelu1[dnnResourceDst]), err);

    srand(1);

    for (int i = 0; i < inputSize[0] * inputSize[1] * inputSize[2] * inputSize[3]; i++)
        user_i[i] = rand() % 10;
    for (int i = 0; i < filterSize[0] * filterSize[1] * filterSize[2] * filterSize[3]; i++)
        user_c1_f[i] = 1;
    for (int i = 0; i < outputSize[2]; i++)
        user_c1_b[i] = 0;

    /*** Execution ***/
    if (cv_user_to_conv1_filt)
        CHECK_ERR(dnnConversionExecute_F64(cv_user_to_conv1_filt, user_c1_f, resConv1[dnnResourceFilter]), err);
    if (cv_user_to_conv1_bias)
        CHECK_ERR(dnnConversionExecute_F64(cv_user_to_conv1_bias, user_c1_b, resConv1[dnnResourceBias]), err);

    if (cv_user_to_conv1_input)
        CHECK_ERR(dnnConversionExecute_F64(cv_user_to_conv1_input, user_i, resConv1[dnnResourceSrc]), err);
    double times[NB_TESTS];
    clock_t start, end;
    for (int i = 0; i < NB_TESTS; i++)
    {
        start = clock();
        CHECK_ERR(dnnExecute_F64(conv1, (void *)resConv1), err);
        printf("ok\n");

        CHECK_ERR(dnnExecute_F64(bn1, (void *)resBn1), err);

        CHECK_ERR(dnnExecute_F64(relu1, (void *)resRelu1), err);
        printf("not ok\n");

        end = clock();

        double time_taken = ((double)(end - start) / CLOCKS_PER_SEC) * 1000;
        times[i] = time_taken;
    }
    printf("Convolution time: %f.\n", median(NB_TESTS, times));

    if (cv_relu1_to_user_output)
        CHECK_ERR(dnnConversionExecute_F64(cv_relu1_to_user_output, resRelu1[dnnResourceDst], user_o), err); // Shift pointers

    FILE *f = fopen("mkl_result.txt", "w");
    if (f == NULL)
    {
        printf("Error opening file!\n");
        exit(1);
    }
    for (int n = 0; n < BATCH_SIZE; ++n)
    {
        for (int z = 0; z < 64; ++z)
        {
            for (int y = 0; y < N; ++y)
            {
                for (int x = 0; x < N; ++x)
                {
                    fprintf(f, "%.2f", user_o[x + y * N + z * N * N + n * N * N * 64]);
                }
                // printf("\n");
            }
            //printf("\n");
        }
        // printf("\n");
    }
    fclose(f);
bail_out:

    dnnDelete_F64(conv1);
    dnnDelete_F64(relu1);
    dnnDelete_F64(bn1);

    dnnDelete_F64(cv_user_to_conv1_input);
    dnnDelete_F64(cv_user_to_conv1_filt);
    dnnDelete_F64(cv_user_to_conv1_bias);
    //dnnDelete_F64(cv_pool1_to_user_output);

    dnnLayoutDelete_F64(lt_user_input);
    dnnLayoutDelete_F64(lt_user_filt);
    dnnLayoutDelete_F64(lt_user_bias);
    dnnLayoutDelete_F64(lt_user_output);
    dnnLayoutDelete_F64(lt_conv1_input);
    dnnLayoutDelete_F64(lt_conv1_filt);
    dnnLayoutDelete_F64(lt_conv1_bias);
    dnnLayoutDelete_F64(lt_conv1_output);
    dnnLayoutDelete_F64(lt_relu1_output);
    dnnLayoutDelete_F64(lt_bn1_output);

    dnnPrimitiveAttributesDestroy_F64(attributes);
    if (resConv1[dnnResourceSrc] != (void *)user_i)
        dnnReleaseBuffer_F64(resConv1[dnnResourceSrc]);
    if (resConv1[dnnResourceFilter] != (void *)user_c1_f)
        dnnReleaseBuffer_F64(resConv1[dnnResourceFilter]);
    if (resConv1[dnnResourceBias] != (void *)user_c1_b)
        dnnReleaseBuffer_F64(resConv1[dnnResourceBias]);
    dnnReleaseBuffer_F64(resConv1[dnnResourceDst]);
    dnnReleaseBuffer_F64(resRelu1[dnnResourceDst]);
    dnnReleaseBuffer_F64(resBn1[dnnResourceDst]);

    free(user_i);
    free(user_c1_f);
    free(user_c1_b);

    return err;
}

int main(int argc, char **argv)
{
    dnnError_t err;
    err = simple_net(0);
    if (err != E_SUCCESS)
    {
        printf("FAILED\n");
        return err;
    }
    err = simple_net(1);
    if (err != E_SUCCESS)
    {
        printf("FAILED\n");
        return err;
    }

    printf("PASSED\n");
    return 0;
}