/*
 * NCWriteMPAS.cpp
 *
 *  Created on: April 9, 2014
 */

#include "NCWriteMPAS.hpp"
#include "MBTagConventions.hpp"

namespace moab {

NCWriteMPAS::~NCWriteMPAS()
{
  // TODO Auto-generated destructor stub
}

ErrorCode NCWriteMPAS::collect_mesh_info()
{
  Interface*& mbImpl = _writeNC->mbImpl;
  std::vector<std::string>& dimNames = _writeNC->dimNames;
  std::vector<int>& dimLens = _writeNC->dimLens;
  Tag& mGlobalIdTag = _writeNC->mGlobalIdTag;

  ErrorCode rval;

  // Look for time dimension
  std::vector<std::string>::iterator vecIt;
  if ((vecIt = std::find(dimNames.begin(), dimNames.end(), "Time")) != dimNames.end())
    tDim = vecIt - dimNames.begin();
  else if ((vecIt = std::find(dimNames.begin(), dimNames.end(), "time")) != dimNames.end())
    tDim = vecIt - dimNames.begin();
  else {
    MB_SET_ERR(MB_FAILURE, "Couldn't find 'Time' or 'time' dimension");
  }
  nTimeSteps = dimLens[tDim];

  // Get number of levels
  if ((vecIt = std::find(dimNames.begin(), dimNames.end(), "nVertLevels")) != dimNames.end())
    levDim = vecIt - dimNames.begin();
  else {
    MB_SET_ERR(MB_FAILURE, "Couldn't find 'nVertLevels' dimension");
  }
  nLevels = dimLens[levDim];

  // Get local vertices
  rval = mbImpl->get_entities_by_dimension(_fileSet, 0, localVertsOwned);MB_CHK_SET_ERR(rval, "Trouble getting local vertices in current file set");
  assert(!localVertsOwned.empty());

  // Get local edges
  rval = mbImpl->get_entities_by_dimension(_fileSet, 1, localEdgesOwned);MB_CHK_SET_ERR(rval, "Trouble getting local edges in current file set");
  // There are no edges if NO_EDGES read option is set

  // Get local cells
  rval = mbImpl->get_entities_by_dimension(_fileSet, 2, localCellsOwned);MB_CHK_SET_ERR(rval, "Trouble getting local cells in current file set");
  assert(!localCellsOwned.empty());

#ifdef MOAB_HAVE_MPI
  bool& isParallel = _writeNC->isParallel;
  if (isParallel) {
    ParallelComm*& myPcomm = _writeNC->myPcomm;
    int rank = myPcomm->proc_config().proc_rank();
    int procs = myPcomm->proc_config().proc_size();
    if (procs > 1) {
#ifndef NDEBUG
      unsigned int num_local_verts = localVertsOwned.size();
#endif
      rval = myPcomm->filter_pstatus(localVertsOwned, PSTATUS_NOT_OWNED, PSTATUS_NOT);MB_CHK_SET_ERR(rval, "Trouble getting owned vertices in current file set");

      // Assume that PARALLEL_RESOLVE_SHARED_ENTS option is set
      // Verify that not all local vertices are owned by the last processor
      if (procs - 1 == rank)
        assert("PARALLEL_RESOLVE_SHARED_ENTS option is set" && localVertsOwned.size() < num_local_verts);

      if (!localEdgesOwned.empty()) {
        rval = myPcomm->filter_pstatus(localEdgesOwned, PSTATUS_NOT_OWNED, PSTATUS_NOT);MB_CHK_SET_ERR(rval, "Trouble getting owned edges in current file set");
      }

      rval = myPcomm->filter_pstatus(localCellsOwned, PSTATUS_NOT_OWNED, PSTATUS_NOT);MB_CHK_SET_ERR(rval, "Trouble getting owned cells in current file set");
    }
  }
#endif

  std::vector<int> gids(localVertsOwned.size());
  rval = mbImpl->tag_get_data(mGlobalIdTag, localVertsOwned, &gids[0]);MB_CHK_SET_ERR(rval, "Trouble getting global IDs on local vertices");

  // Get localGidVertsOwned
  std::copy(gids.rbegin(), gids.rend(), range_inserter(localGidVertsOwned));

  if (!localEdgesOwned.empty()) {
    gids.resize(localEdgesOwned.size());
    rval = mbImpl->tag_get_data(mGlobalIdTag, localEdgesOwned, &gids[0]);MB_CHK_SET_ERR(rval, "Trouble getting global IDs on local edges");

    // Get localGidEdgesOwned
    std::copy(gids.rbegin(), gids.rend(), range_inserter(localGidEdgesOwned));
  }

  gids.resize(localCellsOwned.size());
  rval = mbImpl->tag_get_data(mGlobalIdTag, localCellsOwned, &gids[0]);MB_CHK_SET_ERR(rval, "Trouble getting global IDs on local cells");

  // Get localGidCellsOwned
  std::copy(gids.rbegin(), gids.rend(), range_inserter(localGidCellsOwned));

  return MB_SUCCESS;
}

ErrorCode NCWriteMPAS::collect_variable_data(std::vector<std::string>& var_names, std::vector<int>& tstep_nums)
{
  NCWriteHelper::collect_variable_data(var_names, tstep_nums);

  std::vector<std::string>& dimNames = _writeNC->dimNames;
  std::vector<int>& dimLens = _writeNC->dimLens;

  // Dimension indices for other optional levels
  std::vector<unsigned int> opt_lev_dims;

  unsigned int lev_idx;
  std::vector<std::string>::iterator vecIt;

  // Get index of vertex levels P1
  if ((vecIt = std::find(dimNames.begin(), dimNames.end(), "nVertLevelsP1")) != dimNames.end()) {
    lev_idx = vecIt - dimNames.begin();
    opt_lev_dims.push_back(lev_idx);
  }

  // Get index of vertex levels P2
  if ((vecIt = std::find(dimNames.begin(), dimNames.end(), "nVertLevelsP2")) != dimNames.end()) {
    lev_idx = vecIt - dimNames.begin();
    opt_lev_dims.push_back(lev_idx);
  }

  // Get index of soil levels
  if ((vecIt = std::find(dimNames.begin(), dimNames.end(), "nSoilLevels")) != dimNames.end()) {
    lev_idx = vecIt - dimNames.begin();
    opt_lev_dims.push_back(lev_idx);
  }

  std::map<std::string, WriteNC::VarData>& varInfo = _writeNC->varInfo;

  for (size_t i = 0; i < var_names.size(); i++) {
    std::string varname = var_names[i];
    std::map<std::string, WriteNC::VarData>::iterator vit = varInfo.find(varname);
    if (vit == varInfo.end())
      MB_SET_ERR(MB_FAILURE, "Can't find variable " << varname);

    WriteNC::VarData& currentVarData = vit->second;
    std::vector<int>& varDims = currentVarData.varDims;

    // Skip edge variables, if there are no edges
    if (localEdgesOwned.empty() && currentVarData.entLoc == WriteNC::ENTLOCEDGE)
      continue;

    // If nVertLevels dimension is not found, try other optional levels such as nVertLevelsP1
    if (std::find(varDims.begin(), varDims.end(), levDim) == varDims.end()) {
      for (unsigned int j = 0; j < opt_lev_dims.size(); j++) {
        if (std::find(varDims.begin(), varDims.end(), opt_lev_dims[j]) != varDims.end()) {
          currentVarData.numLev = dimLens[opt_lev_dims[j]];
          break;
        }
      }
    }

    // Skip set variables, which were already processed in NCWriteHelper::collect_variable_data()
    if (WriteNC::ENTLOCSET == currentVarData.entLoc)
      continue;

    // Set up writeStarts and writeCounts (maximum number of dimensions is 3)
    currentVarData.writeStarts.resize(3);
    currentVarData.writeCounts.resize(3);
    unsigned int dim_idx = 0;

    // First: Time
    if (currentVarData.has_tsteps) {
      // Non-set variables with timesteps
      // 3 dimensions like (Time, nCells, nVertLevels)
      // 2 dimensions like (Time, nCells)
      assert(3 == varDims.size() || 2 == varDims.size());

      // Time should be the first dimension
      assert(tDim == varDims[0]);

      currentVarData.writeStarts[dim_idx] = 0; // This value is timestep dependent, will be set later
      currentVarData.writeCounts[dim_idx] = 1;
      dim_idx++;
    }
    else {
      // Non-set variables without timesteps
      // 2 dimensions like (nCells, nVertLevels)
      // 1 dimension like (nCells)
      assert(2 == varDims.size() || 1 == varDims.size());
    }

    // Next: nVertices / nCells / nEdges
    switch (currentVarData.entLoc) {
      case WriteNC::ENTLOCVERT:
        // Vertices
        // Start from the first localGidVerts
        // Actually, this will be reset later for writing
        currentVarData.writeStarts[dim_idx] = localGidVertsOwned[0] - 1;
        currentVarData.writeCounts[dim_idx] = localGidVertsOwned.size();
        break;
      case WriteNC::ENTLOCFACE:
        // Faces
        // Start from the first localGidCells
        // Actually, this will be reset later for writing
        currentVarData.writeStarts[dim_idx] = localGidCellsOwned[0] - 1;
        currentVarData.writeCounts[dim_idx] = localGidCellsOwned.size();
        break;
      case WriteNC::ENTLOCEDGE:
        // Edges
        // Start from the first localGidEdges
        // Actually, this will be reset later for writing
        currentVarData.writeStarts[dim_idx] = localGidEdgesOwned[0] - 1;
        currentVarData.writeCounts[dim_idx] = localGidEdgesOwned.size();
        break;
      default:
        MB_SET_ERR(MB_FAILURE, "Unexpected entity location type for variable " << varname);
    }
    dim_idx++;

    // Finally: nVertLevels or other optional levels, it is possible that there is no
    // level dimension (numLev is 0) for this non-set variable, e.g. (Time, nCells)
    if (currentVarData.numLev > 0) {
      // Non-set variables with levels
      // 3 dimensions like (Time, nCells, nVertLevels)
      // 2 dimensions like (nCells, nVertLevels)
      assert(3 == varDims.size() || 2 == varDims.size());

      currentVarData.writeStarts[dim_idx] = 0;
      currentVarData.writeCounts[dim_idx] = currentVarData.numLev;
      dim_idx++;
    }
    else {
      // Non-set variables without levels
      // 2 dimensions like (Time, nCells)
      // 1 dimension like (nCells)
      assert(2 == varDims.size() || 1 == varDims.size());
    }

    // Get variable size
    currentVarData.sz = 1;
    for (std::size_t idx = 0; idx < dim_idx; idx++)
      currentVarData.sz *= currentVarData.writeCounts[idx];
  }

  return MB_SUCCESS;
}

ErrorCode NCWriteMPAS::write_nonset_variables(std::vector<WriteNC::VarData>& vdatas, std::vector<int>& tstep_nums)
{
  Interface*& mbImpl = _writeNC->mbImpl;

  int success;

  // For each variable tag in the indexed lists, write a time step data
  for (unsigned int i = 0; i < vdatas.size(); i++) {
    WriteNC::VarData& variableData = vdatas[i];

    // Skip edge variables, if there are no edges
    if (localEdgesOwned.empty() && variableData.entLoc == WriteNC::ENTLOCEDGE)
      continue;

    // Get local owned entities of this variable
    Range* pLocalEntsOwned = NULL;
    Range* pLocalGidEntsOwned = NULL;
    switch (variableData.entLoc) {
      case WriteNC::ENTLOCVERT:
        // Vertices
        pLocalEntsOwned = &localVertsOwned;
        pLocalGidEntsOwned = &localGidVertsOwned;
        break;
      case WriteNC::ENTLOCEDGE:
        // Edges
        pLocalEntsOwned = &localEdgesOwned;
        pLocalGidEntsOwned = &localGidEdgesOwned;
        break;
      case WriteNC::ENTLOCFACE:
        // Cells
        pLocalEntsOwned = &localCellsOwned;
        pLocalGidEntsOwned = &localGidCellsOwned;
        break;
      default:
        MB_SET_ERR(MB_FAILURE, "Unexpected entity location type for variable " << variableData.varName);
    }

    unsigned int num_timesteps;
    unsigned int ents_idx = 0;
    if (variableData.has_tsteps) {
      // Non-set variables with timesteps
      // 3 dimensions like (Time, nCells, nVertLevels)
      // 2 dimensions like (Time, nCells)
      num_timesteps = tstep_nums.size();
      ents_idx++;
    }
    else {
      // Non-set variables without timesteps
      // 2 dimensions like (nCells, nVertLevels)
      // 1 dimension like (nCells)
      num_timesteps = 1;
    }

    unsigned int num_lev;
    if (variableData.numLev > 0) {
      // Non-set variables with levels
      // 3 dimensions like (Time, nCells, nVertLevels)
      // 2 dimensions like (nCells, nVertLevels)
      num_lev = variableData.numLev;
    }
    else {
      // Non-set variables without levels
      // 2 dimensions like (Time, nCells)
      // 1 dimension like (nCells)
      num_lev = 1;
    }

    for (unsigned int t = 0; t < num_timesteps; t++) {
      // We will write one time step, and count will be one; start will be different
      // Use tag_get_data instead of tag_iterate to copy tag data, as localEntsOwned
      // might not be contiguous.
      if (tDim == variableData.varDims[0])
        variableData.writeStarts[0] = t; // This is start for time
      std::vector<double> tag_data(pLocalEntsOwned->size() * num_lev);
      ErrorCode rval = mbImpl->tag_get_data(variableData.varTags[t], *pLocalEntsOwned,
                                            &tag_data[0]);MB_CHK_SET_ERR(rval, "Trouble getting tag data on owned entities");

#ifdef MOAB_HAVE_PNETCDF
      size_t nb_writes = pLocalGidEntsOwned->psize();
      std::vector<int> requests(nb_writes), statuss(nb_writes);
      size_t idxReq = 0;
#endif

      // Now write copied tag data
      // Use nonblocking put (request aggregation)
      switch (variableData.varDataType) {
        case NC_DOUBLE: {
          size_t indexInDoubleArray = 0;
          size_t ic = 0;
          for (Range::pair_iterator pair_iter = pLocalGidEntsOwned->pair_begin();
              pair_iter != pLocalGidEntsOwned->pair_end(); ++pair_iter, ic++) {
            EntityHandle starth = pair_iter->first;
            EntityHandle endh = pair_iter->second;
            variableData.writeStarts[ents_idx] = (NCDF_SIZE)(starth - 1);
            variableData.writeCounts[ents_idx] = (NCDF_SIZE)(endh - starth + 1);

            // Do a partial write, in each subrange
#ifdef MOAB_HAVE_PNETCDF
            // Wait outside this loop
            success = NCFUNCREQP(_vara_double)(_fileId, variableData.varId,
                &(variableData.writeStarts[0]), &(variableData.writeCounts[0]),
                           &(tag_data[indexInDoubleArray]), &requests[idxReq++]);
#else
            success = NCFUNCAP(_vara_double)(_fileId, variableData.varId,
                &(variableData.writeStarts[0]), &(variableData.writeCounts[0]),
                           &(tag_data[indexInDoubleArray]));
#endif
            if (success)
              MB_SET_ERR(MB_FAILURE, "Failed to write double data in a loop for variable " << variableData.varName);
            // We need to increment the index in double array for the
            // next subrange
            indexInDoubleArray += (endh - starth + 1) * num_lev;
          }
          assert(ic == pLocalGidEntsOwned->psize());
#ifdef MOAB_HAVE_PNETCDF
          success = ncmpi_wait_all(_fileId, requests.size(), &requests[0], &statuss[0]);
          if (success)
            MB_SET_ERR(MB_FAILURE, "Failed on wait_all");
#endif
          break;
        }
        default:
          MB_SET_ERR(MB_NOT_IMPLEMENTED, "Writing non-double data is not implemented yet");
      }
    }
  }

  return MB_SUCCESS;
}

} /* namespace moab */
