//=========================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt for details.
//
//  This software is distributed WITHOUT ANY WARRANTY; without even
//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
//  PURPOSE.  See the above copyright notice for more information.
//=========================================================================

#include "RadialAverage.h"

#include "smtk/mesh/Collection.h"
#include "smtk/mesh/Manager.h"
#include "smtk/mesh/PointLocator.h"
#include "smtk/mesh/PointSet.h"

#include "smtk/mesh/interpolation/PointCloud.h"
#include "smtk/mesh/interpolation/StructuredGrid.h"

#include <cmath>
#include <utility>

namespace
{
struct RadialAverageForPointCloud
{
  RadialAverageForPointCloud(
    smtk::mesh::CollectionPtr collection, const smtk::mesh::PointCloud& pointcloud, double radius)
    : m_pointcloud(pointcloud)
    , m_radius(radius)
    , m_locator(collection, pointcloud.size(), pointcloud.coordinates())
  {
  }

  double operator()(std::array<double, 3> x)
  {
    smtk::mesh::PointLocator::LocatorResults results;
    m_locator.find(x[0], x[1], 0.0, m_radius, results);

    if (results.pointIds.empty())
    {
      return std::numeric_limits<double>::quiet_NaN();
    }

    //otherwise we need to average the z values
    double sum = 0;
    std::size_t numPointsInRadius = 0;
    for (auto i : results.pointIds)
    {
      sum += m_pointcloud.data()(i);
      ++numPointsInRadius;
    }
    return sum / numPointsInRadius;
  }

  const smtk::mesh::PointCloud m_pointcloud;
  double m_radius;
  smtk::mesh::PointLocator m_locator;
};

struct RadialAverageForStructuredGrid
{
  typedef std::pair<int, int> Coord;

  RadialAverageForStructuredGrid(const smtk::mesh::StructuredGrid& structuredgrid, double radius)
    : m_structuredgrid(structuredgrid)
    , m_radius2(radius * radius)
  {
    m_discreteRadius[0] =
      std::abs(static_cast<int>(std::round(radius / m_structuredgrid.m_spacing[0])));
    m_discreteRadius[1] =
      std::abs(static_cast<int>(std::round(radius / m_structuredgrid.m_spacing[1])));

    m_limits[0] = m_structuredgrid.m_origin[0];
    m_limits[1] = m_structuredgrid.m_origin[0] +
      (m_structuredgrid.m_extent[1] - m_structuredgrid.m_extent[0]) * m_structuredgrid.m_spacing[0];
    if (m_limits[0] > m_limits[1])
    {
      std::swap(m_limits[0], m_limits[1]);
    }
    m_limits[2] = m_structuredgrid.m_origin[1];
    m_limits[3] = m_structuredgrid.m_origin[1] +
      (m_structuredgrid.m_extent[3] - m_structuredgrid.m_extent[2]) * m_structuredgrid.m_spacing[1];
    if (m_limits[2] > m_limits[3])
    {
      std::swap(m_limits[2], m_limits[3]);
    }
  }

  double operator()(std::array<double, 3> x)
  {
    if (x[0] < m_limits[0] || x[0] > m_limits[1] || x[1] < m_limits[2] || x[1] > m_limits[3])
    {
      // this point is outside of our data range
      return std::numeric_limits<double>::quiet_NaN();
    }

    // (ix,iy) represents the closest point in the grid to the query point
    int ix = static_cast<int>(std::round(m_structuredgrid.m_extent[0] +
      ((x[0] - m_structuredgrid.m_origin[0]) / m_structuredgrid.m_spacing[0])));
    int iy = static_cast<int>(std::round(m_structuredgrid.m_extent[2] +
      ((x[1] - m_structuredgrid.m_origin[1]) / m_structuredgrid.m_spacing[1])));

    // Since we allow for different spacing in x and y, our circle maps to an
    // ellipse with axes in the x and y dimensions of m_discreteRadius[0] and
    // m_discreteRadius[1], respectively, in discrete space. To integrate this
    // ellipse, we sweep across the y dimension; for each discrete step in y,
    // we compute the extrema in x and sum the results at each coordinate.
    double sum = 0.;
    std::size_t nCoords = 0;
    int i_extrema[2];
    for (int j = iy - m_discreteRadius[1]; j < iy + m_discreteRadius[1]; j++)
    {
      if (j < m_structuredgrid.m_extent[2] || j > m_structuredgrid.m_extent[3])
      {
        continue;
      }

      int halfChord = int(m_discreteRadius[0] * sin(acos(double(j - iy) / m_discreteRadius[1])));
      i_extrema[0] = ix - halfChord;
      i_extrema[1] = ix + halfChord;

      if (i_extrema[0] < m_structuredgrid.m_extent[0])
      {
        i_extrema[0] = m_structuredgrid.m_extent[0];
      }
      if (i_extrema[1] > m_structuredgrid.m_extent[1])
      {
        i_extrema[1] = m_structuredgrid.m_extent[1];
      }

      // We perform an unweighted average to maintain parity with the
      // unstructured grid version of this operator
      for (int i = i_extrema[0]; i < i_extrema[1]; i++)
      {
        if (m_structuredgrid.containsIndex(i, j))
        {
          sum += m_structuredgrid.data()(i, j);
          nCoords++;
        }
      }
    }
    if (nCoords == 0)
    {
      if (m_structuredgrid.containsIndex(ix, iy))
      {
        sum = m_structuredgrid.data()(ix, iy);
      }
      else
      {
        sum = std::numeric_limits<double>::quiet_NaN();
      }
    }
    else
    {
      sum /= nCoords;
    }

    return sum;
  }

  const smtk::mesh::StructuredGrid m_structuredgrid;
  double m_radius2;
  int m_discreteRadius[2];
  double m_limits[4];
};
}

namespace smtk
{
namespace mesh
{

RadialAverage::RadialAverage(
  smtk::mesh::CollectionPtr collection, const PointCloud& pointcloud, double radius)
  : m_function(RadialAverageForPointCloud(collection, pointcloud, radius))
{
}

RadialAverage::RadialAverage(const StructuredGrid& structuredgrid, double radius)
  : m_function(RadialAverageForStructuredGrid(structuredgrid, radius))
{
}
}
}
