#include <stdio.h>
#include <mpi.h>

#include <iostream>
#include <armadillo>
#include <math.h>
#include "dG_library.h"
#include "matrix_utils.h"
#include "dG_parallel.h"

using namespace std;
using namespace arma;
using namespace dG_global;

// advekce 2D
mat u;
double a_advec, b_advec;

void Problem()
{
  T = 10;
  a_advec = 1;
  b_advec = 1;
  
  filename = directoryname + "P_for_h=" + mesh_size_parameter_string + ".txt";
  P.load(filename);
  Px = P.row(0);
  Py = P.row(1);
  filename = directoryname + "E_for_h=" + mesh_size_parameter_string + ".txt";
  E.load(filename);
  E = E - 1; // kvuli cislovani od 0 v C++
}

void BuildBCMaps()
{
  vec dotproducts = a_advec * n_x(mapB) + b_advec * n_y(mapB);
  mapI = mapB(find(dotproducts < -1e-12));
  vmapI = vmapM(mapI);
  mapO = mapB(find(dotproducts > -1e-12));
  vmapO = vmapM(mapO);
}

int main(int argc, char* argv[])
{
  if (argc != 3)
  {
    cout << "Usage: " << argv[0] << " ORDER_OF_LOCAL_APPROXIMATION MESH_SIZE_PARAMETER" << endl;
    
    return 1;
  }
  
  // pomocne a debugovaci promenne
  char str_from_sprintf_buffer[64];
  string Nplus1_string;
  
  wall_clock timer; // mereni delky vypoctu
  double elapsed_time;
  
  int myrank, size;
  // Initialization of MPI
  MPI_Init(&argc, &argv);
  // myrank will contain the rank of the process
  MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
  // size will contain the number of processes
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  
  timer.tic();
  // parametry diskretizace
  N = atoi(argv[1]);
  mesh_size_parameter = atof(argv[2]);
  DiscretizationParamatersToStrings();
  
  directoryname = "dG_2D_advection_parallel_data/"; // tj. ze ktere slozky se maji nacitat data (uloha Advection 2D parallel)
  
  Problem();
  
  StartUp2D(); // StartUp2D jako z MATLABu
  
  // urceni delky kroku
  sprintf(str_from_sprintf_buffer, "%d", N + 1);
  Nplus1_string = string(str_from_sprintf_buffer);
  filename = "dG_general_data/GLnodes_for_n=" + Nplus1_string + ".txt";
  vec GLnodes; GLnodes.load(filename); double r_min = abs(GLnodes(1) - GLnodes(0));
  dtscale2D(); double CFL = 0.75; dt = CFL * min(dtscale) * (2.0 / 3.0) * r_min / sqrt(pow(a_advec,2) + pow(b_advec,2));
  N_steps = (int)ceil(T / dt); dt = T / N_steps;
  
  elapsed_time = timer.toc();
  if(myrank == 0)
  {
    printf("General preprocessing on process %d took %f to complete\n",myrank,elapsed_time);
  }
  
  
  
  timer.tic();
  // rozdeleni dat podle jednotlivych casti vypocetni site - zmeni se K jiz na lokalni K atd.
  ParallelStartUp();
  
  urowvec local_element_IDs = partitions(myrank,0);
  K = local_element_IDs.n_cols;
  /*if(myrank == 0)
  {
    cout << "lokalne mame " << K << " elementu" << endl;
  }*/
  
  x = x.cols(local_element_IDs);
  y = y.cols(local_element_IDs);
  r_x = r_x.cols(local_element_IDs);
  r_y = r_y.cols(local_element_IDs);
  s_x = s_x.cols(local_element_IDs);
  s_y = s_y.cols(local_element_IDs);
  Fx = Fx.cols(local_element_IDs);
  Fy = Fy.cols(local_element_IDs);
  n_x = n_x.cols(local_element_IDs);
  n_y = n_y.cols(local_element_IDs);
  Fscale = Fscale.cols(local_element_IDs);
  
  mapM = mapM_local(myrank,0);
  vmapM = vmapM_local(myrank,0);
  mapP = mapP_local(myrank,0);
  vmapP = vmapP_local(myrank,0);
  mapB = mapB_local(myrank,0);
  vmapB = vmapB_local(myrank,0);
  mapM_inner = mapM_inner_local(myrank,0);
  vmapM_inner = vmapM_inner_local(myrank,0);
  mapP_inner = mapP_inner_local(myrank,0);
  vmapP_inner = vmapP_inner_local(myrank,0);
  mapM_interfaces = mapM_interfaces_local(myrank,0);
  vmapM_interfaces = vmapM_interfaces_local(myrank,0);
  neighbouring_partitions_IDs = neighbouring_partitions_IDs_local(myrank,0);
  neighbouring_partitions_IDs_unique = neighbouring_partitions_IDs_unique_local(myrank,0);
  permutation_vector = permutation_vector_local(myrank,0);
  
  int number_of_neighbouring_partitions = neighbouring_partitions_IDs_unique.n_rows;
  ivec current_neighbouring_partition_vec = zeros<ivec>(number_of_neighbouring_partitions);
  field<uvec> indeces_current_neighbouring_partition_cell = field<uvec>(number_of_neighbouring_partitions,1);
  ivec number_current_processed_nodes_vec = zeros<ivec>(number_of_neighbouring_partitions);
  vec u_interface_send, u_interface_receive;
  stdvec u_interface_send_std, u_interface_receive_std;
  mat u_neighbouring_partitions = zeros<mat>(mapM_interfaces.n_rows,1);
  
  for(int j = 0; j < number_of_neighbouring_partitions; j++)
  {
    current_neighbouring_partition_vec(j) = neighbouring_partitions_IDs_unique(j);
    indeces_current_neighbouring_partition_cell(j,0) = find(neighbouring_partitions_IDs == current_neighbouring_partition_vec(j));
    number_current_processed_nodes_vec(j) = indeces_current_neighbouring_partition_cell(j,0).n_rows;
  }
  elapsed_time = timer.toc();
  if(myrank == 0)
  {
    printf("Parallelization preprocessing on process %d took %f to complete\n",myrank,elapsed_time);
  }
  
  
  // ODTUD JIZ KOD LOKALNE NA KAZDEM PC
  BuildBCMaps(); // jiz lokalne se napocita mapI, vmapI, mapO a vmapO
  
  timer.tic();
  // Timestepping
  mat u_temp = zeros<mat>(N_p,K);
  
  // pomocne promenne pro ODEs_RHS
  mat u_rhs = zeros<mat>(N_p,K);
  mat flux;
  mat numerical_flux = zeros<mat>(N_faces * N_fp,K);
  mat x_in = x(vmapI);
  mat y_in = y(vmapI);
  mat u_inflow;
  mat u_r, u_s, u_x, u_y;
  
  // pocatecni podminka
  double b = 1; // pro vetsi ulohu voleno b = 3
  u = exp((-log(2) / pow(b,2)) * (square(x) + square(y)));
  double current_time = 0;
  double time_RK;
  for(int timestep = 1; timestep <= N_steps; timestep++)
  {
    for(int i = 1; i <= 5; i++)
    {
      // synchronizace procesu
      MPI_Barrier(MPI_COMM_WORLD);
      
      time_RK = current_time + rk4c(i - 1) * dt;
      
      
      // u_rhs = ODEs_RHS(u,time_RK);
      // predani dat mezi procesory - jedina vyrazna zmena oproti kodu pro 1 procesor
      for(int j = 0; j < number_of_neighbouring_partitions; j++)
      {  
        // pres funkci Sendrecv - pry je to lepsi, kdyz se posilaji vzajemne informace a clovek nemusi davat pozor na razeni informaci
        u_interface_send = u(vmapM_interfaces(indeces_current_neighbouring_partition_cell(j,0)));
        u_interface_send_std = conv_to<stdvec>::from(u_interface_send);
        u_interface_receive_std = conv_to<stdvec>::from(u_interface_send);
        MPI_Sendrecv(&u_interface_send_std[0], number_current_processed_nodes_vec(j), MPI_DOUBLE, current_neighbouring_partition_vec(j), 0, &u_interface_receive_std[0], number_current_processed_nodes_vec(j), MPI_DOUBLE, current_neighbouring_partition_vec(j), 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        u_interface_receive = conv_to<vec>::from(u_interface_receive_std);
        u_neighbouring_partitions(indeces_current_neighbouring_partition_cell(j,0)) = u_interface_receive;
      }
      u_neighbouring_partitions = u_neighbouring_partitions(permutation_vector); // prehazeni prijatych hodnot tak, aby body, kterym tyto hodnoty nalezi na jinem procesoru, geometricky nalehaly na body na okraji vypocetni podoblasti na konkretnim procesoru
      
      
      // zbytek kodu je jako predtim, pouze kontrola, zda mapI nebo mapO nejsou prazdne vektory
      flux = (a_advec * u(vmapM)) % vectorise(n_x) + (b_advec * u(vmapM)) % vectorise(n_y);
      
      // numericky tok v bodech, ktere nesousedi s body z jineho procesoru
      numerical_flux(mapM_inner) = ((a_advec * (u(vmapM_inner) + u(vmapP_inner))) % n_x(mapM_inner) + (b_advec * (u(vmapM_inner) + u(vmapP_inner))) % n_y(mapM_inner) + (u(vmapM_inner) - u(vmapP_inner)) % abs(a_advec * n_x(mapM_inner) + b_advec * n_y(mapM_inner))) / 2;
      
      // numericky tok v bodech, ktere sousedi s body z jineho procesoru
      numerical_flux(mapM_interfaces) = ((a_advec * (u(vmapM_interfaces) + u_neighbouring_partitions)) % n_x(mapM_interfaces) + (b_advec * (u(vmapM_interfaces) + u_neighbouring_partitions)) % n_y(mapM_interfaces) + (u(vmapM_interfaces) - u_neighbouring_partitions) % abs(a_advec * n_x(mapM_interfaces) + b_advec * n_y(mapM_interfaces))) / 2;
      
      if(!mapI.is_empty())
      {
        u_inflow = exp((-log(2) / pow(b,2)) * (square(x_in - a_advec * time_RK) + square(y_in - b_advec * time_RK)));
        numerical_flux(mapI) = ((a_advec * (u(vmapI) + u_inflow)) % n_x(mapI) + (b_advec * (u(vmapI) + u_inflow)) % n_y(mapI) + (u(vmapI) - u_inflow) % abs(a_advec * n_x(mapI) + b_advec * n_y(mapI))) / 2;
      }
      
      if(!mapO.is_empty())
      {
        numerical_flux(mapO) = (a_advec * u(vmapO)) % n_x(mapO) + (b_advec * u(vmapO)) % n_y(mapO);
      }
      
      // prerovnani flux do matice - numerical_flux netreba prerovnavat, byl naalokovan jako matice a je pouzito linearni indexovani
      flux.reshape(N_fp * N_faces,K);
	  
	  u_r = D_r * u; u_s = D_s * u; u_x = r_x % u_r + s_x % u_s; u_y = r_y % u_r + s_y % u_s; // Grad2D
	  u_rhs = -a_advec * u_x - b_advec * u_y + LIFT * (Fscale % (flux - numerical_flux));
      
      
      u_temp = rk4a(i - 1) * u_temp + dt * u_rhs;
	  u = u + rk4b(i - 1) * u_temp;
    }
    current_time = current_time + dt;
    
    if(myrank == 0)
    {
      if(timestep % 100 == 0)
      {
        //cout << timestep << " of total " << N_steps << " timesteps processed." << endl;
      
        elapsed_time = timer.toc();
        cout << timestep << " of total " << N_steps << " timesteps processed; elapsed " << elapsed_time << " seconds so far" << endl;
      }
    }
  }
  
  elapsed_time = timer.toc();
  if(myrank == 0)
  {
    printf("Computation in process %d took %f to complete\n",myrank,elapsed_time);
  }
  
  filename = "advekce_2D_u_spoctene_pro_N=" + N_string + "_a_h=" + mesh_size_parameter_string + "_partition" + to_string(myrank) + ".txt";
  u.save(filename, raw_ascii);
  
  mat u_exact = exp((-log(2) / pow(b,2)) * (square(x - a_advec * T) + square(y - b_advec * T)));
  mat errors = abs(u - u_exact);
  double E_max = max(max(errors));
  double E_avg = accu(errors) / (N_p * K);
  printf("Process %d reports E_max = %e and E_avg = %e\n",myrank,E_max,E_avg);
  
  MPI_Finalize();
  
  return 0;
}
