/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2024, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */
#pragma once

#include <mrpt/core/round.h>

#include <cmath>
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>

namespace mrpt
{
namespace containers
{
namespace internal
{
// Aux class.
struct dynamic_grid_txt_saver
{
  bool saveToTextFile(const std::string& fileName) const;
  virtual unsigned int getSizeX() const = 0;
  virtual unsigned int getSizeY() const = 0;
  virtual float getCellAsFloat(unsigned int cx, unsigned int cy) const = 0;
};
}  // namespace internal

/** A 2D grid of dynamic size which stores any kind of data at each cell.
 * \tparam T The type of each cell in the 2D grid.
 * \ingroup mrpt_containers_grp
 */
template <class T>
class CDynamicGrid
{
 public:
  using grid_data_t = std::vector<T>;
  using iterator = typename grid_data_t::iterator;
  using const_iterator = typename grid_data_t::const_iterator;

  /** Constructor */
  CDynamicGrid(
      double x_min = -10.,
      double x_max = 10.,
      double y_min = -10.,
      double y_max = 10.,
      double resolution = 0.1)
  {
    setSize(x_min, x_max, y_min, y_max, resolution);
  }

  /** Destructor */
  virtual ~CDynamicGrid() = default;
  /** Changes the size of the grid, ERASING all previous contents.
   * If \a fill_value is left as nullptr, the contents of cells may be
   * undefined (some will remain with
   *  their old values, the new ones will have the default cell value, but
   * the location of old values
   *  may change wrt their old places).
   * If \a fill_value is not nullptr, it is assured that all cells will have
   * a copy of that value after resizing.
   * \sa resize, fill
   */
  void setSize(
      const double x_min,
      const double x_max,
      const double y_min,
      const double y_max,
      const double resolution,
      const T* fill_value = nullptr)
  {
    // Adjust sizes to adapt them to full sized cells according to the
    // resolution:
    m_x_min = resolution * round(x_min / resolution);
    m_y_min = resolution * round(y_min / resolution);
    m_x_max = resolution * round(x_max / resolution);
    m_y_max = resolution * round(y_max / resolution);

    // Res:
    m_resolution = resolution;

    // Now the number of cells should be integers:
    m_size_x = round((m_x_max - m_x_min) / m_resolution);
    m_size_y = round((m_y_max - m_y_min) / m_resolution);

    // Cells memory:
    if (fill_value)
      m_map.assign(m_size_x * m_size_y, *fill_value);
    else
      m_map.resize(m_size_x * m_size_y);
  }

  /** Erase the contents of all the cells. */
  void clear()
  {
    m_map.clear();
    m_map.resize(m_size_x * m_size_y);
  }

  /** Fills all the cells with the same value
   */
  inline void fill(const T& value)
  {
    for (auto it = m_map.begin(); it != m_map.end(); ++it) *it = value;
  }

  /** Changes the size of the grid, maintaining previous contents.
   * \sa setSize
   */
  virtual void resize(
      double new_x_min,
      double new_x_max,
      double new_y_min,
      double new_y_max,
      const T& defaultValueNewCells,
      double additionalMarginMeters = 2.0)
  {
    // Is resize really necessary?
    if (new_x_min >= m_x_min && new_y_min >= m_y_min && new_x_max <= m_x_max &&
        new_y_max <= m_y_max)
      return;

    if (new_x_min > m_x_min) new_x_min = m_x_min;
    if (new_x_max < m_x_max) new_x_max = m_x_max;
    if (new_y_min > m_y_min) new_y_min = m_y_min;
    if (new_y_max < m_y_max) new_y_max = m_y_max;

    // Additional margin:
    if (additionalMarginMeters > 0)
    {
      if (new_x_min < m_x_min) new_x_min = floor(new_x_min - additionalMarginMeters);
      if (new_x_max > m_x_max) new_x_max = ceil(new_x_max + additionalMarginMeters);
      if (new_y_min < m_y_min) new_y_min = floor(new_y_min - additionalMarginMeters);
      if (new_y_max > m_y_max) new_y_max = ceil(new_y_max + additionalMarginMeters);
    }

    // Adjust sizes to adapt them to full sized cells according to the
    // resolution:
    if (fabs(new_x_min / m_resolution - round(new_x_min / m_resolution)) > 0.05f)
      new_x_min = m_resolution * round(new_x_min / m_resolution);
    if (fabs(new_y_min / m_resolution - round(new_y_min / m_resolution)) > 0.05f)
      new_y_min = m_resolution * round(new_y_min / m_resolution);
    if (fabs(new_x_max / m_resolution - round(new_x_max / m_resolution)) > 0.05f)
      new_x_max = m_resolution * round(new_x_max / m_resolution);
    if (fabs(new_y_max / m_resolution - round(new_y_max / m_resolution)) > 0.05f)
      new_y_max = m_resolution * round(new_y_max / m_resolution);

    // Change the map size: Extensions at each side:
    unsigned int extra_x_izq = round((m_x_min - new_x_min) / m_resolution);
    unsigned int extra_y_arr = round((m_y_min - new_y_min) / m_resolution);

    unsigned int new_size_x = round((new_x_max - new_x_min) / m_resolution);
    unsigned int new_size_y = round((new_y_max - new_y_min) / m_resolution);

    // Reserve new memory:
    grid_data_t new_map;
    new_map.resize(new_size_x * new_size_y, defaultValueNewCells);

    // Copy previous rows:
    unsigned int x, y;
    iterator itSrc, itDst;
    for (y = 0; y < m_size_y; y++)
    {
      for (x = 0, itSrc = (m_map.begin() + y * m_size_x),
          itDst = (new_map.begin() + extra_x_izq + (y + extra_y_arr) * new_size_x);
           x < m_size_x; ++x, ++itSrc, ++itDst)
      {
        *itDst = *itSrc;
      }
    }

    // Update the new map limits:
    m_x_min = new_x_min;
    m_x_max = new_x_max;
    m_y_min = new_y_min;
    m_y_max = new_y_max;

    m_size_x = new_size_x;
    m_size_y = new_size_y;

    // Keep the new map only:
    m_map.swap(new_map);
  }

  /** Returns a pointer to the contents of a cell given by its coordinates, or
   * nullptr if it is out of the map extensions.
   */
  inline T* cellByPos(double x, double y)
  {
    const int cx = x2idx(x);
    const int cy = y2idx(y);
    if (cx < 0 || cx >= static_cast<int>(m_size_x)) return nullptr;
    if (cy < 0 || cy >= static_cast<int>(m_size_y)) return nullptr;
    return &m_map[cx + cy * m_size_x];
  }
  /** \overload */
  inline const T* cellByPos(double x, double y) const
  {
    const int cx = x2idx(x);
    const int cy = y2idx(y);
    if (cx < 0 || cx >= static_cast<int>(m_size_x)) return nullptr;
    if (cy < 0 || cy >= static_cast<int>(m_size_y)) return nullptr;
    return &m_map[cx + cy * m_size_x];
  }

  /** Returns a pointer to the contents of a cell given by its cell indexes,
   * or nullptr if it is out of the map extensions.
   */
  inline T* cellByIndex(unsigned int cx, unsigned int cy)
  {
    if (cx >= m_size_x || cy >= m_size_y)
      return nullptr;
    else
      return &m_map[cx + cy * m_size_x];
  }

  /** Returns a pointer to the contents of a cell given by its cell indexes,
   * or nullptr if it is out of the map extensions.
   */
  inline const T* cellByIndex(unsigned int cx, unsigned int cy) const
  {
    if (cx >= m_size_x || cy >= m_size_y)
      return nullptr;
    else
      return &m_map[cx + cy * m_size_x];
  }

  /** Returns the horizontal size of grid map in cells count */
  inline size_t getSizeX() const { return m_size_x; }
  /** Returns the vertical size of grid map in cells count */
  inline size_t getSizeY() const { return m_size_y; }
  /** Returns the "x" coordinate of left side of grid map */
  inline double getXMin() const { return m_x_min; }
  /** Returns the "x" coordinate of right side of grid map */
  inline double getXMax() const { return m_x_max; }
  /** Returns the "y" coordinate of top side of grid map */
  inline double getYMin() const { return m_y_min; }
  /** Returns the "y" coordinate of bottom side of grid map */
  inline double getYMax() const { return m_y_max; }
  /** Returns the resolution of the grid map */
  inline double getResolution() const { return m_resolution; }
  /** Transform a coordinate values into cell indexes */
  inline int x2idx(double x) const { return static_cast<int>((x - m_x_min) / m_resolution); }
  inline int y2idx(double y) const { return static_cast<int>((y - m_y_min) / m_resolution); }
  inline int xy2idx(double x, double y) const { return x2idx(x) + y2idx(y) * m_size_x; }

  /** Transform a global (linear) cell index value into its corresponding
   * (x,y) cell indexes. */
  inline void idx2cxcy(int idx, int& cx, int& cy) const
  {
    cx = idx % m_size_x;
    cy = idx / m_size_x;
  }

  /** Transform a cell index into a coordinate value of the cell central point
   */
  inline double idx2x(int cx) const { return m_x_min + (cx + 0.5) * m_resolution; }
  inline double idx2y(int cy) const { return m_y_min + (cy + 0.5) * m_resolution; }

  /** Get the entire grid as a matrix.
   *  \tparam MAT The type of the matrix, typically a
   * mrpt::math::CMatrixDouble.
   *  \param[out] m The output matrix; will be set automatically to the
   * correct size.
   *  Entry (cy,cx) in the matrix contains the grid cell with indices
   * (cx,cy).
   * \note This method will compile only for cell types that can be
   * converted to the type of the matrix elements (e.g. double).
   */
  template <class MAT>
  void getAsMatrix(MAT& m) const
  {
    m.setSize(m_size_y, m_size_x);
    if (m_map.empty()) return;
    const T* c = &m_map[0];
    for (size_t cy = 0; cy < m_size_y; cy++)
      for (size_t cx = 0; cx < m_size_x; cx++) m(cy, cx) = *c++;
  }

  /** The user must implement this in order to provide "saveToTextFile" a way
   * to convert each cell into a numeric value */
  virtual float cell2float(const T&) const { return 0; }
  /** Saves a float representation of the grid (via "cell2float()") to a text
   * file. \return false on error */
  bool saveToTextFile(const std::string& fileName) const
  {
    struct aux_saver : public internal::dynamic_grid_txt_saver
    {
      aux_saver(const CDynamicGrid<T>& obj) : m_obj(obj) {}
      unsigned int getSizeX() const override { return m_obj.getSizeX(); }
      unsigned int getSizeY() const override { return m_obj.getSizeY(); }
      float getCellAsFloat(unsigned int cx, unsigned int cy) const override
      {
        return m_obj.cell2float(m_obj.m_map[cx + cy * m_obj.getSizeX()]);
      }
      const CDynamicGrid<T>& m_obj;
    };
    aux_saver aux(*this);
    return aux.saveToTextFile(fileName);
  }

  /** @name Direct and range-based access to data
   * @{ */
  inline const grid_data_t& data() const { return m_map; }
  inline iterator begin() { return m_map.begin(); }
  inline iterator end() { return m_map.end(); }
  inline const_iterator begin() const { return m_map.begin(); }
  inline const_iterator end() const { return m_map.end(); }
  /** @} */

 protected:
  template <class STREAM>
  void dyngridcommon_writeToStream(STREAM& out) const
  {
    out << m_x_min << m_x_max << m_y_min << m_y_max;
    out << m_resolution;
    out.template WriteAs<uint32_t>(m_size_x).template WriteAs<uint32_t>(m_size_y);
  }
  template <class STREAM>
  void dyngridcommon_readFromStream(STREAM& in, bool cast_from_float = false)
  {
    if (!cast_from_float)
    {
      in >> m_x_min >> m_x_max >> m_y_min >> m_y_max;
      in >> m_resolution;
    }
    else
    {
      float xmin, xmax, ymin, ymax, res;
      in >> xmin >> xmax >> ymin >> ymax >> res;
      m_x_min = xmin;
      m_x_max = xmax;
      m_y_min = ymin;
      m_y_max = ymax;
      m_resolution = res;
    }
    m_size_x = in.template ReadAs<uint32_t>();
    m_size_y = in.template ReadAs<uint32_t>();
    m_map.resize(m_size_x * m_size_y);
  }

 protected:
  /** The cells  */
  grid_data_t m_map;
  /** Used only from logically const method that really need to modify the
   * object */
  inline grid_data_t& m_map_castaway_const() const { return const_cast<grid_data_t&>(m_map); }

  double m_x_min{0}, m_x_max{0}, m_y_min{0}, m_y_max{0}, m_resolution{0};
  size_t m_size_x{0}, m_size_y{0};

};  // end of CDynamicGrid<>

}  // namespace containers
}  // namespace mrpt
