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

  Author: Uwe Schulzweida

*/

#include <algorithm>
#include <span>

#include "cdo_options.h"
#include "cdo_omp.h"
#include "remap_vars.h"
#include "remap_store_link.h"

template <std::ranges::random_access_range Rng>
void
sort_indexWeights(Rng &&rng) noexcept
{
  using T = std::ranges::range_value_t<Rng>;
  std::ranges::sort(rng, {}, &T::index);
}

void
store_weightLinks(size_t numWeights, std::vector<size_t> const &indices, std::vector<double> const &weights,
                  WeightLinks &weightLinks)
{
  weightLinks.nlinks = numWeights;
  auto &indexWeights = weightLinks.indexWeights;
  indexWeights.resize(numWeights);
  for (size_t i = 0; i < numWeights; ++i) { indexWeights[i] = { indices[i], weights[i] }; }

  if (numWeights == 1) return;
  if (!std::ranges::is_sorted(std::span(indices).first(numWeights))) sort_indexWeights(std::span(indexWeights).first(numWeights));
}

void
store_weightLinks(size_t index, double weight, WeightLinks1 &weightLinks)
{
  weightLinks.nlinks = 1;
  weightLinks.indexWeights = { index, weight };
}

void
store_weightLinks(std::array<size_t, 4> const &indices, std::array<double, 4> const &weights, WeightLinksBilinear &weightLinks)
{
  weightLinks.nlinks = 4;
  auto &indexWeights = weightLinks.indexWeights;
  for (int i = 0; i < 4; ++i) { indexWeights[i] = { indices[i], weights[i] }; }
}

void
store_weightLinks(std::array<size_t, 4> const &indices, BicubicWeights const &weights, WeightLinksBicubic &weightLinks)
{
  weightLinks.nlinks = 4;
  auto &indexWeights = weightLinks.indexWeights;
  for (int i = 0; i < 4; ++i) { indexWeights[i] = { indices[i], weights[i] }; }
}

void
weightLinks_to_remapLinks(size_t gridSize, std::vector<WeightLinks> const &weightLinksList, RemapVars &rv)
{
  std::vector<size_t> offsetList(gridSize, 0);

  size_t numLinks = 0;
  for (size_t i = 0; i < gridSize; ++i)
  {
    if (weightLinksList[i].nlinks)
    {
      offsetList[i] = numLinks;
      numLinks += weightLinksList[i].nlinks;
    }
  }

  rv.maxLinks = rv.numLinks = numLinks;
  if (!numLinks) return;

  rv.srcCellIndices.resize(numLinks);
  rv.tgtCellIndices.resize(numLinks);
  rv.weights.resize(numLinks * rv.numWeights);

  auto useLinksArrays = (Threading::ompNumMaxThreads > 1 && rv.numLinksPerValue == -1 && gridSize > cdoMinLoopSize);
  if (useLinksArrays) rv.linksOffset.resize(gridSize);
  if (useLinksArrays) rv.linksPerValue.resize(gridSize);

#ifdef _OPENMP
#pragma omp parallel for if (gridSize > cdoMinLoopSize) schedule(static) default(shared)
#endif
  for (size_t i = 0; i < gridSize; ++i)
  {
    auto &link = weightLinksList[i];
    if (link.nlinks)
    {
      auto offset = offsetList[i];
      if (useLinksArrays) { rv.linksOffset[i] = offset; }
      if (useLinksArrays) { rv.linksPerValue[i] = link.nlinks; }
      auto const &indexWeights = link.indexWeights;
      for (size_t ilink = 0; ilink < link.nlinks; ++ilink)
      {
        auto offset2 = offset + ilink;
        rv.srcCellIndices[offset2] = indexWeights[ilink].index;
        rv.tgtCellIndices[offset2] = i;
        rv.weights[offset2 * rv.numWeights] = indexWeights[ilink].weight;
      }
    }
    else
    {
      if (useLinksArrays) { rv.linksOffset[i] = 0; }
      if (useLinksArrays) { rv.linksPerValue[i] = 0; }
    }
  }
}

void
weightLinks_to_remapLinks(size_t gridSize, std::vector<WeightLinks1> const &weightLinksList, RemapVars &rv)
{
  std::vector<size_t> offsetList(gridSize, 0);

  size_t numLinks = 0;
#ifdef _OPENMP
// #pragma omp parallel for if (gridSize > cdoMinLoopSize) default(shared) schedule(static) reduction(+ : numLinks)
#endif
  for (size_t i = 0; i < gridSize; ++i)
  {
    if (weightLinksList[i].nlinks)
    {
      offsetList[i] = numLinks;
      numLinks += weightLinksList[i].nlinks;
    }
  }

  rv.maxLinks = numLinks;
  rv.numLinks = numLinks;

  if (!numLinks) return;

  rv.srcCellIndices.resize(numLinks);
  rv.tgtCellIndices.resize(numLinks);
  rv.weights.resize(numLinks);

#ifdef _OPENMP
#pragma omp parallel for if (gridSize > cdoMinLoopSize) default(shared) schedule(static)
#endif
  for (size_t i = 0; i < gridSize; ++i)
  {
    auto &link = weightLinksList[i];
    if (link.nlinks)
    {
      auto offset = offsetList[i];
      auto const &indexWeights = link.indexWeights;
      rv.srcCellIndices[offset] = indexWeights.index;
      rv.tgtCellIndices[offset] = i;
      rv.weights[offset] = indexWeights.weight;
    }
  }
}

void
weightLinks_to_remapLinks(size_t gridSize, std::vector<WeightLinksBilinear> const &weightLinksList, RemapVars &rv)
{
  std::vector<size_t> offsetList(gridSize, 0);

  size_t numLinks = 0;
#ifdef _OPENMP
// #pragma omp parallel for if (gridSize > cdoMinLoopSize) default(shared) schedule(static) reduction(+ : numLinks)
#endif
  for (size_t i = 0; i < gridSize; ++i)
  {
    if (weightLinksList[i].nlinks)
    {
      offsetList[i] = numLinks;
      numLinks += weightLinksList[i].nlinks;
    }
  }

  rv.maxLinks = numLinks;
  rv.numLinks = numLinks;

  if (!numLinks) return;

  rv.srcCellIndices.resize(numLinks);
  rv.tgtCellIndices.resize(numLinks);
  rv.weights.resize(numLinks);

#ifdef _OPENMP
#pragma omp parallel for if (gridSize > cdoMinLoopSize) default(shared) schedule(static)
#endif
  for (size_t i = 0; i < gridSize; ++i)
  {
    auto &link = weightLinksList[i];
    if (link.nlinks)
    {
      auto offset = offsetList[i];
      auto const &indexWeights = link.indexWeights;
      for (size_t ilink = 0; ilink < link.nlinks; ++ilink)
      {
        auto offset2 = offset + ilink;
        rv.srcCellIndices[offset2] = indexWeights[ilink].index;
        rv.tgtCellIndices[offset2] = i;
        rv.weights[offset2] = indexWeights[ilink].weight;
      }
    }
  }
}

void
weightLinks_to_remapLinks(size_t gridSize, std::vector<WeightLinksBicubic> const &weightLinksList, RemapVars &rv)
{
  std::vector<size_t> offsetList(gridSize, 0);

  size_t numLinks = 0;
#ifdef _OPENMP
// #pragma omp parallel for if (gridSize > cdoMinLoopSize) default(shared) schedule(static) reduction(+ : numLinks)
#endif
  for (size_t i = 0; i < gridSize; ++i)
  {
    if (weightLinksList[i].nlinks)
    {
      offsetList[i] = numLinks;
      numLinks += weightLinksList[i].nlinks;
    }
  }

  rv.maxLinks = numLinks;
  rv.numLinks = numLinks;

  if (!numLinks) return;

  rv.srcCellIndices.resize(numLinks);
  rv.tgtCellIndices.resize(numLinks);
  rv.weights.resize(4 * numLinks);

#ifdef _OPENMP
#pragma omp parallel for if (gridSize > cdoMinLoopSize) default(shared) schedule(static)
#endif
  for (size_t i = 0; i < gridSize; ++i)
  {
    auto &link = weightLinksList[i];
    if (link.nlinks)
    {
      auto offset = offsetList[i];
      auto const &indexWeights = link.indexWeights;
      for (size_t ilink = 0; ilink < link.nlinks; ++ilink)
      {
        auto offset2 = offset + ilink;
        rv.srcCellIndices[offset2] = indexWeights[ilink].index;
        rv.tgtCellIndices[offset2] = i;
        for (size_t k = 0; k < 4; ++k) { rv.weights[offset2 * 4 + k] = indexWeights[ilink].weight[k]; }
      }
    }
  }
}

void
set_numLinksPerValue(long &numLinksPerValue, long numWeights)
{
  if (numLinksPerValue == -1) { numLinksPerValue = numWeights; }
  else if (numLinksPerValue > 0 && numLinksPerValue != numWeights) { numLinksPerValue = 0; }
}

void
check_numLinksPerValue(std::vector<long> &v)
{
  if (v[0] > 0)
    for (size_t i = 1, n = v.size(); i < n; ++i)
      if (v[0] != v[i])
      {
        v[0] = 0;
        break;
      }
}
