/*
  This file is part of CDO. CDO is a collection of Operators to manipulate and analyse Climate model Data.

  Author: Uwe Schulzweida

*/

#include <cdi.h>

#include "cdo_options.h"
#include "cdo_output.h"
#include "cdo_varlist.h"
#include "compare.h"
#include "cdi_lockedIO.h"
#include "varray.h"

int
cdo_read_timestepmask(const std::string &maskfile, std::vector<bool> &imask)
{
  auto streamID = stream_open_read_locked(maskfile.c_str());
  auto vlistID = streamInqVlist(streamID);
  VarList varList(vlistID);
  const auto &var0 = varList.vars[0];

  if (varList.numVars() > 1) cdo_abort("timestepmask %s contains more than one variable!", maskfile);
  if (var0.nlevels > 1) cdo_abort("timestepmask %s has more than one level!", maskfile);
  if (var0.gridsize > 1) cdo_abort("timestepmask %s has more than one gridpoint!", maskfile);

  auto nts = varList.numSteps();
  if (nts == -1)
    {
      nts = 0;
      while (streamInqTimestep(streamID, nts)) nts++;

      if (Options::cdoVerbose) cdo_print("%s: counted %i timeSteps in %s", __func__, nts, maskfile);

      streamClose(streamID);
      streamID = stream_open_read_locked(maskfile.c_str());
    }
  else if (Options::cdoVerbose) { cdo_print("%s: found %i timeSteps in %s", __func__, nts, maskfile); }

  int n = nts;
  imask.resize(nts);

  int tsID = 0;
  while (true)
    {
      auto numFields = streamInqTimestep(streamID, tsID);
      if (numFields == 0) break;

      if (numFields != 1) cdo_abort("Internal error; unexprected number of fields!");

      int varID, levelID;
      size_t numMissVals;
      double value;
      streamInqField(streamID, &varID, &levelID);
      streamReadField(streamID, &value, &numMissVals);

      imask[tsID] = !(numMissVals || is_equal(value, 0));

      tsID++;
    }

  streamClose(streamID);

  return n;
}

static void
read_one_field(const char *text, const char *filename, Varray<double> &array)
{
  auto streamID = stream_open_read_locked(filename);
  auto vlistID = streamInqVlist(streamID);
  VarList varList(vlistID);
  const auto &var0 = varList.vars[0];

  if (varList.numVars() > 1) cdo_abort("%s file %s contains more than one variable!", text, filename);
  if (var0.nlevels > 1) cdo_abort("%s file %s has more than one level!", text, filename);
  if (varList.maxFields() != 1) cdo_abort("%s file %s contains more than one field!", text, filename);

  array.resize(var0.gridsize);

  int varID, levelID;
  size_t numMissVals;
  streamInqField(streamID, &varID, &levelID);
  streamReadField(streamID, array.data(), &numMissVals);
  streamClose(streamID);
}

size_t
cdo_read_mask(const char *maskfile, std::vector<bool> &imask)
{
  Varray<double> array;
  read_one_field("Mask", maskfile, array);

  auto gridsize = array.size();
  imask.resize(gridsize);

  for (size_t i = 0; i < gridsize; ++i) imask[i] = is_not_equal(array[i], 0);

  return gridsize;
}

size_t
cdo_read_index(const char *indexfile, std::vector<int> &index)
{
  Varray<double> array;
  read_one_field("Index", indexfile, array);

  auto gridsize = array.size();
  index.resize(gridsize);

  for (size_t i = 0; i < gridsize; ++i) index[i] = (int) std::lround(array[i]) - 1;

  return gridsize;
}
