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

  Author: Uwe Schulzweida

*/

/*
   This module contains the following operators:

      Setbox     setclonlatbox   Set lon/lat box to constant
      Setbox     setcindexbox    Set index box to constant
*/

#include <cdi.h>

#include "process_int.h"
#include "param_conversion.h"
#include <mpim_grid.h>
#include "selboxinfo.h"

static void
setcbox(double constant, double *array, int gridID, const SelboxInfo &sbox)
{
  const auto &lat1 = sbox.lat1;
  const auto &lat2 = sbox.lat2;
  const auto &lon11 = sbox.lon11;
  const auto &lon12 = sbox.lon12;
  const auto &lon21 = sbox.lon21;
  const auto &lon22 = sbox.lon22;
  const long nlon = gridInqXsize(gridID);
  const long nlat = gridInqYsize(gridID);

  for (long ilat = 0; ilat < nlat; ilat++)
    for (long ilon = 0; ilon < nlon; ilon++)
      if ((lat1 <= ilat && ilat <= lat2 && ((lon11 <= ilon && ilon <= lon12) || (lon21 <= ilon && ilon <= lon22))))
        {
          array[nlon * ilat + ilon] = constant;
        }
}

void *
Setbox(void *process)
{
  int gridID = -1;
  int index, gridtype;
  SelboxInfo sbox;

  cdo_initialize(process);

  // clang-format off
  const auto SETCLONLATBOX = cdo_operator_add("setclonlatbox", 0, 0, "constant, western and eastern longitude and southern and northern latitude");
  const auto SETCINDEXBOX = cdo_operator_add("setcindexbox", 0, 0, "constant, index of first and last longitude and index of first and last latitude");
  // clang-format on

  const auto operatorID = cdo_operator_id();

  operator_input_arg(cdo_operator_enter(operatorID));

  const auto constant = parameter_to_double(cdo_operator_argv(0));

  const auto streamID1 = cdo_open_read(0);
  const auto vlistID1 = cdo_stream_inq_vlist(streamID1);

  const auto ngrids = vlistNgrids(vlistID1);
  int ndiffgrids = 0;
  for (index = 1; index < ngrids; index++)
    if (vlistGrid(vlistID1, 0) != vlistGrid(vlistID1, index)) ndiffgrids++;

  for (index = 0; index < ngrids; index++)
    {
      gridID = vlistGrid(vlistID1, index);
      gridtype = gridInqType(gridID);
      if (gridtype == GRID_LONLAT || gridtype == GRID_GAUSSIAN) break;
      if (gridtype == GRID_CURVILINEAR) break;
      if (operatorID == SETCINDEXBOX && gridtype == GRID_GENERIC && gridInqXsize(gridID) > 0 && gridInqYsize(gridID) > 0) break;
    }

  if (gridInqType(gridID) == GRID_GAUSSIAN_REDUCED)
    cdo_abort("Gaussian reduced grid found. Use option -R to convert it to a regular grid!");

  if (index == ngrids) cdo_abort("No regular grid found!");
  if (ndiffgrids > 0) cdo_abort("Too many different grids!");

  operator_input_arg(cdo_operator_enter(operatorID));

  if (operatorID == SETCLONLATBOX)
    genlonlatbox(1, gridID, sbox);
  else
    genindexbox(1, gridID, sbox);

  const auto vlistID2 = vlistDuplicate(vlistID1);

  const auto taxisID1 = vlistInqTaxis(vlistID1);
  const auto taxisID2 = taxisDuplicate(taxisID1);
  vlistDefTaxis(vlistID2, taxisID2);

  const auto nvars = vlistNvars(vlistID1);
  std::vector<bool> vars(nvars);
  for (int varID = 0; varID < nvars; varID++) vars[varID] = (gridID == vlistInqVarGrid(vlistID1, varID));

  const auto streamID2 = cdo_open_write(1);

  cdo_def_vlist(streamID2, vlistID2);

  const auto gridsize = gridInqSize(gridID);
  Varray<double> array(gridsize);

  int tsID = 0;
  while (true)
    {
      const auto nrecs = cdo_stream_inq_timestep(streamID1, tsID);
      if (nrecs == 0) break;

      cdo_taxis_copy_timestep(taxisID2, taxisID1);
      cdo_def_timestep(streamID2, tsID);

      for (int recID = 0; recID < nrecs; recID++)
        {
          int varID, levelID;
          cdo_inq_record(streamID1, &varID, &levelID);

          if (vars[varID])
            {
              size_t nmiss;
              cdo_read_record(streamID1, array.data(), &nmiss);

              setcbox(constant, array.data(), gridID, sbox);

              const auto missval = vlistInqVarMissval(vlistID1, varID);
              nmiss = varray_num_mv(gridsize, array, missval);
              cdo_def_record(streamID2, varID, levelID);
              cdo_write_record(streamID2, array.data(), nmiss);
            }
        }

      tsID++;
    }

  cdo_stream_close(streamID2);
  cdo_stream_close(streamID1);

  cdo_finish();

  return nullptr;
}
