#include <cassert>
#include <chrono>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <string>

#ifdef USE_MPI
/* MPI header */
#include <mpi.h>
#endif

#include "charmap.h"

// ./generator <a> <m> <pixel graphics output file name> <vector graphics output file name>
//
// size a, default a = 1024, the figure size is n = a^2, with a pixels in x direction and a pixels in y direction
// number of disks m used to generate the image, default m = 10
//
// default pixel output file name: benchmark.pxl
// default vector output file name: benchmark.vct
//
int main(int argc, char** argv)
{
   auto t0 = std::chrono::high_resolution_clock::now();
#ifdef USE_MPI
   MPI_Init(&argc, &argv);
   
   // how many processes are there?
   int size = 0;  
   MPI_Comm_size(MPI_COMM_WORLD, &size); 
   
   // what is the rank of this process?
   int rank = 0;
   MPI_Comm_rank(MPI_COMM_WORLD, &rank); 
#else
   int size = 1;
   int rank = 0;
#endif
   auto t1 = std::chrono::high_resolution_clock::now();  // after MPI launch

   int a = 1024;
   if(argc > 1) a = std::atoi(argv[1]);
   assert(a > 1);
   if(!rank) std::cout << "Image edge size:\t" << a << " pixels (" << a << " x " << a << ")\n";
   
   int m = 10;
   if(argc > 2) m = std::atoi(argv[2]);
   assert(m > 0);
   if(!rank) std::cout << "No. random disks:\t" << m << "\n";

   std::string pixout_fname = "benchmark.pxl";
   if(argc > 3) pixout_fname = argv[3];
   if(!rank) std::cout << "Pixel output to:\t" << pixout_fname << "\n";
   
   std::string vctout_fname = "benchmark.vct";
   if(argc > 4) vctout_fname = argv[4];
   if(!rank) std::cout << "Vector output to:\t" << vctout_fname << "\n";

   // random seed from chrono
   long random_seed = std::chrono::duration_cast<std::chrono::nanoseconds>(
         std::chrono::high_resolution_clock::now().time_since_epoch()
      ).count();
#ifdef USE_MPI
   // make all ranks use the same random seed, taken from rank 0
   MPI_Bcast(&random_seed, 1, MPI_LONG, 0, MPI_COMM_WORLD);
#endif
   std::srand(random_seed);
   
   // create the random figure
   //
   diskgraphics::DiskVector dv;
   dv.set_size(a, a);
   dv.set_background(255);
   float disk_radius = a;
   unsigned char disk_colour = 0;
   for(int i = 0; i < m; i++)
   {
      diskgraphics::Disk dsk;
      disk_radius /= 1.125 + 0.125*std::rand()/(float)RAND_MAX;  // shrink down by random factor between 1.125 and 1.25
      
      dsk.x = a * (0.75*(std::rand()/(float)RAND_MAX) + 0.125);
      dsk.y = a * (0.75*(std::rand()/(float)RAND_MAX) + 0.125);
      dsk.r = disk_radius;
      dsk.colour = disk_colour;
      dv.add_disk(dsk);
      
      // switch colour from one disk to the next with 2/3 probability
      if(std::rand() > RAND_MAX/3.0) disk_colour = 255 - disk_colour;
   }
   
   // vector file output
   if(!rank)
   {
      std::ofstream vctostrm(vctout_fname);
      if(!vctostrm) return EXIT_FAILURE;
      vctostrm << dv;
      vctostrm.close();
   }
   
   // what part of the figure is this process in charge of?
   //
   int ny = 1 + (a-1)/size;
   int yoffset = rank*ny;
   if(rank == size-1) ny = a - yoffset;
   // std::cout << "\t\tRank " << rank << " in charge of " << ny << " rows, starting from y = " << yoffset << "\n";
   
   // pixel conversion
   //
   diskgraphics::Charmap cm(dv, 0, a, yoffset, ny);
   auto t2 = std::chrono::high_resolution_clock::now();  // after computational task

   // pixel output
   //
   if(32 >= a)
   {
#ifdef USE_MPI
      for(int i = 0; i < rank; i++) MPI_Barrier(MPI_COMM_WORLD);
      std::cout << "Rank " << rank << ":\n===\n";
#endif
      cm.out(&std::cout, true, false);  // if a small figure, also show on cout
#ifdef USE_MPI
      std::cout << "===\n\n";
      std::cout.flush();
      for(int i = rank; i < size; i++) MPI_Barrier(MPI_COMM_WORLD);
#endif
   }
   
   constexpr size_t header_size = 20 * sizeof(char);
   if(!rank)
   {
      std::ofstream pixostrm(pixout_fname);
      if(!pixostrm) return EXIT_FAILURE;
      
      // rank 0 write out header (figure size)
      pixostrm.width(9);
      pixostrm.fill('0');
      pixostrm << a << " ";
      pixostrm.width(9);
      pixostrm.fill('0');
      pixostrm << a << "\n";
#ifndef USE_MPI
      // standard sequential output
      cm.out(&pixostrm, false, false);
#endif
      pixostrm.close();
   }
#ifdef USE_MPI
   size_t elements = cm.get_sizex() * cm.get_sizey();
   unsigned char* content = cm.access_data();
   
   // displacement, etype, and file type information
   MPI_Offset displacement = header_size + (yoffset*a)*sizeof(unsigned char);
   MPI_Datatype etype = MPI_UNSIGNED_CHAR;
   MPI_Datatype etype_array;
   MPI_Type_contiguous(elements, etype, &etype_array);
   MPI_Type_commit(&etype_array);
   
   // open the file
   MPI_File fh;
   MPI_File_open(MPI_COMM_WORLD, pixout_fname.c_str(), MPI_MODE_WRONLY, MPI_INFO_NULL, &fh);
   
   /*
   for(int i = 0; i < rank; i++) MPI_Barrier(MPI_COMM_WORLD);
   std::cout << "\trank " << rank << "\t\tdisplacement " << displacement << "\t\telements " << elements << "\n";
   std::cout.flush();
   for(int i = rank; i < size; i++) MPI_Barrier(MPI_COMM_WORLD);
    */
   
   // now create a "view" consisting of the displacement, the etype, and the file type
   MPI_File_set_view(fh, displacement, etype, etype_array, "native", MPI_INFO_NULL);
   MPI_File_write(fh, content, elements, etype, MPI_STATUS_IGNORE);
   MPI_File_close(&fh);
#endif
   
   auto t3 = std::chrono::high_resolution_clock::now();  // after file output

#ifdef USE_MPI
   MPI_Finalize();
#endif
   auto t4 = std::chrono::high_resolution_clock::now();  // after MPI cleanup
   
   if(!rank)
   {
      auto delta_t01 = std::chrono::duration_cast<std::chrono::microseconds>(t1-t0).count();
      auto delta_t12 = std::chrono::duration_cast<std::chrono::microseconds>(t2-t1).count();
      auto delta_t23 = std::chrono::duration_cast<std::chrono::microseconds>(t3-t2).count();
      auto delta_t34 = std::chrono::duration_cast<std::chrono::microseconds>(t4-t3).count();
      auto totaltime = std::chrono::duration_cast<std::chrono::microseconds>(t4-t0).count();
      
      std::cout << "\n===\n"
                << "Parallel environment setup:\t"   << 1.0e-06*delta_t01 << " s\n"
                << "Character map generation:\t"   << 1.0e-06*delta_t12 << " s\n"
                << "Character map file output:\t"   << 1.0e-06*delta_t23 << " s\n"
                << "Parallel environment cleanup:\t" << 1.0e-06*delta_t34 << " s\n===\n"
                << "Total program execution time:\t" << 1.0e-06*totaltime << " s\n";
      std::cout.flush();
   }
}