#include <iostream>
#include <cstdlib>
#include <vector>
#include <string>
#include <map>
using namespace std;

#include <boost/program_options.hpp>
namespace bpo = boost::program_options;

#include "legion.hpp"
#include "cudamaniple.hpp"
#include "cudapinnedmemorytask.hpp"
#include "cudacheck.hpp"


// ==============================================================

const unsigned int nManiplesDefault = 2;
const unsigned int nAllocsDefault = 20;
const size_t nAllocBytesDefault = 1024;
const int selectGPUDefault = -1;

unsigned int nManiples;
unsigned int nAllocs;
size_t nAllocBytes;
int selectGPU;


void ReadCommandLine( int ac, char* av[] ) {

  try {
    bpo::options_description generic("Generic options");
    generic.add_options()
      ("help", "Produce help message" )
      ;

    bpo::options_description control("Control options");
    control.add_options()
      ("maniples",bpo::value<unsigned int>(&nManiples)->default_value(nManiplesDefault),"Number of maniples to enlist")
      ("allocs",bpo::value<unsigned int>(&nAllocs)->default_value(nAllocsDefault),"Number allocations to perform")
      ("bytes",bpo::value<size_t>(&nAllocBytes)->default_value(nAllocBytesDefault),"Number of bytes per allocation" )
      ("target",bpo::value<int>(&selectGPU)->default_value(selectGPUDefault),"Put all allocations on this maniple (-ve for any)" )
      ;

    bpo::options_description cmdLine;
    cmdLine.add(generic).add(control);

    bpo::variables_map vm;
    bpo::store( bpo::command_line_parser(ac, av).
		options(cmdLine).run(), vm );
    bpo::notify( vm );
    
    if( vm.count( "help" ) ) {
      cout << cmdLine << endl;
      exit( EXIT_SUCCESS );
    }
  }
  catch( exception& e ) {
    cerr << "Error: " << e.what() << endl;
    exit( EXIT_FAILURE );
  }
  catch( ... ) {
    cerr << "Unknown exception" << endl;
    exit( EXIT_FAILURE );
  }

  if( selectGPU >= static_cast<int>(nManiples) ) {
    cerr << "Must have selected GPU in range!" << endl;
    exit( EXIT_FAILURE );
  }

}

  



// ==============================================================

int main( int argc, char *argv[] ) {

  ReadCommandLine( argc, argv );

  cout << "Legion CUDA Pinned Memory" << endl;
  cout << "=========================" << endl << endl;

  // Create the Legion
  SciGPU::Legion::Legion myLegion;

  // Map for maniples and threads
  map<int,boost::thread::id> manipleIDs;

  // Give the Legion some Maniples
  for( unsigned int i=0; i<nManiples; i++ ) {
    SciGPU::Legion::CUDAmaniple myManiple;
    myManiple.gpuID = i;
    manipleIDs[i] = myLegion.AddManiple( myManiple );
  }

  // Declare a list of host pointers
  vector<char*> h_ptr(nAllocs,NULL);
  
  // Create the list of allocation tasks
  vector<SciGPU::Legion::CUDAhostAllocTask> memAllocs;

  for( unsigned int i=0; i<nAllocs; i++ ) {
    SciGPU::Legion::CUDAhostAllocTask nextTask;

    nextTask.ptr = (void**)&(h_ptr.at(i));
    nextTask.bytes = nAllocBytes;
    if( selectGPU >= 0 ) {
      nextTask.tid = manipleIDs.find( selectGPU )->second;
    }

    memAllocs.push_back( nextTask );
  } // Task list complete

  // Enqueue the tasks
  for( unsigned int i=0; i<memAllocs.size(); i++ ) {
    myLegion.Enqueue( &(memAllocs.at(i)) );
  }

  cout << "Tasks enqueued" << endl;

  // Wait for completion
  volatile size_t nComplete;
  do {
    boost::posix_time::seconds slp(1);  
    boost::this_thread::sleep( slp );
    nComplete = myLegion.CountComplete();
  } while( nComplete < nAllocs );
  
  cout << "Allocations complete" << endl;


  // Create the tasks to free memory
  vector<SciGPU::Legion::CUDAhostFreeTask> memFrees;  

  for( unsigned int i=0; i<h_ptr.size(); i++ ) {
    SciGPU::Legion::CUDAhostFreeTask nextTask;

    nextTask.tag = 1;
    nextTask.ptr = reinterpret_cast<void**>(&(h_ptr.at(i)));
    if( selectGPU >=0 ) {
      nextTask.tid = manipleIDs.find( selectGPU )->second;
    }

    memFrees.push_back( nextTask );
  } // Task list complete

  // Enqueue them
  for( unsigned int i=0; i<memFrees.size(); i++ ) {
    myLegion.Enqueue( &(memFrees.at(i)) );
  }

  // Wait for completion
  do {
    boost::posix_time::seconds slp(1);  
    boost::this_thread::sleep( slp );
    nComplete = myLegion.CountComplete();
  } while( nComplete < nAllocs );
  
  cout << "Release complete" << endl;
  

  // Show some stats
  cout << endl;
  myLegion.PrintCompletedStats( cout );
  cout << endl;

  return( EXIT_SUCCESS );
}
