// -*- mode: c++ -*-

#include "mesh.hpp"
#include "agent.hpp"
#include "parser.hpp"

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

  int rank;
  int size;

  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  if (rank == 0) {
    
    std::cout << std::endl;
    std::cout << "       /  /^    agent provocateur" << std::endl;
    std::cout << "      o--o      =================" << std::endl;
    std::cout << std::endl;

    std::cout << "[~>] Read in input data" << std::endl;

  }

  Parser parser;
  parser.read_options("input/options.input");

  if (rank == 0) {

    std::cout << "[~~>] Number of agents: " << parser.n_agent << std::endl;
    std::cout << "[~~>] Number of iterations: " << parser.n_iteration << std::endl;
    std::cout << "[~~>] Random seed: " << parser.seed << std::endl;
    std::cout << "[~~>] Geometric fuzz: " << parser.fuzz << std::endl;

    std::cout << "[~~>] Surface mesh file: " << parser.surface_mesh_file << std::endl;
    std::cout << "[~~>] Surface data file: " << parser.surface_data_file << std::endl;
    std::cout << "[~~>] Surface time: ";
    for (uint i = 0; i < parser.surface_time.size(); i ++)
      std::cout << parser.surface_time[i] << " ";
    std::cout <<std::endl;
    std::cout << "[~~>] Surface group: ";
    for (uint i = 0; i < parser.surface_data_group.size(); i ++)
      std::cout << parser.surface_data_group[i] << " ";
    std::cout <<std::endl;

    std::cout << "[~~>] Subsurface mesh file: " << parser.subsurface_mesh_file << std::endl;
    std::cout << "[~~>] Subsurface data file: " << parser.subsurface_data_file << std::endl;
    std::cout << "[~~>] Subsurface time: ";
    for (uint i = 0; i < parser.subsurface_time.size(); i ++)
      std::cout << parser.subsurface_time[i] << " ";
    std::cout <<std::endl;
    std::cout << "[~~>] Subsurface group: ";
    for (uint i = 0; i < parser.subsurface_data_group.size(); i ++)
      std::cout << parser.subsurface_data_group[i] << " ";
    std::cout << std::endl;
    std::cout << std::endl;
    
    std::cout << "[~>] Building the subsurface mesh" << std::endl;
    
  }

  Mesh mesh = Mesh();

  mesh.set_fuzz(parser.fuzz);
  
  H5::H5File fp_subsurface = H5::H5File(parser.subsurface_mesh_file.c_str(), H5F_ACC_RDONLY);
  H5::H5File dp_subsurface = H5::H5File(parser.subsurface_data_file.c_str(), H5F_ACC_RDONLY);
  H5::H5File fp_surface = H5::H5File(parser.surface_mesh_file.c_str(), H5F_ACC_RDONLY);
  H5::H5File dp_surface = H5::H5File(parser.surface_data_file.c_str(), H5F_ACC_RDONLY);
  
  mesh.build_subsurface_mesh(fp_subsurface);
  fp_subsurface.close();

  if (rank == 0) {
    std::cout << "[~~>] Number of cells: " << mesh.get_n_cells(SUBSURFACE_FLAG) << std::endl;
    std::cout << std::endl;
    
    std::cout << "[~>] Building the surface mesh (includes coupling with subsurface and may take a while)" << std::endl;
  }
  
  mesh.build_surface_mesh(fp_surface, dp_surface, parser.surface_data_group[0]);
  fp_surface.close();

  if (rank == 0) {

    int len_surf = mesh.get_n_cells(SURFACE_FLAG);

    std::cout << "[~~>] Number of cells: " << len_surf << std::endl;
    std::cout << std::endl;
    
    FILE *mesh_fp = fopen("mesh.csv", "w");
    fprintf(mesh_fp, "x (m), y (m), z (m), index\n");
    Cell *c;
    
    for (int n = 0; n < len_surf; n ++) {
      c = mesh.get_cell(n, SURFACE_FLAG);
      fprintf(mesh_fp, "%f, %f, %f, %d\n",
	      c->get_coordinate(0),
	      c->get_coordinate(1),
	      c->get_coordinate(2), n);
    }
  
    fclose(mesh_fp);

    std::cout << "[~>] Reading the subsurface environment" << std::endl;
  
  }

  // Place agents in the domain and advance N_ITERATION steps

  // Bootstrap 
  srand(parser.seed + rank);
  int len = mesh.get_n_cells(SUBSURFACE_FLAG);
  int index;

  std::cout << "[~>] Releasing " << parser.n_agent << " agents" << std::endl;

  size_t partition = parser.n_agent / size;
  size_t low_bound = rank * partition;
  size_t high_bound = (rank + 1) * partition;
  if (rank == size - 1) // top processor picks up extras
    high_bound = parser.n_agent;
  
  for (size_t i = low_bound; i < high_bound; i ++) {

    // for each agent, start with a fresh environment
    H5::H5File _dp_subsurface = H5::H5File(parser.subsurface_data_file.c_str(), H5F_ACC_RDONLY);
	  
    mesh.build_environment(_dp_subsurface,
			   parser.subsurface_data_group[0],
			   parser.subsurface_data_group[1],
			   SUBSURFACE_FLAG);
    _dp_subsurface.close();

    H5::H5File _dp_surface = H5::H5File(parser.surface_data_file.c_str(), H5F_ACC_RDONLY);
	  
    mesh.build_environment(_dp_surface,
			   parser.surface_data_group[0],
			   parser.surface_data_group[1],
			   SURFACE_FLAG);
    _dp_surface.close();

    // create agent

    Agent agent;
    agent.init(parser.n_iteration);
    
    index = rand()%len;
    agent.set_cell_index(index);
    agent.set_path(mesh.get_cell(index, SUBSURFACE_FLAG)->get_coordinate(0),
		   mesh.get_cell(index, SUBSURFACE_FLAG)->get_coordinate(1),
		   mesh.get_cell(index, SUBSURFACE_FLAG)->get_coordinate(2), 0);
  
    agent.set_domain(SUBSURFACE_FLAG, 0);
    agent.set_domain(SUBSURFACE_FLAG, 1);

    agent.set_path(mesh.get_cell(index, SUBSURFACE_FLAG)->get_coordinate(0),
		   mesh.get_cell(index, SUBSURFACE_FLAG)->get_coordinate(1),
		   mesh.get_cell(index, SUBSURFACE_FLAG)->get_coordinate(2), 1);

    int step;
    int surface_time_counter = 0;
    int subsurface_time_counter = 0;
    int n_levels_subsurface = parser.subsurface_time.size();
    int n_levels_surface    = parser.surface_time.size();

    // if the size of the vectors is larger than 1, there are multiple
    // time levels and the environment is dynamic.  we increase the
    // counter to the next time step to linearly interpolate between
    // values
    if (n_levels_subsurface > 1)
      subsurface_time_counter ++;
    if (n_levels_surface > 1)
      surface_time_counter ++;
    
    double time_horizon_surface    = parser.surface_time[surface_time_counter];
    double time_horizon_subsurface = parser.subsurface_time[subsurface_time_counter];
    
    char output_fname[1024];

    // move agent
    for (step = 1; step < parser.n_iteration; step ++) {

      double rel_dt_surface = (agent.get_time(step) - parser.surface_time[surface_time_counter-1]) / (parser.surface_time[surface_time_counter] - parser.surface_time[surface_time_counter-1]);
      double rel_dt_subsurface = (agent.get_time(step) - parser.subsurface_time[subsurface_time_counter-1]) / (parser.subsurface_time[subsurface_time_counter] - parser.subsurface_time[subsurface_time_counter-1]);

      rel_dt_surface    = std::max(0.0, rel_dt_surface);
      rel_dt_subsurface = std::max(0.0, rel_dt_subsurface);
      
      agent.advance(&mesh, rel_dt_surface, rel_dt_subsurface, step);

      // update environment if necessary
      double current_time = agent.get_time(step);
      double dt_max       = time_horizon_subsurface - current_time;
      agent.set_timestepsize(dt_max);

      if ((step - 1) % (parser.n_iteration / 5) == 0 && i % (parser.n_agent / 5) == 0) {
	std::cout << std::endl;
	std::cout << "Agent #" << i << std::endl;
	std::cout << "Current time: " << current_time << "/" << time_horizon_subsurface << " s";
	std::cout << std::endl;
	std::cout << "Current counter: " << subsurface_time_counter << "/" << n_levels_subsurface - 1 << std::endl;
	std::cout << std::endl;
      }

      // check the subsurface
      if (subsurface_time_counter < n_levels_subsurface - 1) {
	
	// once the simulation time exceeds the time frame of the
	// environment, we assume steady state

	if (current_time >= time_horizon_subsurface) {

	  if (step % (parser.n_iteration / 5) == 0 && i % (parser.n_agent / 5) == 0) {
	    std::cout << "  Advance to group " << parser.subsurface_data_group[subsurface_time_counter] << std::endl;
	    std::cout << std::endl;
	  }
	  
	  time_horizon_subsurface = parser.subsurface_time[subsurface_time_counter+1];
	  
	  H5::H5File _dp = H5::H5File(parser.subsurface_data_file.c_str(), H5F_ACC_RDONLY);
	  mesh.build_environment(_dp,
				 parser.subsurface_data_group[subsurface_time_counter],
				 parser.subsurface_data_group[subsurface_time_counter+1],
				 SUBSURFACE_FLAG);
	  _dp.close();
	  
	  subsurface_time_counter ++;
	  
	}
	
      }

      if (surface_time_counter < n_levels_surface - 1) {

	if (current_time >= time_horizon_surface) {

	  if (step % (parser.n_iteration / 5) == 0 && i % (parser.n_agent / 5) == 0) {
	    std::cout << "  Advance to group " << parser.surface_data_group[surface_time_counter] << std::endl;
	    std::cout << std::endl;
	  }
	  
	  time_horizon_surface = parser.surface_time[surface_time_counter+1];
	  
	  H5::H5File _dp = H5::H5File(parser.surface_data_file.c_str(), H5F_ACC_RDONLY);
	  mesh.build_environment(_dp,
				 parser.surface_data_group[surface_time_counter],
				 parser.surface_data_group[surface_time_counter+1],
				 SURFACE_FLAG);
	  _dp.close();
	  
	  surface_time_counter ++;
	  
	}
	
      }
      
    }

    // write out results
    sprintf(output_fname, "results/agent-%06lu.csv", i);
    FILE *result_fp = fopen(output_fname, "w");

    fprintf(result_fp, "time (s), x (m), y (m), z (m), domain\n");
    for (step = 0; step < parser.n_iteration; step ++) {
      fprintf(result_fp, "%f, %f, %f, %f, %d\n",
	      agent.get_time(step),
	      agent.get_path(step, 0),
	      agent.get_path(step, 1),
	      agent.get_path(step, 2),
	      agent.get_domain(step));
    }
    fclose(result_fp);
      
  }

  MPI_Finalize();
  
  return MPI_SUCCESS;

}