#include "macrodef.h"

#ifdef USE_GPU

#include <cmath>
#include <vector>

#include "bssn_class.h"
#include "bssn_cuda_ops.h"
#include "bssn_gpu.h"
#include "bssn_macro.h"
#include "rungekutta4_rout.h"

void bssn_class::Step_MainPath_GPU(int lev, int YN)
{
#ifdef WithShell
#error "Step_MainPath_GPU currently supports Patch grids only."
#endif

  if (bssn_gpu_bind_process_device(myrank))
  {
    cerr << "GPU device bind failure on MPI rank " << myrank << endl;
    MPI_Abort(MPI_COMM_WORLD, 1);
  }
  bssn_gpu_clear_cached_device_buffers();

  setpbh(BH_num, Porg0, Mass, BH_num_input);

  const double dT_lev = dT * pow(0.5, Mymax(lev, trfls));

#if (MAPBH == 1)
  if (BH_num > 0 && lev == GH->levels - 1)
  {
    compute_Porg_rhs(Porg0, Porg_rhs, Sfx0, Sfy0, Sfz0, lev);
    for (int ithBH = 0; ithBH < BH_num; ithBH++)
    {
      for (int ith = 0; ith < 3; ith++)
        Porg1[ithBH][ith] = Porg0[ithBH][ith] + Porg_rhs[ithBH][ith] * dT_lev;
      if (Symmetry > 0)
        Porg1[ithBH][2] = fabs(Porg1[ithBH][2]);
      if (Symmetry == 2)
      {
        Porg1[ithBH][0] = fabs(Porg1[ithBH][0]);
        Porg1[ithBH][1] = fabs(Porg1[ithBH][1]);
      }
    }
  }

  if (lev == a_lev)
    AnalysisStuff(lev, dT_lev);
#endif

#ifdef With_AHF
  AH_Step_Find(lev, dT_lev);
#endif

  const bool BB = fgt(PhysTime, StartTime, dT_lev / 2);
  (void)BB;
#if (MAPBH == 0)
  const bool need_host_stage_sync = (BH_num > 0 && lev == GH->levels - 1);
#else
  const bool need_host_stage_sync = false;
#endif
  double ndeps = (lev < GH->movls) ? numepsb : numepss;
  double TRK4 = PhysTime;
  int iter_count = 0;
  int pre = 0, cor = 1;
  int ERROR = 0;

  auto run_stage_on_block =
      [&](Block *cg, Patch *patch, MyList<var> *state0_list,
          MyList<var> *boundary_src_list, MyList<var> *stage_data_list,
          MyList<var> *rhs_list, int rk_stage) {
        MyList<var> *varl0 = state0_list;
        MyList<var> *varlb = boundary_src_list;
        MyList<var> *varls = stage_data_list;
        MyList<var> *varlr = rhs_list;

        while (varl0)
        {
          if (bssn_cuda_rk4_boundary_var(cg->shape, dT_lev,
                                         cg->X[0], cg->X[1], cg->X[2],
                                         patch->bbox[0], patch->bbox[1], patch->bbox[2],
                                         patch->bbox[3], patch->bbox[4], patch->bbox[5],
                                         cg->fgfs[varl0->data->sgfn],
                                         cg->fgfs[varlb->data->sgfn],
                                         cg->fgfs[varls->data->sgfn],
                                         cg->fgfs[varlr->data->sgfn],
                                         varl0->data->propspeed,
                                         varl0->data->SoA,
                                         Symmetry, lev, rk_stage, false))
          {
            cerr << "GPU rk4/boundary failure: lev=" << lev
                 << " rk_stage=" << rk_stage
                 << " var=" << varl0->data->name
                 << " bbox=(" << cg->bbox[0] << ":" << cg->bbox[3] << ","
                 << cg->bbox[1] << ":" << cg->bbox[4] << ","
                 << cg->bbox[2] << ":" << cg->bbox[5] << ")" << endl;
            ERROR = 1;
            break;
          }
          varl0 = varl0->next;
          varlb = varlb->next;
          varls = varls->next;
          varlr = varlr->next;
        }
      };

  auto stage_download_var_list =
      [&](Block *cg, MyList<var> *var_list) {
        while (var_list)
        {
          if (bssn_cuda_download_buffer(cg->shape, cg->fgfs[var_list->data->sgfn]))
          {
            cerr << "GPU stage download failure: lev=" << lev
                 << " var=" << var_list->data->name
                 << " bbox=(" << cg->bbox[0] << ":" << cg->bbox[3] << ","
                 << cg->bbox[1] << ":" << cg->bbox[4] << ","
                 << cg->bbox[2] << ":" << cg->bbox[5] << ")" << endl;
            ERROR = 1;
            break;
          }
          var_list = var_list->next;
        }
      };

  auto stage_upload_var_list =
      [&](Block *cg, MyList<var> *var_list) {
        const int n = cg->shape[0] * cg->shape[1] * cg->shape[2];
        while (var_list)
        {
          if (bssn_gpu_stage_upload_buffer(cg->fgfs[var_list->data->sgfn], n))
          {
            cerr << "GPU state upload failure: lev=" << lev
                 << " var=" << var_list->data->name
                 << " bbox=(" << cg->bbox[0] << ":" << cg->bbox[3] << ","
                 << cg->bbox[1] << ":" << cg->bbox[4] << ","
                 << cg->bbox[2] << ":" << cg->bbox[5] << ")" << endl;
            ERROR = 1;
            break;
          }
          var_list = var_list->next;
        }
      };

  auto ensure_stage_device_var_list =
      [&](Block *cg, MyList<var> *var_list) {
        const int n = cg->shape[0] * cg->shape[1] * cg->shape[2];
        while (var_list)
        {
          double *host_ptr = cg->fgfs[var_list->data->sgfn];
          if (!bssn_gpu_find_device_buffer(host_ptr) &&
              bssn_gpu_stage_upload_buffer(host_ptr, n))
          {
            cerr << "GPU state ensure failure: lev=" << lev
                 << " var=" << var_list->data->name
                 << " bbox=(" << cg->bbox[0] << ":" << cg->bbox[3] << ","
                 << cg->bbox[1] << ":" << cg->bbox[4] << ","
                 << cg->bbox[2] << ":" << cg->bbox[5] << ")" << endl;
            ERROR = 1;
            break;
          }
          var_list = var_list->next;
        }
      };

  auto refresh_synced_device_regions =
      [&](Block *cg, MyList<var> *var_list, Parallel::SyncCache &cache) {
        std::vector<Parallel::gridseg *> local_segments;
        for (int node = 0; node < cache.cpusize; ++node)
        {
          MyList<Parallel::gridseg> *seg = cache.combined_dst[node];
          while (seg)
          {
            if (seg->data && seg->data->Bg == cg)
              local_segments.push_back(seg->data);
            seg = seg->next;
          }
        }

        if (local_segments.empty())
          return;

        const int n = cg->shape[0] * cg->shape[1] * cg->shape[2];
        while (var_list)
        {
          double *host_ptr = cg->fgfs[var_list->data->sgfn];
          if (!bssn_gpu_find_device_buffer(host_ptr))
          {
            if (bssn_gpu_stage_upload_buffer(host_ptr, n))
            {
              cerr << "GPU sync refresh upload failure: lev=" << lev
                   << " var=" << var_list->data->name
                   << " bbox=(" << cg->bbox[0] << ":" << cg->bbox[3] << ","
                   << cg->bbox[1] << ":" << cg->bbox[4] << ","
                   << cg->bbox[2] << ":" << cg->bbox[5] << ")" << endl;
              ERROR = 1;
              break;
            }
          }
          else
          {
            for (size_t i = 0; i < local_segments.size(); ++i)
            {
              Parallel::gridseg *seg = local_segments[i];
              if (bssn_gpu_stage_upload_region(host_ptr,
                                               cg->shape,
                                               cg->bbox,
                                               cg->bbox + dim,
                                               seg->shape,
                                               seg->llb))
              {
                cerr << "GPU sync region refresh failure: lev=" << lev
                     << " var=" << var_list->data->name
                     << " bbox=(" << cg->bbox[0] << ":" << cg->bbox[3] << ","
                     << cg->bbox[1] << ":" << cg->bbox[4] << ","
                     << cg->bbox[2] << ":" << cg->bbox[5] << ")" << endl;
                ERROR = 1;
                break;
              }
            }
            if (ERROR)
              break;
          }
          var_list = var_list->next;
        }
      };

  auto refresh_stage_device_after_sync =
      [&](MyList<var> *var_list, Parallel::SyncCache &cache) {
        MyList<Patch> *patch_it = GH->PatL[lev];
        while (patch_it)
        {
          MyList<Block> *block_it = patch_it->data->blb;
          while (block_it)
          {
            Block *cg = block_it->data;
            if (myrank == cg->rank)
              refresh_synced_device_regions(cg, var_list, cache);

            if (block_it == patch_it->data->ble)
              break;
            block_it = block_it->next;
          }
          if (ERROR)
            break;
          patch_it = patch_it->next;
        }
      };

  auto refresh_stage_host_before_sync =
      [&](MyList<var> *var_list, Parallel::SyncCache &cache) -> bool {
        if (!cache.valid || !cache.combined_src || myrank < 0 || myrank >= cache.cpusize)
          return false;

        MyList<Patch> *patch_it = GH->PatL[lev];
        while (patch_it)
        {
          MyList<Block> *block_it = patch_it->data->blb;
          while (block_it)
          {
            Block *cg = block_it->data;
            if (myrank == cg->rank)
            {
              std::vector<Parallel::gridseg *> local_segments;
              MyList<Parallel::gridseg> *seg = cache.combined_src[myrank];
              while (seg)
              {
                if (seg->data && seg->data->Bg == cg)
                  local_segments.push_back(seg->data);
                seg = seg->next;
              }

              if (!local_segments.empty())
              {
                MyList<var> *var_it = var_list;
                while (var_it)
                {
                  double *host_ptr = cg->fgfs[var_it->data->sgfn];
                  for (size_t i = 0; i < local_segments.size(); ++i)
                  {
                    Parallel::gridseg *src_seg = local_segments[i];
                    if (bssn_gpu_stage_download_region(host_ptr,
                                                       cg->shape,
                                                       cg->bbox,
                                                       cg->bbox + dim,
                                                       src_seg->shape,
                                                       src_seg->llb))
                    {
                      cerr << "GPU sync region download failure: lev=" << lev
                           << " var=" << var_it->data->name
                           << " bbox=(" << cg->bbox[0] << ":" << cg->bbox[3] << ","
                           << cg->bbox[1] << ":" << cg->bbox[4] << ","
                           << cg->bbox[2] << ":" << cg->bbox[5] << ")" << endl;
                      ERROR = 1;
                      return true;
                    }
                  }
                  var_it = var_it->next;
                }
              }
            }

            if (block_it == patch_it->data->ble)
              break;
            block_it = block_it->next;
          }
          patch_it = patch_it->next;
        }

        return true;
      };

  auto can_pack_sync_from_device =
      [&](MyList<var> *var_list, Parallel::SyncCache &cache) -> bool {
        if (!cache.valid || !cache.combined_src || myrank < 0 || myrank >= cache.cpusize)
          return false;

        MyList<Parallel::gridseg> *seg = cache.combined_src[myrank];
        while (seg)
        {
          MyList<var> *var_it = var_list;
          while (var_it)
          {
            if (!bssn_gpu_find_device_buffer(seg->data->Bg->fgfs[var_it->data->sgfn]))
              return false;
            var_it = var_it->next;
          }
          seg = seg->next;
        }
        return true;
      };

  MyList<Patch> *Pp = GH->PatL[lev];
  while (Pp)
  {
    MyList<Block> *BP = Pp->data->blb;
    while (BP)
    {
      Block *cg = BP->data;
      if (myrank == cg->rank)
      {
        stage_upload_var_list(cg, StateList);
        if (gpu_rhs(CALLED_BY_STEP, myrank, RHS_PARA_CALLED_FIRST_TIME))
          ERROR = 1;

        run_stage_on_block(cg, Pp->data, StateList, StateList, SynchList_pre, RHSList, iter_count);

        if (bssn_cuda_lowerbound(cg->shape, cg->fgfs[phi->sgfn], chitiny, false))
        {
          cerr << "GPU lowerbound failure: lev=" << lev
               << " rk_stage=" << iter_count
               << " var=" << phi->name
               << " bbox=(" << cg->bbox[0] << ":" << cg->bbox[3] << ","
               << cg->bbox[1] << ":" << cg->bbox[4] << ","
               << cg->bbox[2] << ":" << cg->bbox[5] << ")" << endl;
          ERROR = 1;
        }
        if (!ERROR && !sync_cache_pre[lev].valid)
          stage_download_var_list(cg, SynchList_pre);
      }
      if (BP == Pp->data->ble)
        break;
      BP = BP->next;
    }
    Pp = Pp->next;
  }

  if (!ERROR && sync_cache_pre[lev].valid && !can_pack_sync_from_device(SynchList_pre, sync_cache_pre[lev]))
    refresh_stage_host_before_sync(SynchList_pre, sync_cache_pre[lev]);

  MPI_Request err_req_pre;
  {
    int erh = ERROR;
    MPI_Iallreduce(&erh, &ERROR, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD, &err_req_pre);
  }

  Parallel::AsyncSyncState async_pre;
  Parallel::Sync_start(GH->PatL[lev], SynchList_pre, Symmetry, sync_cache_pre[lev], async_pre);
  Parallel::Sync_finish(sync_cache_pre[lev], async_pre, SynchList_pre, Symmetry, need_host_stage_sync);
  if (!ERROR && need_host_stage_sync)
    refresh_stage_device_after_sync(SynchList_pre, sync_cache_pre[lev]);

  MPI_Wait(&err_req_pre, MPI_STATUS_IGNORE);
  if (ERROR)
  {
    Parallel::Dump_Data(GH->PatL[lev], StateList, 0, PhysTime, dT_lev);
    if (myrank == 0)
    {
      if (ErrorMonitor->outfile)
        ErrorMonitor->outfile << "find NaN in state variables at t = " << PhysTime
                              << ", lev = " << lev << endl;
      MPI_Abort(MPI_COMM_WORLD, 1);
    }
  }

#if (MAPBH == 0)
  if (BH_num > 0 && lev == GH->levels - 1)
  {
    compute_Porg_rhs(Porg0, Porg_rhs, Sfx0, Sfy0, Sfz0, lev);
    for (int ithBH = 0; ithBH < BH_num; ithBH++)
    {
      f_rungekutta4_scalar(dT_lev, Porg0[ithBH][0], Porg[ithBH][0], Porg_rhs[ithBH][0], iter_count);
      f_rungekutta4_scalar(dT_lev, Porg0[ithBH][1], Porg[ithBH][1], Porg_rhs[ithBH][1], iter_count);
      f_rungekutta4_scalar(dT_lev, Porg0[ithBH][2], Porg[ithBH][2], Porg_rhs[ithBH][2], iter_count);
      if (Symmetry > 0)
        Porg[ithBH][2] = fabs(Porg[ithBH][2]);
      if (Symmetry == 2)
      {
        Porg[ithBH][0] = fabs(Porg[ithBH][0]);
        Porg[ithBH][1] = fabs(Porg[ithBH][1]);
      }
    }
  }

  if (lev == a_lev)
    AnalysisStuff(lev, dT_lev);
#endif

  for (iter_count = 1; iter_count < 4; iter_count++)
  {
    if (iter_count == 1 || iter_count == 3)
      TRK4 += dT_lev / 2;

    Pp = GH->PatL[lev];
    while (Pp)
    {
      MyList<Block> *BP = Pp->data->blb;
      while (BP)
      {
        Block *cg = BP->data;
        if (myrank == cg->rank)
        {
          ensure_stage_device_var_list(cg, SynchList_pre);
          if (gpu_rhs(CALLED_BY_STEP, myrank, RHS_PARA_CALLED_THEN))
            ERROR = 1;

          run_stage_on_block(cg, Pp->data, StateList, SynchList_pre, SynchList_cor, RHSList, iter_count);

          if (bssn_cuda_lowerbound(cg->shape, cg->fgfs[phi1->sgfn], chitiny, false))
          {
          cerr << "GPU lowerbound failure: lev=" << lev
               << " rk_stage=" << iter_count
               << " var=" << phi1->name
                 << " bbox=(" << cg->bbox[0] << ":" << cg->bbox[3] << ","
                 << cg->bbox[1] << ":" << cg->bbox[4] << ","
                 << cg->bbox[2] << ":" << cg->bbox[5] << ")" << endl;
            ERROR = 1;
          }
          if (!ERROR && (!sync_cache_cor[lev].valid || iter_count == 3))
            stage_download_var_list(cg, SynchList_cor);
        }

        if (BP == Pp->data->ble)
          break;
        BP = BP->next;
      }
      Pp = Pp->next;
    }

    if (!ERROR && sync_cache_cor[lev].valid && iter_count < 3 &&
        !can_pack_sync_from_device(SynchList_cor, sync_cache_cor[lev]))
      refresh_stage_host_before_sync(SynchList_cor, sync_cache_cor[lev]);

    MPI_Request err_req_cor;
    {
      int erh = ERROR;
      MPI_Iallreduce(&erh, &ERROR, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD, &err_req_cor);
    }

    Parallel::AsyncSyncState async_cor;
    Parallel::Sync_start(GH->PatL[lev], SynchList_cor, Symmetry, sync_cache_cor[lev], async_cor);
    const bool unpack_cor_to_host = (iter_count == 3) || need_host_stage_sync;
    Parallel::Sync_finish(sync_cache_cor[lev], async_cor, SynchList_cor, Symmetry, unpack_cor_to_host);
    if (!ERROR && iter_count < 3 && unpack_cor_to_host)
      refresh_stage_device_after_sync(SynchList_cor, sync_cache_cor[lev]);

    MPI_Wait(&err_req_cor, MPI_STATUS_IGNORE);
    if (ERROR)
    {
      Parallel::Dump_Data(GH->PatL[lev], SynchList_pre, 0, PhysTime, dT_lev);
      if (myrank == 0)
      {
        if (ErrorMonitor->outfile)
          ErrorMonitor->outfile << "find NaN in RK4 substep#" << iter_count
                                << " variables at t = " << PhysTime
                                << ", lev = " << lev << endl;
        MPI_Abort(MPI_COMM_WORLD, 1);
      }
    }

#if (MAPBH == 0)
    if (BH_num > 0 && lev == GH->levels - 1)
    {
      compute_Porg_rhs(Porg, Porg1, Sfx, Sfy, Sfz, lev);
      for (int ithBH = 0; ithBH < BH_num; ithBH++)
      {
        f_rungekutta4_scalar(dT_lev, Porg0[ithBH][0], Porg1[ithBH][0], Porg_rhs[ithBH][0], iter_count);
        f_rungekutta4_scalar(dT_lev, Porg0[ithBH][1], Porg1[ithBH][1], Porg_rhs[ithBH][1], iter_count);
        f_rungekutta4_scalar(dT_lev, Porg0[ithBH][2], Porg1[ithBH][2], Porg_rhs[ithBH][2], iter_count);
        if (Symmetry > 0)
          Porg1[ithBH][2] = fabs(Porg1[ithBH][2]);
        if (Symmetry == 2)
        {
          Porg1[ithBH][0] = fabs(Porg1[ithBH][0]);
          Porg1[ithBH][1] = fabs(Porg1[ithBH][1]);
        }
      }
    }
#endif

    if (iter_count < 3)
    {
      Pp = GH->PatL[lev];
      while (Pp)
      {
        MyList<Block> *BP = Pp->data->blb;
        while (BP)
        {
          BP->data->swapList(SynchList_pre, SynchList_cor, myrank);
          if (BP == Pp->data->ble)
            break;
          BP = BP->next;
        }
        Pp = Pp->next;
      }

#if (MAPBH == 0)
      if (BH_num > 0 && lev == GH->levels - 1)
      {
        for (int ithBH = 0; ithBH < BH_num; ithBH++)
        {
          Porg[ithBH][0] = Porg1[ithBH][0];
          Porg[ithBH][1] = Porg1[ithBH][1];
          Porg[ithBH][2] = Porg1[ithBH][2];
        }
      }
#endif
    }
  }

#if (RPS == 0)
  RestrictProlong(lev, YN, BB);
#endif

  bssn_gpu_clear_cached_device_buffers();

  Pp = GH->PatL[lev];
  while (Pp)
  {
    MyList<Block> *BP = Pp->data->blb;
    while (BP)
    {
      Block *cg = BP->data;
      cg->swapList(StateList, SynchList_cor, myrank);
      cg->swapList(OldStateList, SynchList_cor, myrank);
      if (BP == Pp->data->ble)
        break;
      BP = BP->next;
    }
    Pp = Pp->next;
  }

  if (BH_num > 0 && lev == GH->levels - 1)
  {
    for (int ithBH = 0; ithBH < BH_num; ithBH++)
    {
      Porg0[ithBH][0] = Porg1[ithBH][0];
      Porg0[ithBH][1] = Porg1[ithBH][1];
      Porg0[ithBH][2] = Porg1[ithBH][2];
    }
  }
}

#endif