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

  Copyright (C) 2003-2021 Uwe Schulzweida, <uwe.schulzweida AT mpimet.mpg.de>
  See COPYING file for copying and redistribution conditions.

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; version 2 of the License.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
*/
#include <cdi.h>

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

int
cdo_read_timestepmask(const char *maskfile, std::vector<bool> &imask)
{
  int n = -1;

  auto streamID = stream_open_read_locked(maskfile);
  const auto vlistID = streamInqVlist(streamID);

  const auto nvars = vlistNvars(vlistID);
  if (nvars > 1) cdo_abort("timestepmask %s contains more than one variable!", maskfile);

  const auto gridsize = gridInqSize(vlistInqVarGrid(vlistID, 0));
  if (gridsize > 1) cdo_abort("timestepmask %s has more than one gridpoint!", maskfile);

  const auto nlev = zaxisInqSize(vlistInqVarZaxis(vlistID, 0));
  if (nlev > 1) cdo_abort("timestepmask %s has more than one level!", maskfile);

  auto nts = vlistNtsteps(vlistID);
  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);
    }
  else if (Options::cdoVerbose)
    cdo_print("%s: found %i timeSteps in %s", __func__, nts, maskfile);

  n = nts;
  imask.resize(nts);

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

      if (nrecs != 1) cdo_abort("Internal error; unexprected number of records!");

      int varID, levelID;
      size_t nmiss;
      double value;
      streamInqRecord(streamID, &varID, &levelID);
      streamReadRecord(streamID, &value, &nmiss);

      imask[tsID] = !(nmiss || IS_EQUAL(value, 0));

      tsID++;
    }

  streamClose(streamID);

  return n;
}

size_t
cdo_read_mask(const char *maskfile, std::vector<bool> &imask)
{
  const auto streamID = stream_open_read_locked(maskfile);
  const auto vlistID = streamInqVlist(streamID);

  const auto nvars = vlistNvars(vlistID);
  if (nvars > 1) cdo_abort("Mask %s contains more than one variable!", maskfile);

  const auto gridsize = gridInqSize(vlistInqVarGrid(vlistID, 0));

  const auto nlev = zaxisInqSize(vlistInqVarZaxis(vlistID, 0));
  if (nlev > 1) cdo_abort("Mask %s has more than one level!", maskfile);

  imask.resize(gridsize);
  Varray<double> array(gridsize);

  const auto nrecs = streamInqTimestep(streamID, 0);
  if (nrecs != 1) cdo_abort("Internal error; unexprected number of records!");

  int varID, levelID;
  size_t nmiss;
  streamInqRecord(streamID, &varID, &levelID);
  streamReadRecord(streamID, array.data(), &nmiss);
  streamClose(streamID);

  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)
{
  const auto streamID = stream_open_read_locked(indexfile);
  const auto vlistID = streamInqVlist(streamID);

  const auto nvars = vlistNvars(vlistID);
  if (nvars > 1) cdo_abort("Index file %s contains more than one variable!", indexfile);

  const auto gridsize = gridInqSize(vlistInqVarGrid(vlistID, 0));

  const auto nlev = zaxisInqSize(vlistInqVarZaxis(vlistID, 0));
  if (nlev > 1) cdo_abort("Index file %s has more than one level!", indexfile);

  index.resize(gridsize);
  Varray<double> array(gridsize);

  const auto nrecs = streamInqTimestep(streamID, 0);
  if (nrecs != 1) cdo_abort("Internal error; unexprected number of records!");

  int varID, levelID;
  size_t nmiss;
  streamInqRecord(streamID, &varID, &levelID);
  streamReadRecord(streamID, array.data(), &nmiss);
  streamClose(streamID);

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

  return gridsize;
}
