/*
 * (C) Copyright 2013 ECMWF.
 *
 * This software is licensed under the terms of the Apache Licence Version 2.0
 * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
 * In applying this licence, ECMWF does not waive the privileges and immunities
 * granted to it by virtue of its status as an intergovernmental organisation
 * nor does it submit to any jurisdiction.
 */

#pragma once

#include "atlas/util/Geometry.h"
#include "atlas/util/ObjectHandle.h"
#include "atlas/util/detail/KDTree.h"

namespace atlas {
namespace util {

//------------------------------------------------------------------------------------------------------

/// @brief k-dimensional tree constructable both with 2D (lon,lat) points as with 3D (x,y,z) points
///
/// The implementation is based on eckit::KDTreeMemory with 3D (x,y,z) points.
/// 2D points (lon,lat) are converted when needed to 3D during insertion, and during search, so that
/// a search always happens with 3D cartesian points.
///
/// ### Example:
///
/// Construct a KDTree, given a `list_of_lonlat_points` (e.g. `std::vector<PointLonLat>` or other container)
/// A payload given is the index in this list.
/// @code{.cpp}
///     KDTree<idx_t> search;
///     search.reserve( list_of_lonlat_points.size() );
///     idx_t n{0};
///     for ( auto& point : list_of_lonlat_points ) {
///         search.insert( point, n++ );
///     }
///     search.build();
/// @endcode
/// We can now do e.g. a search for the nearest 4 neighbours (k=4) sorted by shortest distance
/// @code{.cpp}
///     idx_t k = 4;
///     auto neighbours = search.closestPoints( PointLonLat{180., 45.}, k ).payloads();
/// @endcode
/// The variable `neighbours` is now a container of indices (the payloads) of the 4 nearest points

template <typename PayloadT, typename PointT = Point3>
class KDTree : public ObjectHandle<detail::KDTreeBase<PayloadT, PointT>> {
public:
    using Handle         = ObjectHandle<detail::KDTreeBase<PayloadT, PointT>>;
    using Implementation = typename Handle::Implementation;
    using Point          = typename Implementation::Point;
    using Payload        = typename Implementation::Payload;
    using PayloadList    = typename Implementation::PayloadList;
    using Value          = typename Implementation::Value;
    using ValueList      = typename Implementation::ValueList;

    using Handle::get;
    using Handle::Handle;

public:
    //--------------------------------------------------------------------------------------
    // Constructors

    KDTree( const KDTree& handle ) : Handle( handle ) {}

    /// @brief Construct an empty kd-tree with default geometry (Earth)
    KDTree() : Handle( new detail::KDTreeMemory<Payload, Point>() ) {}

    /// @brief Construct an empty kd-tree with custom geometry
    KDTree( const Geometry& geometry ) : Handle( new detail::KDTreeMemory<Payload, Point>( geometry ) ) {}

    /// @brief Construct a shared kd-tree with default geometry (Earth)
    template <typename Tree>
    KDTree( const std::shared_ptr<Tree>& kdtree ) :
        Handle( new detail::KDTree_eckit<Tree, Payload, Point>( kdtree ) ) {}

    /// @brief Construct a shared kd-tree with custom geometry
    template <typename Tree>
    KDTree( const std::shared_ptr<Tree>& kdtree, const Geometry& geometry ) :
        Handle( new detail::KDTree_eckit<Tree, Payload, Point>( kdtree, geometry ) ) {}

    //--------------------------------------------------------------------------------------
    // Methods to build the KDTree

    /// @brief Reserve memory for building the kd-tree in one shot (optional, at cost of extra memory)
    void reserve( idx_t size ) { get()->reserve( size ); }

    /// @brief Insert spherical point (lon,lat) or 3D cartesian point (x,y,z)
    /// @warning If memory has been reserved with reserve(), insertion will be delayed until build() is called.
    template <typename Point>
    void insert( const Point& p, const Payload& payload ) {
        get()->insert( p, payload );
    }

    /// @brief Insert kd-tree value in tree
    /// @warning If memory has been reserved with reserve(), insertion will be delayed until build() is called.
    void insert( const Value& value ) { get()->insert( value ); }

    /// @brief Build the kd-tree in one shot, if memory has been reserved.
    /// This will need to be called before all search functions like closestPoints().
    /// @post The KDTree is ready to be used
    void build() { get()->build(); }

    /// @brief Build with spherical points (lon,lat) where longitudes, latitudes, and payloads are separate containers.
    /// Memory will be reserved with reserve() to match the size
    /// @post The KDTree is ready to be used
    template <typename Longitudes, typename Latitudes, typename Payloads>
    void build( const Longitudes& longitudes, const Latitudes& latitudes, const Payloads& payloads ) {
        get()->build( longitudes, latitudes, payloads );
    }

    /// @brief Build with spherical points (lon,lat) given separate iterator ranges for longitudes, latitudes, and payloads.
    /// Memory will be reserved with reserve() to match the size
    /// @post The KDTree is ready to be used
    template <typename LongitudesIterator, typename LatitudesIterator, typename PayloadsIterator>
    void build( const LongitudesIterator& longitudes_begin, const LongitudesIterator& longitudes_end,
                const LatitudesIterator& latitudes_begin, const LatitudesIterator& latitudes_end,
                const PayloadsIterator& payloads_begin, const PayloadsIterator& payloads_end ) {
        get()->build( longitudes_begin, longitudes_end, latitudes_begin, latitudes_end, payloads_begin, payloads_end );
    }

    /// @brief Build with templated points. Points can be either 2D (lon,lat) or 3D (x,y,z)
    /// Memory will be reserved with reserve() to match the size
    /// @post The KDTree is ready to be used
    template <typename Points, typename Payloads>
    void build( const Points& points, const Payloads& payloads ) {
        get()->build( points, payloads );
    }

    /// @brief Build with spherical points (lon,lat) given separate iterator ranges for longitudes, latitudes, and payloads.
    /// Memory will be reserved with reserve() to match the size
    /// @post The KDTree is ready to be used
    template <typename PointIterator, typename PayloadsIterator>
    void build( const PointIterator& points_begin, const PointIterator& points_end,
                const PayloadsIterator& payloads_begin, const PayloadsIterator& payloads_end ) {
        get()->build( points_begin, points_end, payloads_begin, payloads_end );
    }

    /// @brief Build with vector of Value
    /// @post The KDTree is ready to be used
    void build( const std::vector<Value>& values ) { get()->build( values ); }

    //--------------------------------------------------------------------------------------
    // Methods to access the KDTree

    bool empty() const { return get()->empty(); }

    size_t size() const { return get()->size(); }

    size_t footprint() const { return get()->footprint(); }

    /// @brief Find k closest points given a 3D cartesian point (x,y,z) or 2D lonlat point(lon,lat)
    template <typename Point>
    ValueList closestPoints( const Point& p, size_t k ) const {
        return get()->closestPoints( p, k );
    }

    /// @brief Find closest point given a 3D cartesian point (x,y,z) or 2D lonlat point(lon,lat)
    template <typename Point>
    Value closestPoint( const Point& p ) const {
        return get()->closestPoint( p );
    }

    /// @brief Find all points within a distance of given radius from a given point 3D cartesian point (x,y,z)
    /// or a 2D (lon,lat) point
    template <typename Point>
    ValueList closestPointsWithinRadius( const Point& p, double radius ) const {
        return get()->closestPointsWithinRadius( p, radius );
    }

    /// @brief Return geometry used to convert (lon,lat) to (x,y,z) coordinates
    const Geometry& geometry() const { return get()->geometry(); }
};

//------------------------------------------------------------------------------------------------------

using IndexKDTree2D = KDTree<idx_t, Point2>;  // 2D search: implementation is using 2D points only
using IndexKDTree3D = KDTree<idx_t, Point3>;  // 3D search: lonlat (2D) to xyz (3D) conversion is done internally
using IndexKDTree   = IndexKDTree3D;

// ------------------------------------------------------------------
// C wrapper interfaces to C++ routines

extern "C" {
IndexKDTree::Implementation* atlas__IndexKDTree__new();
IndexKDTree::Implementation* atlas__IndexKDTree__new_geometry( const Geometry::Implementation* geometry );
void atlas__IndexKDTree__delete( IndexKDTree::Implementation* This );
void atlas__IndexKDTree__reserve( IndexKDTree::Implementation* This, const idx_t size );
void atlas__IndexKDTree__insert( IndexKDTree::Implementation* This, const double lon, const double lat,
                                 const idx_t index );
void atlas__IndexKDTree__build( IndexKDTree::Implementation* This );
void atlas__IndexKDTree__closestPoints( const IndexKDTree::Implementation* This, const double plon, const double plat,
                                        const size_t k, double*& lon, double*& lat, idx_t*& indices,
                                        double*& distances );
void atlas__IndexKDTree__closestPoint( const IndexKDTree::Implementation* This, const double plon, const double plat,
                                       double& lon, double& lat, idx_t& index, double& distance );
void atlas__IndexKDTree__closestPointsWithinRadius( const IndexKDTree::Implementation* This, const double plon,
                                                    const double plat, const double radius, size_t& k, double*& lon,
                                                    double*& lat, idx_t*& indices, double*& distances );
const Geometry::Implementation* atlas__IndexKDTree__geometry( const IndexKDTree::Implementation* This );
int atlas__IndexKDTree__empty( const IndexKDTree::Implementation* This );
idx_t atlas__IndexKDTree__size( const IndexKDTree::Implementation* This );
}

//------------------------------------------------------------------------------------------------------

}  // namespace util
}  // namespace atlas
