#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <algorithm>
#include "ChemPlugin.h"
#include <mpi.h>

int exit_client(int status)
{
   if (status != 0)
      MPI_Abort(MPI_COMM_WORLD, status);
   MPI_Finalize();
   return status;
}

void open_input(std::ifstream& input, int argc, char** argv) {
   std::string filename;
   if (argc < 2) {
      std::cerr << "You must specify an input file as the first program argument." << std::endl;
      exit_client(-1);
   }
   else {
      filename = argv[1];
   }

   input.open(filename);

   if (!input.is_open()) {
      std::cerr << "The input file does not exist" << std::endl;
      exit_client(-1);
   }
}

void write_results(std::vector<ChemPlugin>& cp, int nx, double gap, double& then)
{
   double now = cp[0].MpiReport1("Time", "years");
   int step = nx / 100;
   if ((then - now) < gap / 1e4) {
      int last_rank = cp[0].MpiRank();
      for (int i=0; i<nx; i+=step) {
         if (cp[i].MpiRank() != last_rank) {
            last_rank = cp[i].MpiRank();
            MPI_Barrier(MPI_COMM_WORLD);
         }
         std::string label = "Instance " + std::to_string(i);
         cp[i].PrintOutput("RTM.txt", label);
      }
      then += gap;
   }
}

int main(int argc, char** argv) {
   // Initialize MPI
   MPI_Init(NULL, NULL);
   // Get the number of processes
   int world_size;
   MPI_Comm_size(MPI_COMM_WORLD, &world_size);
   // Get the rank of this process
   int world_rank;
   MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);

   if (world_rank == 0)
      std::cout << std::endl << "Model reactive transport in one dimension using MPI" << std::endl << std::endl;

   // Simulation parameters.
   int nx = 4000; // number of instances along x
   double length = 100; // m
   double deltax = length / nx; // m
   double deltay = 1.0, deltaz = 1.0; // m

   double time_end = 10.0;  // years
   double delta_years = time_end / 5; // years
   double next_output = 0.0;  // years

   // Create the inlet and interior instances.
   ChemPlugin cp_inlet;
   std::vector<ChemPlugin> cp(nx);

   // Assign instances to ranks
   cp_inlet.MpiAssign(0);
   for (int i=0; i<nx; i++)
      cp[i].MpiAssign(i * world_size / nx);

   cp_inlet.Console("stdout");
   cp[0].Console("stdout");
   cp[0].Config("pluses = banner");

   std::string cmd = "volume = " + std::to_string(deltax * deltay * deltaz) +
                     " m3; time end = " + std::to_string(time_end) + " years";
   for (int i=0; i<nx; i++)
      cp[i].Config(cmd);

   // Set input file and configure inlet, initial domain.
   std::ifstream input;
   open_input(input, argc, argv);
   int scope = 0;
   while (!input.eof()) {
      std::string line;
      std::getline(input, line);
      if (line == "go")
         break;
      else if (line == "scope inlet")
         scope = 1;
      else if (line == "scope initial")
         scope = 2;
      else if (scope == 1)
         cp_inlet.Config(line);
      else if (scope == 2)
         for (int i=0; i<nx; i++)
            cp[i].Config(line);
   }

   // Initialize the inlet and interior instances; write out initial conditions.
   if (cp_inlet.Initialize()) {
      std::cout << "Inlet failed to initialize" << std::endl;
      return exit_client(-1);
   }
   cp_inlet.PrintOutput("RTM.txt", "Inlet fluid");

   for (int i=0; i<nx; i++) {
      if (cp[i].Initialize()) {
         std::cout << "Instance " << i << " failed to initialize" << std::endl;
         return exit_client(-1);
      }
   }
   write_results(cp, nx, delta_years, next_output);

   // Set velocity; calculate flow rate and transmissivities.
   double veloc_in; // m/yr
   if (argc < 3) {
      std::cerr << "You must specify fluid velocity in m/yr as the second program argument." << std::endl;
      exit_client(-1);
   }
   veloc_in = atof(argv[2]);
   
   double velocity = veloc_in / 31557600.;  // m/s
   double porosity = cp[0].MpiReport1("porosity");  // volume fraction
   double flow = deltay * deltaz * porosity * velocity;  // m3/s

   double diffcoef = 1e-10;  // m2/s
   double dispersivity = 1.0; // m
   double dispcoef = velocity * dispersivity + diffcoef; // m2/s
   double trans = deltay * deltaz * porosity * dispcoef / deltax;  // m3/s

   double tcond = 2.0;  // W/m/K
   double ttrans = deltay * deltaz * tcond / deltax;  // W/K

   // Link the instances.
   CpiLink link = cp[0].Link(cp_inlet);
   link.Transmissivity(trans, "m3/s");
   link.HeatTrans(ttrans, "W/K");
   link.FlowRate(flow, "m3/s");

   for (int i=1; i<nx; i++) {
      link = cp[i].Link(cp[i-1]);
      link.Transmissivity(trans, "m3/s");
      link.HeatTrans(ttrans, "W/K");
      link.FlowRate(flow, "m3/s");
   }

   link = cp[nx-1].Link();
   link.FlowRate(-flow, "m3/s");

   // Time marching loop.
   while (true) {
      double deltat = 1e99, deltat_world;
      for (int i=0; i<nx; i++)
         deltat = std::min(deltat, cp[i].ReportTimeStep());
      MPI_Allreduce(&deltat, &deltat_world, 1, MPI_DOUBLE, MPI_MIN, MPI_COMM_WORLD);

      for (int i=0; i<nx; i++)
         if (cp[i].AdvanceTimeStep(deltat_world)) return exit_client(0);

      for (int i=0; i<nx; i++)
         for (int l=0; l<cp[i].nLinks(); l++)
            if (cp[i].Link(l).MpiUpdateLink()) return exit_client(-1);

      for (int i=0; i<nx; i++)
         if (cp[i].AdvanceTransport()) return exit_client(-1);

      for (int i=0; i<nx; i++)
         if (cp[i].AdvanceHeatTransport()) return exit_client(-1);

      for (int i=0; i<nx; i++)
         if (cp[i].AdvanceChemical()) return exit_client(-1);

      write_results(cp, nx, delta_years, next_output);
   }

   // Never gets here.
   return 0;
}