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

  Author: Uwe Schulzweida

*/

#include <cmath>

#include "cdo_math.h"
#include "knndata.h"
#include "pointsearch_reg2d.h"
#include "grid_convert.h"

static size_t
fill_src_indices(bool isCyclic, long nx, long ny, long ii, long jj, long k, std::span<size_t> psrcIndices)
{
  k /= 2;

  auto j0 = jj - k;
  auto jn = jj + k;
  auto i0 = ii - k;
  auto in = ii + k;
  if (j0 < 0) j0 = 0;
  if (jn >= ny) jn = ny - 1;
  if ((in - i0) > nx)
  {
    i0 = 0;
    in = nx - 1;
  }

  size_t numIndices = 0;

  for (long j = j0; j <= jn; ++j)
    for (long i = i0; i <= in; ++i)
    {
      auto ix = i;
      if (isCyclic && ix < 0) ix += nx;
      if (isCyclic && ix >= nx) ix -= nx;
      if (ix >= 0 && ix < nx && j < ny) psrcIndices[numIndices++] = j * nx + ix;
    }

  return numIndices;
}

void
PointsearchReg2d::compute_point(size_t index, std::span<double, 3> xyz)
{
  auto iy = index / m_nx;
  auto ix = index - iy * m_nx;
  xyz[0] = m_cosLats[iy] * m_cosLons[ix];
  xyz[1] = m_cosLats[iy] * m_sinLons[ix];
  xyz[2] = m_sinLats[iy];
}

void
PointsearchReg2d::store_distance_reg2d(double plon, double plat, KnnData &knnData, size_t numIndices, std::span<size_t> indices,
                                       std::span<double> distances, double searchRadius)
{
  std::array<double, 3> tgtCoord;
  lonlat_to_xyz(plon, plat, tgtCoord);
  auto sqrSearchRadius = cdo::sqr(searchRadius);

  std::array<double, 3> srcCoord;
  size_t numWeights = 0;
  for (size_t i = 0; i < numIndices; ++i)
  {
    compute_point(indices[i], srcCoord);
    // Find distance to this point
    double sqrDistance = static_cast<float>(cdo::sqr_distance(tgtCoord, srcCoord));
    // Store the index and distance if this is one of the smallest so far
    if (sqrDistance <= sqrSearchRadius)
    {
      indices[numWeights] = indices[i];
      distances[numWeights] = std::sqrt(sqrDistance);
      numWeights++;
    }
  }

  auto maxNeighbors = std::min(numWeights, knnData.maxNeighbors());
  for (size_t i = 0; i < numWeights; ++i) { knnData.store_distance(indices[i], distances[i], maxNeighbors); }

  knnData.check_distance();

  if (knnData.m_needCoords)
  {
    lonlat_to_xyz(plon, plat, knnData.m_tgtCoord);
    auto numNeighbors = knnData.m_numNeighbors;
    for (size_t i = 0; i < numNeighbors; ++i) { compute_point(knnData.m_indices[i], knnData.m_srcCoords[i]); }
  }
}

bool
PointsearchReg2d::search(double plon, double plat, KnnData &knnData, double searchRadius, bool isCyclic)
{
  auto nxm = isCyclic ? m_nx + 1 : m_nx;
  auto numNeighbors = knnData.maxNeighbors();

  size_t ii, jj;
  auto pointFound = rect_grid_search(ii, jj, plon, plat, nxm, m_ny, m_lonsReg2d, m_latsReg2d);
  if (not pointFound) return pointFound;

  if (isCyclic && ii == (nxm - 1)) { ii = 0; }

  constexpr size_t MAX_SEARCH_CELLS{ 25 };
  std::array<size_t, MAX_SEARCH_CELLS> srcIndices{};
  std::array<double, MAX_SEARCH_CELLS> srcDist{};

  size_t k{ 3 };
  for (; k < 10000; k += 2)
    if (numNeighbors <= (k - 2) * (k - 2)) break;

  auto maxCells = k * k;
  std::vector<size_t> tmpIndices;
  std::vector<double> tmpDist;
  std::span<size_t> psrcIndices;
  std::span<double> psrcDist;
  if (maxCells > MAX_SEARCH_CELLS)
  {
    tmpIndices.resize(maxCells);
    tmpDist.resize(maxCells);
    psrcIndices = tmpIndices;
    psrcDist = tmpDist;
  }
  else
  {
    psrcIndices = srcIndices;
    psrcDist = srcDist;
  }

  auto numIndices = fill_src_indices(isCyclic, m_nx, m_ny, ii, jj, k, psrcIndices);
  store_distance_reg2d(plon, plat, knnData, numIndices, psrcIndices, psrcDist, searchRadius);

  return pointFound;
}
