/*
 * (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.
 */

#include <array>
#include <cmath>
#include <limits>
#include <numeric>
#include <vector>

#include "atlas/grid.h"
#include "atlas/util/KDTree.h"

#include "tests/AtlasTestEnvironment.h"

using namespace atlas::util;

namespace atlas {
namespace test {

//------------------------------------------------------------------------------------------------
namespace {  // helpers

template <typename T>
std::string to_string( const std::vector<T>& vector ) {
    std::stringstream s;
    s << vector;
    return s.str();
}
template <typename T>
std::string to_string( const T& v ) {
    return std::to_string( v );
}

class PayloadGenerator {
public:
    class iterator {
        friend class PayloadGenerator;

    public:
        long int operator*() const { return i_; }
        const iterator& operator++() {
            ++i_;
            return *this;
        }
        /*
        // unused
        iterator operator++( int ) {
            iterator copy( *this );
            ++i_;
            return copy;
        }*/

        bool operator==( const iterator& other ) const { return i_ == other.i_; }
        bool operator!=( const iterator& other ) const { return i_ != other.i_; }

    protected:
        iterator( long int start ) : i_( start ) {}

    private:
        unsigned long i_;
    };

    iterator begin() const { return begin_; }
    iterator end() const { return end_; }
    PayloadGenerator( long int end ) : begin_( 0 ), end_( end ) {}

    template <typename Container, typename Value = typename Container::value_type>
    void fill( Container& container ) {
        std::iota( container.begin(), container.end(), Value( *begin_ ) );
    }

    template <typename Value, size_t size>
    std::array<Value, size> make_array() {
        std::array<Value, size> array;
        fill( array );
        return array;
    }

private:
    iterator begin_;
    iterator end_;
};

static double radius() {
    return util::Earth::radius();
}

static Geometry& geometry() {
    static Geometry _geometry( radius() );
    return _geometry;
}

static PointXYZ make_xyz( const PointLonLat& lonlat ) {
    return geometry().xyz( lonlat );
}

static std::array<double, 7>& test_lon() {
    static std::array<double, 7> lon = {0., 30., 60., 90., 120., 150., 180.};
    return lon;
}

static std::array<double, 7>& test_lat() {
    static std::array<double, 7> lat = {90., 60., 30., 0., -30., -60., -90.};
    return lat;
}

static std::array<PointLonLat, 7>& test_lonlat() {
    static std::array<PointLonLat, 7> lonlat{PointLonLat{0., 90.},   PointLonLat{30., 60.},   PointLonLat{60., 30.},
                                             PointLonLat{90., 0.},   PointLonLat{120., -30.}, PointLonLat{150., -60.},
                                             PointLonLat{180., -90.}};
    return lonlat;
}

static std::array<PointXYZ, 7>& test_xyz() {
    static std::array<PointXYZ, 7> xyz{make_xyz( PointLonLat{0., 90.} ),    make_xyz( PointLonLat{30., 60.} ),
                                       make_xyz( PointLonLat{60., 30.} ),   make_xyz( PointLonLat{90., 0.} ),
                                       make_xyz( PointLonLat{120., -30.} ), make_xyz( PointLonLat{150., -60.} ),
                                       make_xyz( PointLonLat{180., -90.} )};
    return xyz;
}

static std::array<idx_t, 7>& test_payloads() {
    static auto payloads = PayloadGenerator( 7 ).make_array<idx_t, 7>();
    EXPECT_EQ( std::distance( payloads.begin(), payloads.end() ), 7 );
    return payloads;
}

void validate( const KDTree<idx_t>& tree ) {
    EXPECT_NO_THROW( tree.closestPoint( PointLonLat{180., 45.} ) );
    // Search 4 nearest neighbours (k=4), sorted by shortest distance
    auto neighbours        = tree.closestPoints( PointLonLat{89.9, 44.9}, 4 );
    auto expected_payloads = std::vector<idx_t>{2, 1, 3, 0};
    EXPECT_EQ( neighbours.payloads(), expected_payloads );
}

static const IndexKDTree& search() {
    static IndexKDTree kdtree = []() {
        IndexKDTree kdtree( geometry() );
        auto grid = Grid{"O32"};
        kdtree.build( grid.lonlat(), PayloadGenerator( grid.size() ) );
        return kdtree;
    }();
    return kdtree;
}

}  // namespace
//------------------------------------------------------------------------------------------------

static bool ECKIT_515_implemented = ( ATLAS_ECKIT_VERSION_INT >= 11302 );  // version 1.13.2
// --> implements eckit::KDTree::size() and eckit::KDTree::empty()

CASE( "test kdtree" ) {
    auto grid = Grid{"O32"};

    IndexKDTree search( geometry() );
    EXPECT( search.empty() );

    search.reserve( grid.size() );
    idx_t n{0};
    for ( auto& point : grid.lonlat() ) {
        search.insert( point, n++ );
        if ( ECKIT_515_implemented ) {
            EXPECT( search.empty() );
        }
    }
    search.build();
    EXPECT_EQ( search.size(), grid.size() );
    EXPECT_NO_THROW( search.closestPoint( PointLonLat{180., 45.} ) );
    // ...
    // Search 4 nearest neighbours (k=4), sorted by shortest distance
    auto neighbours          = search.closestPoints( PointLonLat{180., 45.}, 4 ).payloads();
    auto expected_neighbours = std::vector<idx_t>{760, 842, 759, 761};
    EXPECT_EQ( neighbours, expected_neighbours );
}

CASE( "test assertion" ) {
    auto grid = Grid{"O32"};

    IndexKDTree search( geometry() );
    search.reserve( grid.size() );
    idx_t n{0};
    for ( auto& point : grid.lonlat() ) {
        search.insert( point, n++ );
    }
    // Forgot to call search.build() --> assertion thrown when trying to access
    EXPECT_THROWS_AS( search.closestPoint( PointLonLat{180., 45.} ), eckit::AssertionFailed );
}

CASE( "test no assertion" ) {
    // Like case "test assertion", but without reserving size
    auto grid = Grid{"O32"};

    IndexKDTree search( geometry() );
    // No search.reserve() --> build() will not be necessary.
    EXPECT( search.empty() );
    idx_t n{0};
    for ( auto& point : grid.lonlat() ) {
        search.insert( point, n++ );
        if ( ECKIT_515_implemented ) {
            EXPECT_EQ( search.size(), n );
        }
    }
    EXPECT_EQ( search.size(), grid.size() );
    // search.build() Not required
    EXPECT_NO_THROW( search.closestPoint( PointLonLat{180., 45.} ) );
}

CASE( "test kdtree building with separate lon and lat and payload arrays" ) {
    IndexKDTree search( geometry() );
    search.build( test_lon(), test_lat(), test_payloads() );
    validate( search );
}

CASE( "test kdtree building with separate lon and lat and raw payload iterators" ) {
    IndexKDTree search( geometry() );
    auto lon       = test_lon();
    auto lat       = test_lat();
    auto payloads_ = test_payloads();
    search.build( lon.begin(), lon.end(), lat.begin(), lat.end(), payloads_.begin(), payloads_.end() );
    validate( search );
}

CASE( "test kdtree building with separate PointLonLat and payload containers" ) {
    IndexKDTree search( geometry() );
    search.build( test_lonlat(), test_payloads() );
    validate( search );
}

CASE( "test kdtree building with separate PointXYZ and payload containers" ) {
    IndexKDTree search( geometry() );
    search.build( test_xyz(), test_payloads() );
    validate( search );
}

CASE( "test assignment" ) {
    IndexKDTree search;
    search = IndexKDTree();
    search.build( test_lonlat(), test_payloads() );
    validate( search );
}

CASE( "test closestPoint" ) {
    auto neighbour          = search().closestPoint( PointLonLat{180., 45.} ).payload();
    auto expected_neighbour = 760;
    EXPECT_EQ( neighbour, expected_neighbour );
}

CASE( "test closestPoints" ) {
    auto neighbours          = search().closestPoints( PointLonLat{180., 45.}, 4 ).payloads();
    auto expected_neighbours = std::vector<idx_t>{760, 842, 759, 761};
    EXPECT_EQ( neighbours, expected_neighbours );
}

CASE( "test closestPointsWithinRadius" ) {
    double km                = 1000. * radius() / util::Earth::radius();
    auto neighbours          = search().closestPointsWithinRadius( PointLonLat{180., 45.}, 500 * km ).payloads();
    auto expected_neighbours = std::vector<idx_t>{760, 842, 759, 761, 841, 843, 682};
    EXPECT_EQ( neighbours, expected_neighbours );
}

CASE( "test compatibility with external eckit KDTree" ) {
    // External world
    struct ExternalKDTreeTraits {
        using Point   = Point3;
        using Payload = size_t;
    };
    using ExternalKDTree = typename eckit::KDTreeMemory<ExternalKDTreeTraits>;
    auto external_kdtree = std::make_shared<ExternalKDTree>();

    // Atlas world
    IndexKDTree search( external_kdtree, geometry() );

    // Construction of tree; could happen in Atlas world or External world separately
    auto grid = Grid{"O32"};
    search.build( grid.lonlat(), PayloadGenerator( grid.size() ) );

    // Usage in External world
    {
        auto neighbours = external_kdtree->kNearestNeighbours( make_xyz( {180., 45.} ), 4 );
        auto payloads   = std::vector<ExternalKDTree::Payload>{};
        for ( auto& neighbour : neighbours ) {
            payloads.push_back( neighbour.payload() );
        }
        auto expected_payloads = std::vector<ExternalKDTree::Payload>{760, 842, 759, 761};
        EXPECT_EQ( payloads, expected_payloads );
    }

    // Usage in Atlas world
    {
        auto payloads          = search.closestPoints( PointLonLat{180., 45.}, 4 ).payloads();
        auto expected_payloads = std::vector<IndexKDTree::Payload>{760, 842, 759, 761};
        EXPECT_EQ( payloads, expected_payloads );
    }
}

CASE( "test IndexKDTree 2D vs 3D" ) {
    IndexKDTree2D search2d( geometry() );
    IndexKDTree3D search3d( geometry() );
    search2d.build( test_lonlat(), test_payloads() );
    search3d.build( test_lonlat(), test_payloads() );
    auto payloads2d = search2d.closestPoints( PointLonLat{89.9, 44.9}, 4 ).payloads();
    auto payloads3d = search3d.closestPoints( PointLonLat{89.9, 44.9}, 4 ).payloads();
    EXPECT_EQ( payloads2d, ( std::vector<idx_t>{2, 3, 1, 4} ) );
    EXPECT_EQ( payloads3d, ( std::vector<idx_t>{2, 1, 3, 0} ) );
    // Note that the expected values are different whether 2D search or 3D search is used
}

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

}  // namespace test
}  // namespace atlas

int main( int argc, char** argv ) {
    return atlas::test::run( argc, argv );
}
