#include "femgrp.h"
#include "debug.hpp"
#include <iostream>
#include <iomanip>
#include <string>
#include <iomanip>
#include "dgtd-performance.hpp"

#ifdef DGTD_USE_OPENMP
  #include <omp.h>
#endif

using std::cout;
using std::cin;
using std::cerr;
using std::endl;
using std::flush; // STDIO
using std::ofstream;
using std::streambuf;
using std::stringstream; // stream
using std::string;


#if (defined DGTD_USE_CUDA)
  void CUBLAS_ROUTINE(FemGrp& FEM){
    FEM.GetMatrices(); 
    cout << "GetMatrices" << endl;
    FEM.PrepareGPUcuBLAS();
    cout << "GPU set up correctly" << endl;
    FEM.TimeSteppingCuBLAS();
  }
#else
  void CPU_ROUTINE(FemGrp& FEM){
    FEM.GetMatrices(); 
    cout << "GetMatrices" << endl; // [M]_eps, [M]_mu
    fflush(stdout);
    MEM_USAGE();

    cout<<"-----------------------------------------Begin LTS_TimeUpdateGlobal_MatrixFree------------------------------------\n";
    FEM.LTS_TimeUpdateGlobal_MatrixFree();
    cout<<"----------------------------------Done LTS_TimeUpdateGlobal_MatrixFree------------------------------------\n";
  }
#endif

// thread-level parallel, OpenMP only, MPI-free
int main(int argc, char **argv){
  timer_start("The Maxwell TD simulation", ' '); // seconds
  FemGrp FEM;
  
  int threads = 0;

  cout << " " << endl;
  #ifdef DGTD_USE_OPENMP
    omp_set_dynamic(0);
    omp_set_nested(0);
    omp_set_num_threads(12);
    #pragma omp parallel
    {
      #pragma omp master
      {
        threads = omp_get_num_threads();
        cout << "===================== " << endl;
        cout << " OpenMP Threads = "  << threads << endl;
        cout << "===================== " << endl;
      }
    }
  #else
    threads = 1;
    cout << "===========" << endl;
    cout << " No OpenMP " << endl;
    cout << "===========" << endl;
  #endif
  cout << " " << endl;

  MEM_USAGE();

  // Input parameters
  bool WriteFields = false;
  bool WriteProbes = false;
  bool WriteSurfFlag = false;
  bool write_AnalyticalIncidentProbes = false;
  bool regularAreas = false;
  bool usePadeBool = false;
  int padeCNT = 0;
  fp_t padeTime = -1;
  fp_t enegyDecayFactor = 0.0;
  int numberOfEnergyPoints = 0;

  // ================================================================================
  // 2. Read inputs from arguments
  if(argc < 18)
  {
    cout << "==================================================================================================================" << endl;
    cout << " Not Enough input arguments" << endl;
    cout << " [EXE] FileName Freq ExcitationType WaveformType Tdelay Twidth FinalTime FluxType PolyOrder MatrixFree ScalaStudy " << endl;
    cout << " 1. Filename: Simulation Filename" << endl;
    cout << " 2. Freq: Central Modulation Frequency (MHz)" << endl;
    cout << " 3. Planewave Waveform: (0)Time Harmonic (1)Gauss (2)Neumann (3) Modulated Gauss" << endl;
    cout << " 4. Port Waveform (0)TH (1)Gauss" << endl;
    cout << " 5. Tdelay (sec): Delay of excitation pulse" << endl;
    cout << " 6. Tau (sec): Parameter for pulse width" << endl;
    cout << " 7. FinalTime (sec): Termination Time (Set to -1 to activate energy convergence termination)" << endl;
    cout << " 8. Gamma " << endl;
    cout << " 9. VTU Output Flag (0) True and (1) False" << endl;
    cout << " 10. Flag to Write Currents / Fields on Surface (0) Off (1) On" << endl;
    cout << " 11. Poly Order: (1) First order (2) Second order" << endl;
    cout << " 12. Free Space Comparison: (0) True and (1) False" << endl;
    cout << " 13. Sampling Rate: sampling frequency / nyquist frequency" << endl;
    cout << " 14. Padé Acceleration: Number of points used to calculate the Pade Aprox" << endl;
    cout << " 15. Time to perform Pade in ns (set to -1 if want to do convergence study)" << endl;
    cout << " 16. Energy Convergence Factor (%): factor in % required to the energy to decay finish the simulation" << endl;
    cout << " 17. Energy Convergence Number of Points: number of Points used to the determine energy convergence (set to 0 for all)" << endl;
    cout << "==================================================================================================================" << endl;
    return 0;
  }

  // Input
  string FileName = argv[1];
  fp_t Freq = atoi(argv[2]);

  int ExcitationType = atoi(argv[3]);
  int WaveformType = atoi(argv[4]);
  fp_t Tdelay = atof(argv[5]);
  fp_t Tau = atof(argv[6]);
  fp_t FinalTime = atof(argv[7]);
  fp_t gamma = atof(argv[8]);

  if(atoi(argv[9]) == 0){
    cout << "WriteFields activated" << endl;
    WriteFields = true;
  }

  int WriteSurfBCFlag = atoi(argv[10]);


  //We obtain the basis function order
  int PolyOrderFlag;

  if(atoi(argv[11]) == 1 || atoi(argv[11]) == 2 || atoi(argv[11]) == 3){
    PolyOrderFlag = atoi(argv[11]);
  }else{
    cout << "==================================================================" << endl;
    cout << " 12. Poly Order: (1) First order (2) Second order (3) Third order" << endl;
    cout << "==================================================================" << endl;
    return 0;
  }

  if(atoi(argv[12]) == 0){
    cout << "Free Space Comparison activated" << endl;
    write_AnalyticalIncidentProbes = true;
  }

  fp_t SamplingRate = atof(argv[13]);

  if(WriteSurfBCFlag > 0){
    cout << "-------------------------" << endl;
    cout << " WriteSurfFlag activated " << endl;
    WriteSurfFlag = true;
    cout << "-------------------------" << endl;
  }

  ifstream ifile;
  ifile.open(FileName + ".probe");
  if(ifile){
    cout << ".probe file exists" << endl;
    WriteProbes = true;
    padeCNT = atoi(argv[14]);
    if(atoi(argv[14]) > 0 && FinalTime > 0){
      cout << "Pade Acceleration activated" << endl;
      usePadeBool = true;
      padeTime = atof(argv[15]);
    }else{
      padeCNT = 0;
      padeTime = -1.0;
    }
  }else{
    cout<<".probe file is missing" << endl;
    if(atoi(argv[14]) > 0){
      cout << "Pade Acceleration not activated because of the lack of probes " << endl;
      usePadeBool = false;
      padeCNT = 0;
      padeTime = -1.0;
    }
    return 0;
  }
  ifile.close();

  if(FinalTime <= 0){
    enegyDecayFactor = atof(argv[16]);
    numberOfEnergyPoints = atoi(argv[17]);
  }

  ifile.open(FileName + ".regular");
  if(ifile){
    cout << ".regular file exists" << endl;
    regularAreas = true;
  }else{
    cout << FileName <<".regular file is missing" << endl;
  }
  ifile.close();

  // Info
  int output_width = 50;
  cout.setf(ios::scientific,ios_base::floatfield);
  cout.precision(10);
  cout << std::setiosflags(std::ios::left) << std::setw(output_width) << std::boolalpha;
  cout << " " << endl;
  cout << "===================================================================" << endl;
  cout << "                            INPUT DATA                             " << endl;
  cout << "===================================================================" << endl;
  cout << std::setw(output_width) << " Penalty: " << gamma << endl;
  cout << std::setw(output_width) << " Basis Polynomial Order: " << 1 << endl;
  cout << std::setw(output_width) << " PlaneWave Excitation (0)TH (1)Gauss (2)Neuman: " << ExcitationType << endl;
  cout << std::setw(output_width) << " Port Waveform (0)TH (1)Gauss: " << WaveformType << endl;
  cout << std::setw(output_width) << " Group File Name: " << FileName << endl;
  cout << std::setw(output_width) << " Maximum Frequency (MHz): " << Freq << "" << endl;
  cout << std::setw(output_width) << " Tdelay: " << Tdelay << endl;
  cout << std::setw(output_width) << " Tau: " << Tau << endl;
  cout << std::setw(output_width) << " Final Time: " << FinalTime << endl;
  cout << std::setw(output_width) << " Polynomial basis order: " << PolyOrderFlag << endl;
  cout << std::setw(output_width) << " Regular Mesh File: " << regularAreas << endl;

  cout << std::setw(output_width) << " WriteFields Mode: " << WriteFields << endl;
  cout << std::setw(output_width) << " WriteProbes Mode: " << WriteProbes << endl;
  cout << std::setw(output_width) << " write_AnalyticalIncidentProbes Mode: " << write_AnalyticalIncidentProbes << endl;

  cout << std::setw(output_width) << " Sampling Rate: " << SamplingRate << endl;

  #if (defined DGTD_USE_CUDA)
    cout << std::setw(output_width) << " GPU Mode (FinalTime): " << FinalTime << endl;
  #endif

  cout << "===================================================================" << endl;
  cout << " " << endl;

  // Reset Folder
  int a;
  a = system(" rm *.vtu Port.domain");
  a = system(" rm *.uvtr");
  a = system(" rm *.TD_Vinc *.TD_Vrefl *.TD_Vtot");
  a = system(" rm *.curJ *.curM");  // Ensure that .tri file is not removed
  a = system(" rm *.TDerrorE *.TDerrorH");
  a = system(" rm *.FD_errorH *.FD_errorE");
  a = system(" rm *.TDerrorEPlot *.TDerrorHPlot");

  // In the results reported in the papers, gamma can be chosen to be a small number 0.025 or 0.1 to have both (attenuating spurious mode) and (significantly reduced numerical dissipation).
  FEM.setFluxFactor(gamma);  
  FEM.setPolyFlag(PolyOrderFlag);//set the Basis Order used in the analysis
  FEM.setWriteFields(WriteFields);
  FEM.setWriteProbes(WriteProbes);
  FEM.setWriteSurfFlag(WriteSurfFlag);
  FEM.setwriteAnalyticalIncidentProbes(write_AnalyticalIncidentProbes);
  FEM.setSamplingRate(SamplingRate);
  FEM.setPade(usePadeBool);
  FEM.setPadeCNT(padeCNT);
  FEM.setPadeTime(padeTime);
  FEM.setEnergyDecayFactor(enegyDecayFactor / 100);
  FEM.setNumberOfEnergyPoints(numberOfEnergyPoints);
    
  #ifdef DGTD_USE_LTS
    FEM.setscalbSty(0); // 0 to turn on LTS
  #else
    FEM.setscalbSty(1); // 1 to turn off LTS    
  #endif
  
  FEM.setTimeDist(WaveformType);
  FEM.setExciFlag(ExcitationType);
  FEM.set_fname(const_cast<char*>(FileName.c_str()));
  FEM.set_freq(Freq);
  FEM.setTo(Tdelay);
  FEM.setTau(Tau);
  FEM.setFinalTime(FinalTime);

  // read geometry model
  timer_start("DGTD: read geometry model", 'm');
  FEM.readNODE(); 
  cout << "readNODE " << endl;
  FEM.readBC(); 
  cout << "readBC " << endl;
  FEM.readTETRA(); 
  cout << "readTETRA " << endl;
  FEM.readMaterial(); 
  cout << "readMaterial " << endl;


  MEM_USAGE();
  cout << " " << endl;
  cout << "===========================" << endl;
  cout << "        Make Arrays        " << endl;
  cout << "===========================" << endl;
  FEM.makeEdgeArray(); 
  cout << " makeEdgeArray " << endl;
  cout << " " << endl;
    
  FEM.makeFaceArray(); 
  cout << " makeFaceArray " << endl;
  
  if(FEM.getNonConformal())
  {
    FEM.makeNonConformalArray();
    cout << " makeNonConformalArray: " << FEM.getNCCnt() << " tetras" << endl;
    cout << " " << endl;
  } 
  cout << "===========================" << endl;
  cout << " " << endl;

  MEM_USAGE();



  // This is the point when the tetrahedral mesh is ready
  // Added by Qi Jian 
  // Create octree
  cout << "Initialize Octree that organize teterhedrals elements in space" << endl;
  FEM.initializeOctree(FileName, FEM.getNonConformal());

  MEM_USAGE();


  if(WriteProbes)
  {
    auto start = std::chrono::high_resolution_clock::now();
    cout << "readPROBE " << endl;
    FEM.readPROBE(); 
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> duration = end - start;
    std::cout << "readPROBE took " << duration.count() << " seconds." << std::endl;    
  }

  if(regularAreas){
    FEM.setRegularRegionFlag(true);
    FEM.readREGULAR(); 
    cout << "readRegular " << endl;
  }


  // Create Surface for output fields or currents
  if(WriteSurfBCFlag > 0)
  {
      cout << endl;
      cout << "--------------------------------------" << endl; 
      cout << "Write Surface Fields / Currents" << endl;
      cout << "makeOutputSurfMesh" << endl;
      auto start = std::chrono::high_resolution_clock::now();      
      FEM.makeOutputSurfMesh(FileName);
      auto end = std::chrono::high_resolution_clock::now();
      std::chrono::duration<double> duration = end - start;
      std::cout << "makeOutputSurfMesh took " << duration.count() << " seconds." << std::endl;      
      cout << "--------------------------------------" << endl;     
  }
  
  if (FEM.PlaneWaveBCFlag == true){
    cout << " " << endl;
    cout << "======================================== " << endl;
    cout << "         PlaneWave BC Detected           " << endl;
    cout << "======================================== " << endl;
    FEM.makeInterSurfMesh(pecType); 
    cout << "   makeInterSurfMesh" << endl;
    FEM.makePlaneWaveMesh(); 
    cout << "   makePlaneWaveMesh" << endl;
    cout << "======================================== " << endl;
    cout << " " << endl;
  }

  // Port Procedures
  if(FEM.PortBCFlag == true)
  {    
    cout << "======================================== " << endl;
    cout << "           Port BC Detected            " << endl;
    cout << "======================================== " << endl;
    FEM.makePortMeshes(); 
    cout << "  makePortMeshes"<<endl;
    
    if (FEM.PEC_PMC_port_flag != 1)
    {
      FEM.solveWaveguidePorts(); 
      cout << "  solveWaveguidePorts DONE" << endl;
      FEM.WriteWaveguidePortFields(); 
      cout << "  WriteWaveguidePortFields " << endl;
    }
    cout << "======================================== " << endl;
    cout << " " << endl;
  }

  MEM_USAGE();

  FEM.AssignExcitParamToFace(); 
  cout << "AssignExcitParamToFace " << endl;
  FEM.AssignMaterialProperties(); 
  cout << "AssignMaterialProperties" << endl;
  FEM.AssignTetraFlags(); 
  cout << "AssignTetraFlags" << endl;
  
  //TODO: review if correct for higher order functions
  if (FEM.PortBCFlag == true)
  {
    if (FEM.PEC_PMC_port_flag == 1)
    {
      FEM.AssignPortFieldsInFaces_TEM(); 
      cout << "AssignPortFieldsInFaces_TEM" << endl;
    }
    else
    {
      FEM.AssignPortFieldsInFaces(); 
      cout << "AssignPortFieldsInFaces" << endl;
    }
  }

  timer_stop('m');

  timer_start("DGTD: prepare for computation", 'm');
  FEM.DG_AssignOffsets(); 
  cout << "DG_AssignOffsets" << endl;
  FEM.SetUpMatrixVector(); 
  cout << "SetUpMatrixFree" << endl;
  FEM.numberDofs(); 
  cout << "numberDofs" << endl;
  FEM.LocalTimeSteppingClassPartioning(); 
  cout << "LocalTimeSteppingClassPartioning" << endl;
  MEM_USAGE();
  timer_stop('m');

  #if (defined DGTD_USE_CUDA)
    // DEVICE
    CUBLAS_ROUTINE(FEM);  
    
  #else
    // CPU
    cout << "Using CPU_ROUTINE" << endl;
    CPU_ROUTINE(FEM);
  #endif

  timer_stop(' ');
  return 0;
}