/* FLASH usecase on multiple ranks but all the writes and
   mappings happen on rank 0, all the processes have their 
   local handles, and local information, we need to see how
   we want to do this for larger runs.
*/

#include <assert.h>
#include <stdio.h>
#include <unistd.h>

//#include "test-flash.h"
#include "damsel-hdf5.h"
#include "damsel.h"
#include "damsel-internal.h"
#include "damsel-ndc.h"
#include "damsel-debug.h"
#include "damsel-connectivity.h"
#include "damsel-types.h"
#include "mpi.h"

#define TOTAL_BLOCKS 17
#define FILENAME "output/test-flash-parallel.h5"

/*Since its just a test case I am creating three arrays here and each rank will use corresponding subarray for tagging local blocks, 
 * I think it is ok to do like that because we dont need any communication for these three arrays. I just want to use correct values
 * instead of made up values I have for coords, size, unk and bounding box. nodetype 1 = root, 3 = leaf and 2 is all the others*/
int nodetype[17] = {1, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 3, 3, 3, 3};
int level[17] = {1, 2, 3, 4, 4, 4, 4, 3, 3, 3, 2, 2, 2, 3, 3, 3, 3};
int which_child[17] = {-1, 1, 1, 1, 2, 3, 4, 2, 3, 4, 2, 3, 4, 1, 2, 3, 4};

damsel_err_t define_tags(damsel_model model, damsel_data_type *types, damsel_handle *tag_handles);
void set_parenthandles(int myrank, int *proc_ranks, damsel_handle *block_id, damsel_handle *local_parent);

int main(int argc, char **argv) {
#ifndef DAMSEL_SINGLETHREAD
  int res;
#endif
  int myrank, i, j, local_blocks;
    
#ifndef DAMSEL_SINGLETHREAD
  MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &res);
  assert(res == MPI_THREAD_MULTIPLE);
#else
  MPI_Init(&argc, &argv);
#endif

  MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
  if(myrank == 0)
	local_blocks = 5; //nothing special about it, only following our 17 block example tree/node using 4 ranks
  else
	local_blocks = 4;

 
  /*Local arrays for holding metadata, etc*/
 double unk[local_blocks][24];
 double coords[local_blocks][2];
 double size[local_blocks][2];
 double bndbox[local_blocks][4];
 int local_nodetype[local_blocks];
 int local_level[local_blocks];
 int local_which_child[local_blocks];
 damsel_handle local_parent[local_blocks];
 damsel_handle block_id[local_blocks];
 int num;
 
 /*just a quick and dirty way to ensure the block number sequence on each rank*/
 if(myrank == 0) num =1;
 else num = 2;

 for(i=0;i<local_blocks;i++){ 
       block_id[i] = (myrank*local_blocks) + num + i;
       coords[i][0] = 0.0*myrank+i;
       coords[i][1] = 0.0*myrank+i;

       size[i][0] = 0.25*myrank+i;
       size[i][1] = 0.25*myrank+i;

       bndbox[i][0] = 0.0*myrank;
       bndbox[i][1] = 0.0*myrank;
       bndbox[i][2] = 0.0*myrank+i;
       bndbox[i][3] = 0.0*myrank+i;

       local_nodetype[i] = nodetype[myrank*local_blocks+num-1+i];
       local_level[i] = level[myrank*local_blocks+num-1+i];
       local_which_child[i] = which_child[myrank*local_blocks+num-1+i];

       for(j=0;j<24;j++)
                unk[i][j]=i+j; //solution data array, considering one variable and 2D data here
  }

  damsel_library lib;
  damsel_err_t err = DMSLlib_init(&lib);
  
  /*creating array data type for unk, and other block metadata tags*/

  size_t unk_array_dim[] = {24};
  damsel_data_type unk_array_type = DMSLtype_array_create(DAMSEL_DATA_TYPE_DOUBLE, 1, unk_array_dim);


  size_t coords_array_dim[] = {2}; //for coordinates and size in x, y dimension
  damsel_data_type coords_array_type = DMSLtype_array_create(DAMSEL_DATA_TYPE_DOUBLE, 1, coords_array_dim);

  size_t bndbox_array_dim[] = {4};
  damsel_data_type bndbox_array_type = DMSLtype_array_create(DAMSEL_DATA_TYPE_DOUBLE, 1, bndbox_array_dim);


  damsel_data_type types[] = {unk_array_type, DAMSEL_DATA_TYPE_HANDLE, DAMSEL_DATA_TYPE_HANDLE,
                              DAMSEL_DATA_TYPE_HANDLE, coords_array_type, coords_array_type, bndbox_array_type, 
			      DAMSEL_DATA_TYPE_INTEGER, DAMSEL_DATA_TYPE_INTEGER, DAMSEL_DATA_TYPE_INTEGER};
 
 /*end of array type definition*/

  damsel_model model;
  err = DMSLmodel_create(DAMSEL_HANDLE_TYPE_HANDLE64, &model);
  assert(model != NULL);

  /* Assume each rank will create a sequence handle collection with 4 handles in it. */
  damsel_container block_collection_id;
  err = DMSLcontainer_create_vector(model, block_id, local_blocks, &block_collection_id);

  damsel_handle tag_handles[10];


  damsel_container parent_collection;
  int proc_ranks[4]; //will not be used at rank 0

  
  set_parenthandles(myrank, proc_ranks, block_id, local_parent);
  /*Each rank creates a collection of parent handles for the block it owns*/
  
  err = DMSLcontainer_create_vector(model, local_parent, local_blocks, &parent_collection);
  /*At this point all ranks know who are the parents of blocks owned by each rank, we will map them later*/  
  /*Now we need to use some communication for neighbor blocks and child blocks to create gid array*/
  /*Fill in later*/

  define_tags(model, types, tag_handles);  

  /* now set the values of the metadata (parent, size, coords, etc) and data (UNK) tag on the handles */
  DMSLmodel_map_tag(local_parent, block_collection_id, &tag_handles[1]);
  DMSLmodel_map_tag(coords, block_collection_id, &tag_handles[4]);
  DMSLmodel_map_tag(size, block_collection_id, &tag_handles[5]);
  DMSLmodel_map_tag(bndbox, block_collection_id, &tag_handles[6]);
  DMSLmodel_map_tag(local_nodetype, block_collection_id, &tag_handles[7]);
  DMSLmodel_map_tag(local_level, block_collection_id, &tag_handles[8]);
  DMSLmodel_map_tag(local_which_child, block_collection_id, &tag_handles[9]);
  DMSLmodel_map_tag(unk, block_collection_id, &tag_handles[0]);

  /* at this point, the in-memory store in each rank knows about all the data that
     we have. next we will write the data out to disk.
     we need to invent some handles for the file store too. They can overlap
     with memstore handles, though here I make sure that they do not.
  */

  if(myrank == 0) unlink(FILENAME); // maybe it doesn't exist, so we'll just silently ignore any error
  damsel_trait_bundle tb;
  err = DMSLtrait_create(&tb);
  DMSLtrait_put(tb, "mode", "WRITE");
#ifdef DAMSEL_SUBFILE
  DMSLtrait_put(tb, "dmsl_num_subfiles", "2");
#endif

  DMSLmodel_attach(model, FILENAME, MPI_COMM_WORLD, tb);
  DMSLtrait_release(tb);

  DMSLmodel_map_handles_inventing_file_handles(block_collection_id);
 /*Now map handles on rank 0 to rank1 and rank2, and handles on rank2 to rank3*/
  if(myrank != 0)
  	DMSLmodel_map_global_handles(parent_collection, proc_ranks, local_parent);

  /* Now we can tell damsel to execute everything. */
  damsel_request_t req;
  DMSLmodel_transfer_async(model, DAMSEL_TRANSFER_TYPE_WRITE, &req);
  damsel_status_t statusR;
  DMSLmodel_wait(req, &statusR);

  DMSLmodel_close(model);
  DMSLcontainer_close(block_collection_id);

  DMSLlib_finalize(lib);

  MPI_Finalize();
  return 0;
}

/*Function to determine parent handles*/
void set_parenthandles(int myrank, int *proc_ranks, damsel_handle *block_id, damsel_handle *local_parent){
  int i;
  MPI_Status status;
  /*Rank 0 has blocks 1, 2, 3, 4 and 5. All the blocks have their parents on rank0*/
  if(myrank == 0) {
        local_parent[0] = -1;
        local_parent[4] = 3;
     for(i=1;i<4;i++){
        local_parent[i] = block_id[i-1];
        if(i<3) {
             MPI_Send(&block_id[i], 1, MPI_LONG_LONG_INT, 1, 10+i, MPI_COMM_WORLD);
             MPI_Send(&block_id[i-1], 1, MPI_LONG_LONG_INT, 2, 13+i, MPI_COMM_WORLD);
                }
             }
        }

  damsel_handle local_block_id[2];
  /*Rank 1 has blocks 6, 7, 8 and 9. Parent blocks of 6, 7 is 3 on rank 0 and for 8, 9 is 2 on rank 0. It receives these handles from Rank 0*/
  if(myrank == 1) {
        MPI_Recv(&local_block_id[0], 1, MPI_LONG_LONG_INT, 0, 11, MPI_COMM_WORLD, &status);
        MPI_Recv(&local_block_id[1], 1, MPI_LONG_LONG_INT, 0, 12, MPI_COMM_WORLD, &status);
        for(i=0;i<4;i++) {
          proc_ranks[i] = 0;
          if(i<2)
                local_parent[i] = local_block_id[1];
          else
                local_parent[i] = local_block_id[0];
        }
        }
  /*Rank 2 has blocks 10, 11, 12, and 13. Parent of 10 is 2 on rank 0, and parent of 11, 12 and 13 is 1 on rank 0. So rank 2 receives handles from rank 0*/
  if(myrank == 2) {
        MPI_Recv(&local_block_id[0], 1, MPI_LONG_LONG_INT, 0, 14, MPI_COMM_WORLD, &status);
        MPI_Recv(&local_block_id[1], 1, MPI_LONG_LONG_INT, 0, 15, MPI_COMM_WORLD, &status);
        local_parent[0] = local_block_id[1];
        proc_ranks[0] = 0;
        for(i=1;i<4;i++){
           local_parent[i] = local_block_id[0];
           proc_ranks[i] = 0;
        }
        MPI_Send(&block_id[3], 1, MPI_LONG_LONG_INT, 3, 17, MPI_COMM_WORLD);
  }

  /*Rank 3 has blocks 14, 15, 16, 17. The parent handle is 13 on rank 2, it receives 13 from rank 2*/
  if(myrank == 3) {
        MPI_Recv(&local_block_id[0], 1, MPI_LONG_LONG_INT, 2, 17, MPI_COMM_WORLD, &status);
        for(i=0;i<4;i++){
           local_parent[i] = local_block_id[0];
           proc_ranks[i] = 2;
        }
  }

}

damsel_err_t define_tags(damsel_model model, damsel_data_type *types, damsel_handle *tag_handles)
{
  int i, n=10;
  char *names[] = {"UNK", "Parents", "Children", "Neighbors", "COordinates", "Size", "bounding box", "Node type", "Level", "Which Child"};

  for (i=0; i<n; i++) {
    damsel_err_t err = DMSLtag_define(model, (damsel_handle_ptr)(tag_handles+i), types[i], names[i]);
        if (err.id != 0) return DMSL_ERROR_GENERAL;
    DMSLmodel_map_single_handle_inventing_file_handle(model, tag_handles[i]);
  }

  return DMSL_OK;
}