/*
  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:

      Getgridcell     gridcellindex    Get grid cell index
*/

#include <algorithm>
#include <cdi.h>

#include "cdo_options.h"
#include "process_int.h"
#include "param_conversion.h"
#include "pmlist.h"
#include <mpim_grid.h>
#include "grid_healpix.h"
#include "grid_point_search.h"

struct GridPoint
{
  double lon = 0.0;
  double lat = 0.0;
  double arc_radius = 0.0;
  double radius = 1.0;
};

static size_t
lonlat_to_index(int gridID, const GridPoint &gridPoint)
{
  auto gridID0 = gridID;
  auto gridsize = gridInqSize(gridID);

  gridID = generate_full_point_grid(gridID);
  if (!gridHasCoordinates(gridID)) cdo_abort("Cell center coordinates missing!");

  Varray<double> xvals(gridsize), yvals(gridsize);
  gridInqXvals(gridID, xvals.data());
  gridInqYvals(gridID, yvals.data());

  // Convert lat/lon units if required
  cdo_grid_to_radian(gridID, CDI_XAXIS, xvals, "grid center lon");
  cdo_grid_to_radian(gridID, CDI_YAXIS, yvals, "grid center lat");

  bool xIsCyclic = false;
  size_t dims[2] = { gridsize, 0 };
  GridPointSearch gps;
  grid_point_search_create(gps, xIsCyclic, dims, gridsize, xvals, yvals);

  if (gridPoint.arc_radius > 0.0)
    grid_point_search_set_arc_radius(gps, gridPoint.arc_radius);
  else
    grid_point_search_set_chord_radius(gps, gridPoint.radius);

  constexpr size_t numNeighbors = 1;
  knnWeightsType knnWeights(numNeighbors);
  grid_search_point(gps, DEG2RAD * gridPoint.lon, DEG2RAD * gridPoint.lat, knnWeights);
  auto cellIdx = knnWeights.m_addr[0];

  grid_point_search_delete(gps);

  if (gridID0 != gridID) gridDestroy(gridID);

  return cellIdx;
}

static GridPoint
get_parameter()
{
  GridPoint gridPoint;

  auto pargc = cdo_operator_argc();
  if (pargc)
    {
      auto &pargv = cdo_get_oper_argv();

      KVList kvlist;
      kvlist.name = cdo_module_name();
      if (kvlist.parse_arguments(pargc, pargv) != 0) cdo_abort("Parse error!");
      if (Options::cdoVerbose) kvlist.print();

      for (const auto &kv : kvlist)
        {
          const auto &key = kv.key;
          if (kv.nvalues > 1) cdo_abort("Too many values for parameter key >%s<!", key);
          if (kv.nvalues < 1) cdo_abort("Missing value for parameter key >%s<!", key);
          const auto &value = kv.values[0];

          // clang-format off
          if      (key == "lon")        gridPoint.lon = parameter_to_double(value);
          else if (key == "lat")        gridPoint.lat = parameter_to_double(value);
          else if (key == "radius")     gridPoint.radius = radius_str_to_deg(value);
          else if (key == "arc_radius") gridPoint.arc_radius = radius_str_to_deg(value);
          else cdo_abort("Invalid parameter key >%s<!", key);
          // clang-format on
        }
    }

  return gridPoint;
}

static void
print_parameter(const GridPoint &gridPoint)
{
  std::stringstream outbuffer;
  outbuffer << "lon=" << gridPoint.lon << ", lat=" << gridPoint.lat;
  cdo_print("%s", outbuffer.str());
}

static void
check_radius_range(double radius, const char *name)
{
  if (radius < 0.0 || radius > 180.0) cdo_abort("%s=%g out of bounds (0-180 deg)!", name, radius);
}

class ModuleGetgridcell
{
  CdoStreamID streamID1;
  GridPoint gridPoint;
  int gridID1;

public:
  void
  init(void *process)
  {
    cdo_initialize(process);

    cdo_operator_add("gridcellindex", 0, 0, "lon/lat coordinate of a single cell");

    gridPoint = get_parameter();

    check_radius_range(gridPoint.radius, "radius");
    check_radius_range(gridPoint.arc_radius, "arc_radius");

    if (Options::cdoVerbose) print_parameter(gridPoint);

    streamID1 = cdo_open_read(0);

    auto vlistID1 = cdo_stream_inq_vlist(streamID1);

    auto ngrids = vlistNgrids(vlistID1);
    if (ngrids != 1) cdo_abort("Too many different grids!");

    gridID1 = vlistGrid(vlistID1, 0);
  }

  void
  run()
  {
    int64_t cellIdx = -1;
    if (is_healpix_grid(gridID1))
      {
        auto [nside, order] = cdo::get_healpix_params(gridID1);
        cellIdx = hp_lonlat_to_index(order, nside, DEG2RAD * gridPoint.lon, DEG2RAD * gridPoint.lat);
      }
    else { cellIdx = lonlat_to_index(gridID1, gridPoint); }

    printf("%ld\n", (long) cellIdx + 1);
  }

  void
  close()
  {
    cdo_stream_close(streamID1);

    cdo_finish();
  }
};

void *
Getgridcell(void *process)
{
  ModuleGetgridcell getgridcell;
  getgridcell.init(process);
  getgridcell.run();
  getgridcell.close();

  return nullptr;
}
