#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include "mystery.h"

const char mystery_op0_type[] = "custom";
const char mystery_op1_type[] = "bev_pool_v2";


void get_output_shape(const int** input_shapes, const int* input_shape_sizes, const int* input_data_types, int input_num, char** attr_names, int* attr_types, int* attr_sizes, void** attr_values, const int attr_num, struct DimensionInfo* out_info){
    out_info->num = 1;
    out_info->shape_sizes = (int *)malloc(out_info->num * sizeof(int));
    out_info->shapes = (int **)malloc(out_info->num * sizeof(int *));
    out_info->types = (int *)malloc(out_info->num * sizeof(int));
    int shape_sizes[] = {4};
    int shapes[1][4] = {{1, 128, 128, input_shapes[1][3]}};
    int types[] = {input_data_types[0]};
    for (int i = 0; i < 1; i++){
        out_info->shape_sizes[i] = shape_sizes[i];
        out_info->shapes[i] = (int*)malloc(shape_sizes[i] * sizeof(int));
        for (int j = 0; j < shape_sizes[i]; j++){
            out_info->shapes[i][j] = shapes[i][j];
        }
        out_info->types[i] = types[i]; 
    }
        //Though currently all inputs outputs are converted to float by piano. It is still reserved to future usage.
}

void custom_shape(const int** input_shapes, const int* input_shape_sizes, const int* input_data_types, int input_num, char** attr_names, int* attr_types, int* attr_sizes, void** attr_values, const int attr_num, struct DimensionInfo* out_info){
    out_info->num = 1;
    out_info->shape_sizes = (int *)malloc(out_info->num * sizeof(int));
    out_info->shapes = (int **)malloc(out_info->num * sizeof(int *));
    out_info->types = (int *)malloc(out_info->num * sizeof(int));
    int shape_sizes[] = {2};
    int shapes[1][2] = {{3, 4}};
    int types[] = {1};
    for (int i = 0; i < 1; i++){
        out_info->shape_sizes[i] = shape_sizes[i];
        out_info->shapes[i] = (int*)malloc(shape_sizes[i] * sizeof(int));
        for (int j = 0; j < shape_sizes[i]; j++){
            out_info->shapes[i][j] = shapes[i][j];
        }
        out_info->types[i] = types[i]; 
    }
        //Though currently all inputs outputs are converted to float by piano. It is still reserved to future usage.
    for (int i = 0; i < attr_num; i++){
        printf("current attr name is %s\n", attr_names[i]);
        if (strcmp(attr_names[i], "add_0") == 0){
            int64_t* curr_value = (int64_t*) attr_values[i];
            for (int k = 0; k < attr_sizes[i]; k++){
                printf("%lld ",curr_value[k]);
            }
            printf("\n");
        }
        else if (strcmp(attr_names[i], "add_1") == 0){
            int64_t* curr_value = (int64_t*) attr_values[i];
            for (int k = 0; k < attr_sizes[i]; k++){
                printf("%lld ",curr_value[k]);
            }
            printf("\n");
        }
        else if (strcmp(attr_names[i], "add_2") == 0){
            char* curr_value = (char*) attr_values[i];
            printf("str found, got %s\n", curr_value);
        }

    }
}

void cumtom_op(const void* const* inputs, const int* const* input_shapes, const int* input_shape_sizes, int* input_data_types, int input_num,
                char** attr_names, int* attr_types, int* attr_sizes, void** attr_values, const int attr_num, void** outputs, const int* const* output_shapes, const int* output_shape_sizes, int* output_data_types, int output_num){
    int add_0 = 0;
    int add_1 = 0;
    printf("attr num is %i\n", attr_num);
    for (int i = 0; i < attr_num; i++){
        printf("compute: current attr name is %s\n", attr_names[i]);
        if (strcmp(attr_names[i], "add_0") == 0){
            int64_t* curr_value = (int64_t*) attr_values[i];
            add_0 = curr_value[0];
            printf("add_0 found, got %i\n", add_0);
        }
        else if (strcmp(attr_names[i], "add_1") == 0){
            int64_t* curr_value = (int64_t*) attr_values[i];
            printf("add_1 found, got ");
            for (int k = 0; k < attr_sizes[i]; k++){
                printf("%lld ",curr_value[k]);
            }
            printf("\n");
        }
        else if (strcmp(attr_names[i], "add_2") == 0){
            char* curr_value = (char*) attr_values[i];
            printf("str found, got %s\n", curr_value);
        }

    }
    float *input_0 = inputs[0];
    int *shape_input_0 = input_shapes[0];
    float *input_1 = inputs[1];
    int *shape_input_1 = input_shapes[1];
    float *output_0 = (float*)outputs[0];
    printf("sizes of input shapes:\n%i\n", input_shape_sizes[0]);
    printf("%i\n", input_shape_sizes[1]);
    printf("input_num: %i\n", input_num);
    int output0_size = 1;
    for (int i = 0; i < output_shape_sizes[0]; i++){
        output0_size *= output_shapes[0][i];
    }
    for (int i = 0; i < output0_size; i++){
        output_0[i] = input_0[i] + input_1[i];
        printf("output %i:%f, is sum of %f and %f\n", i, output_0[i], input_0[i], input_1[i]);
    }
}

void bev_pool_v2(const void* const* inputs, const int* const* input_shapes, const int* input_shape_sizes, int* input_data_types, int input_num,
                char** attr_names, int* attr_types, int* attr_sizes, void** attr_values, const int attr_num, void** outputs, const int* const* output_shapes, const int* output_shape_sizes, int* output_data_types, int output_num) {
    float *depth = (float*)inputs[0];
    float *feat = (float*)inputs[1];
    float *ranks_depth = (float*)inputs[2];
    float *ranks_feat = (float*)inputs[3];
    float *ranks_bev = (float*)inputs[4];
    float *interval_starts = (float*)inputs[5];
    float *interval_lengths = (float*)inputs[6];
    float *f_output = (float*)outputs[0];
    int out_height = 128;
    int out_width = 128;

    for (int i = 0; i < attr_num; i++){
        int* curr_value = (int*) attr_values[i];
        if (strcmp(attr_names[i], "out_height") == 0){
            out_height = curr_value[0];
        }
        if (strcmp(attr_names[i], "out_width") == 0){
            out_width = curr_value[0];
        }
    }
    int n_intervals = input_shapes[6][0];
    int bev_feat_shape[] = {1, out_height, out_width};
    int c2 = input_shapes[1][3];
    int output0_size = 1;
    for (int i = 0; i < output_shape_sizes[0]; i++){
        output0_size *= output_shapes[0][i];
    }
    for (int i = 0; i < output0_size; i++){
        f_output[i] = 0.;
    }

    for (int i = 0; i < n_intervals; i++){
        int interval_start = interval_starts[i];
        int interval_length = interval_lengths[i];
        for (int j = 0; j < interval_length; j++){
            for (int k = 0; k < c2; k++){
                f_output[(int)ranks_bev[interval_start]*c2 + k] += depth[(int)ranks_depth[interval_start+j]] * feat[(int)ranks_feat[interval_start+j]*c2 + k];
            }
        }
    }
}

void build_mystery_fn(struct MysteryFuMap mystery_fns[]) {
    mystery_fns[0].op_type = mystery_op0_type;
    mystery_fns[0].Mystery_fn = cumtom_op;
    mystery_fns[1].op_type = mystery_op1_type;
    mystery_fns[1].Mystery_fn = bev_pool_v2;
}