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

  Author: Uwe Schulzweida

*/

#include "remap_grid_cell_search.h"
#include "cdo_output.h"
#include "compare.h"
#include "varray.h"
#include "cdo_options.h"
#include "cimdOmp.h"
#include <mpim_grid.h>

extern "C"
{
#include "lib/yac/sphere_part.h"
}

CellSearchMethod cellSearchMethod(CellSearchMethod::spherepart);

void
set_cell_search_method(const char *methodstr)
{
  // clang-format off
  if      (cdo_cmpstr(methodstr, "spherepart")) cellSearchMethod = CellSearchMethod::spherepart;
  else if (cdo_cmpstr(methodstr, "latbins"))    cellSearchMethod = CellSearchMethod::latbins;
  else cdo_abort("Grid cell search method %s not available!", methodstr);
  // clang-format on
}

static void
gridBoundboxReg2d(size_t nx, size_t ny, const Varray<double> &reg2d_corner_lon, const Varray<double> &reg2d_corner_lat,
                  double *grid_bound_box)
{
  grid_bound_box[0] = reg2d_corner_lat[0];
  grid_bound_box[1] = reg2d_corner_lat[ny];
  if (grid_bound_box[0] > grid_bound_box[1])
    {
      grid_bound_box[0] = reg2d_corner_lat[ny];
      grid_bound_box[1] = reg2d_corner_lat[0];
    }
  grid_bound_box[2] = reg2d_corner_lon[0];
  grid_bound_box[3] = reg2d_corner_lon[nx];
}

void
grid_cell_search_create_reg_2d(GridCellSearch &gcs, size_t dims[2], const Varray<double> &reg2d_corner_lon,
                               const Varray<double> &reg2d_corner_lat)
{
  gcs.is_reg2d = true;
  gcs.dims[0] = dims[0];
  gcs.dims[1] = dims[1];
  const auto nx = dims[0];
  const auto ny = dims[1];
  const auto nxp1 = nx + 1;
  const auto nyp1 = ny + 1;

  gcs.reg2d_corner_lon.resize(nxp1);
  gcs.reg2d_corner_lat.resize(nyp1);

  varray_copy(nxp1, reg2d_corner_lon, gcs.reg2d_corner_lon);
  varray_copy(nyp1, reg2d_corner_lat, gcs.reg2d_corner_lat);

  gridBoundboxReg2d(nx, ny, gcs.reg2d_corner_lon, gcs.reg2d_corner_lat, gcs.gridBoundboxReg2d);

  gcs.in_use = true;
}

void
grid_cell_search_create(GridCellSearch &gcs, size_t numCells, size_t numCellCorners, Varray<double> &cellCornerLon,
                        Varray<double> &cellCornerLat)
{
  gcs.method = cellSearchMethod;

  std::vector<struct grid_cell> cells(Threading::ompNumThreads);
  for (int i = 0; i < Threading::ompNumThreads; ++i)
    {
      cells[i].coordinates_xyz = (double(*)[3]) malloc(numCellCorners * sizeof(*(cells[i].coordinates_xyz)));
      cells[i].edge_type = (enum yac_edge_type *) malloc(numCellCorners * sizeof(*(cells[i].edge_type)));
      cells[i].num_corners = numCellCorners;
      cells[i].array_size = numCellCorners;
      for (size_t k = 0; k < numCellCorners; ++k) cells[i].edge_type[k] = GREAT_CIRCLE;
    }

  struct bounding_circle *bnd_circles = (struct bounding_circle *) malloc(numCells * sizeof(struct bounding_circle));

#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (size_t i = 0; i < numCells; ++i)
    {
      const auto ompthID = cdo_omp_get_thread_num();
      auto &cell = cells[ompthID];
      auto xyz = cell.coordinates_xyz;

      for (size_t k = 0; k < numCellCorners; ++k)
        gcLLtoXYZ(cellCornerLon[i * numCellCorners + k], cellCornerLat[i * numCellCorners + k], xyz[k]);

      if (numCellCorners == 3)
        yac_get_cell_bounding_circle_unstruct_triangle(xyz[0], xyz[1], xyz[2], &bnd_circles[i]);
      else
        yac_get_cell_bounding_circle(cell, &bnd_circles[i]);
    }

  gcs.yacSearch = yac_bnd_sphere_part_search_new(bnd_circles, numCells);
  gcs.yacBndCircles = bnd_circles;

  for (int i = 0; i < Threading::ompNumThreads; ++i)
    {
      free(cells[i].edge_type);
      free(cells[i].coordinates_xyz);
    }

  gcs.in_use = true;
}

void
grid_cell_search_delete(GridCellSearch &gcs)
{
  if (gcs.in_use)
    {
      varray_free(gcs.reg2d_corner_lon);
      varray_free(gcs.reg2d_corner_lat);

      if (gcs.yacSearch) yac_bnd_sphere_part_search_delete((bnd_sphere_part_search *) gcs.yacSearch);
      if (gcs.yacBndCircles) free(gcs.yacBndCircles);

      gcs.in_use = false;
    }
}

static size_t
doGridCellSearchYac(bnd_sphere_part_search *yacSearch, bounding_circle *yacBndCircles, bool isReg2dCell, const grid_cell &gridCell,
                    Varray<size_t> &srchAddr)
{
  size_t numCellCorners = gridCell.num_corners;
  bounding_circle bndCircle;
  double(*xyz)[3] = gridCell.coordinates_xyz;

  if (numCellCorners == 4 && isReg2dCell)
    yac_get_cell_bounding_circle_reg_quad(xyz[0], xyz[1], xyz[2], &bndCircle);
  else if (numCellCorners == 3)
    yac_get_cell_bounding_circle_unstruct_triangle(xyz[0], xyz[1], xyz[2], &bndCircle);
  else
    yac_get_cell_bounding_circle(gridCell, &bndCircle);

  size_t numSrchCells;
  size_t *currNeighs;
  yac_bnd_sphere_part_search_do_bnd_circle_search(yacSearch, &bndCircle, 1, &currNeighs, &numSrchCells);

  if (srchAddr.size() < numSrchCells) srchAddr.resize(numSrchCells);

  size_t k = 0;
  // for (size_t i = 0; i < numSrchCells; ++i) srchAddr[i] = currNeighs[i];
  for (size_t i = 0; i < numSrchCells; ++i)
    {
      if (yac_extents_overlap(&bndCircle, &yacBndCircles[currNeighs[i]])) srchAddr[k++] = currNeighs[i];
    }
  numSrchCells = k;
  free(currNeighs);

  return numSrchCells;
}

size_t
do_grid_cell_search(GridCellSearch &gcs, bool isReg2dCell, grid_cell &gridCell, Varray<size_t> &srchAddr)
{
  if (gcs.in_use)
    {
      // clang-format off
      if (gcs.method == CellSearchMethod::spherepart)
        return doGridCellSearchYac((bnd_sphere_part_search*)gcs.yacSearch, (bounding_circle *)gcs.yacBndCircles,
                                   isReg2dCell, gridCell, srchAddr);
      else
        cdo_abort("%s::method undefined!", __func__);
      // clang-format on
    }

  return 0;
}
