/*
 * (C) Copyright 1996- 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 <tuple>
#include <vector>

#include "eckit/linalg/Matrix.h"
#include "eckit/linalg/Vector.h"

#include "atlas/linalg/sparse.h"

#include "tests/AtlasTestEnvironment.h"


using namespace atlas::linalg;

namespace atlas {
namespace test {

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

// strings to be used in the tests
static std::string eckit_linalg = sparse::backend::eckit_linalg::type();
static std::string omp          = sparse::backend::omp::type();

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

// Only reason to define these derived classes is for nicer constructors and convenience in the tests

class Vector : public eckit::linalg::Vector {
public:
    using Scalar = eckit::linalg::Scalar;
    using eckit::linalg::Vector::Vector;
    Vector( const std::initializer_list<Scalar>& v ) : eckit::linalg::Vector::Vector( v.size() ) {
        size_t i = 0;
        for ( auto& s : v ) {
            operator[]( i++ ) = s;
        }
    }
};

class Matrix : public eckit::linalg::Matrix {
public:
    using Scalar = eckit::linalg::Scalar;
    using eckit::linalg::Matrix::Matrix;

    Matrix( const std::initializer_list<std::vector<Scalar>>& m ) :
        eckit::linalg::Matrix::Matrix( m.size(), m.size() ? m.begin()->size() : 0 ) {
        size_t r = 0;
        for ( auto& row : m ) {
            for ( size_t c = 0; c < cols(); ++c ) {
                operator()( r, c ) = row[c];
            }
            ++r;
        }
    }
};

// 2D array constructable from eckit::linalg::Matrix
// Indexing/memorylayout and data type can be customized for testing
template <typename Value, Indexing indexing = Indexing::layout_left>
struct ArrayMatrix {
    array::ArrayView<Value, 2>& view() { return view_; }
    ArrayMatrix( const eckit::linalg::Matrix& m ) : ArrayMatrix( m.rows(), m.cols() ) {
        for ( int r = 0; r < m.rows(); ++r ) {
            for ( int c = 0; c < m.cols(); ++c ) {
                auto& v = layout_left ? view_( r, c ) : view_( c, r );
                v       = m( r, c );
            }
        }
    }
    ArrayMatrix( int r, int c ) : array( make_shape( r, c ) ), view_( array::make_view<Value, 2>( array ) ) {}

private:
    static constexpr bool layout_left = ( indexing == Indexing::layout_left );
    static array::ArrayShape make_shape( int rows, int cols ) {
        return layout_left ? array::make_shape( rows, cols ) : array::make_shape( cols, rows );
    }
    array::ArrayT<Value> array;
    array::ArrayView<Value, 2> view_;
};

// 1D array constructable from eckit::linalg::Vector
template <typename Value>
struct ArrayVector {
    array::ArrayView<Value, 1>& view() { return view_; }
    ArrayVector( const eckit::linalg::Vector& v ) : ArrayVector( v.size() ) {
        for ( int n = 0; n < v.size(); ++n ) {
            view_[n] = v[n];
        }
    }
    ArrayVector( int size ) : array( size ), view_( array::make_view<Value, 1>( array ) ) {}

private:
    array::ArrayT<Value> array;
    array::ArrayView<Value, 1> view_;
};


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

template <typename T>
void expect_equal( T* v, T* r, size_t s ) {
    EXPECT(
        is_approximately_equal( eckit::testing::make_view( v, s ), eckit::testing::make_view( r, s ), T( 1.e-5 ) ) );
}

template <class T1, class T2>
void expect_equal( const T1& v, const T2& r ) {
    expect_equal( v.data(), r.data(), std::min( v.size(), r.size() ) );
}

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

CASE( "test introspection" ) {
    SECTION( "ArrayView" ) {
        array::ArrayT<float> array( 4, 3 );
        auto view  = array::make_view<const float, 2>( array );
        using View = decltype( view );
        static_assert( introspection::has_contiguous<View>::value, "ArrayView expected to have contiguous" );
        static_assert( introspection::has_rank<View>::value, "ArrayView expected to have rank" );
        static_assert( introspection::has_shape<View>::value, "ArrayView expected to have shape" );
        EXPECT( introspection::contiguous( view ) );
        EXPECT_EQ( introspection::shape<0>( view ), 4 );
        EXPECT_EQ( introspection::shape<1>( view ), 3 );
    }
    SECTION( "std::vector" ) {
        using View = std::vector<double>;
        static_assert( not introspection::has_contiguous<View>::value, "std::vector does not have contiguous" );
        static_assert( not introspection::has_rank<View>::value, "std::vector does not have rank" );
        static_assert( not introspection::has_shape<View>::value, "std::vector does not have shape" );
        static_assert( introspection::rank<View>() == 1, "std::vector is of rank 1" );
        auto v = View{1, 2, 3, 4};
        EXPECT( introspection::contiguous( v ) );
        EXPECT_EQ( introspection::shape<0>( v ), 4 );
        EXPECT_EQ( introspection::shape<1>( v ), 4 );
    }
    SECTION( "eckit::linlag::Vector" ) {
        using View = Vector;
        static_assert( not introspection::has_contiguous<View>::value,
                       "eckit::linalg::Vector does not have contiguous" );
        static_assert( not introspection::has_rank<View>::value, "eckit::linalg::Vector does not have rank" );
        static_assert( not introspection::has_shape<View>::value, "seckit::linalg::Vector does not have shape" );
        static_assert( introspection::rank<View>() == 1, "eckit::linalg::Vector is of rank 1" );
        auto v = Vector{1, 2, 3, 4};
        EXPECT( introspection::contiguous( v ) );
        EXPECT_EQ( introspection::shape<0>( v ), 4 );
    }
    SECTION( "eckit::linlag::Matrix" ) {
        using View = Matrix;
        static_assert( not introspection::has_contiguous<View>::value,
                       "eckit::linalg::Matrix does not have contiguous" );
        static_assert( not introspection::has_rank<View>::value, "eckit::linalg::Matrix does not have rank" );
        static_assert( not introspection::has_shape<View>::value, "seckit::linalg::Matrix does not have shape" );
        static_assert( introspection::rank<View>() == 2, "eckit::linalg::Matrix is of rank 1" );
        auto m = Matrix{{1., 2.}, {3., 4.}, {5., 6.}};
        EXPECT( introspection::contiguous( m ) );
        // Following is reversed because M has column-major ordering
        EXPECT_EQ( introspection::shape<0>( m ), 3 );
        EXPECT_EQ( introspection::shape<1>( m ), 2 );
    }
}

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

CASE( "test backend functionalities" ) {
    sparse::current_backend( omp );
    EXPECT_EQ( sparse::current_backend().type(), omp );
    EXPECT_EQ( sparse::current_backend().getString( "backend", "undefined" ), "undefined" );

    sparse::current_backend( eckit_linalg );
    EXPECT_EQ( sparse::current_backend().type(), "eckit_linalg" );
    EXPECT_EQ( sparse::current_backend().getString( "backend", "undefined" ), "undefined" );
    sparse::current_backend().set( "backend", "default" );
    EXPECT_EQ( sparse::current_backend().getString( "backend" ), "default" );

    sparse::current_backend( omp );
    EXPECT_EQ( sparse::current_backend().getString( "backend", "undefined" ), "undefined" );
    EXPECT_EQ( sparse::default_backend( eckit_linalg ).getString( "backend" ), "default" );

    sparse::default_backend( eckit_linalg ).set( "backend", "generic" );
    EXPECT_EQ( sparse::default_backend( eckit_linalg ).getString( "backend" ), "generic" );

    const sparse::Backend backend_default      = sparse::Backend();
    const sparse::Backend backend_omp          = sparse::backend::omp();
    const sparse::Backend backend_eckit_linalg = sparse::backend::eckit_linalg();
    EXPECT_EQ( backend_default.type(), omp );
    EXPECT_EQ( backend_omp.type(), omp );
    EXPECT_EQ( backend_eckit_linalg.type(), eckit_linalg );

    EXPECT_EQ( std::string( backend_omp ), omp );
    EXPECT_EQ( std::string( backend_eckit_linalg ), eckit_linalg );
}

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

CASE( "sparse_matrix vector multiply (spmv)" ) {
    // "square" matrix
    // A =  2  . -3
    //      .  2  .
    //      .  .  2
    // x = 1 2 3
    // y = 1 2 3
    SparseMatrix A{3, 3, {{0, 0, 2.}, {0, 2, -3.}, {1, 1, 2.}, {2, 2, 2.}}};

    for ( std::string backend : {omp, eckit_linalg} ) {
        sparse::current_backend( backend );

        SECTION( "test_identity [backend=" + sparse::current_backend().type() + "]" ) {
            {
                Vector y1( 3 );
                SparseMatrix B;
                B.setIdentity( 3, 3 );
                sparse_matrix_multiply( B, Vector{1., 2., 3.}, y1 );
                expect_equal( y1, Vector{1., 2., 3.} );
            }

            {
                SparseMatrix C;
                C.setIdentity( 6, 3 );
                Vector y2( 6 );
                sparse_matrix_multiply( C, Vector{1., 2., 3.}, y2 );
                expect_equal( y2, Vector{1., 2., 3.} );
                expect_equal( y2.data() + 3, Vector{0., 0., 0.}.data(), 3 );
            }

            {
                SparseMatrix D;
                D.setIdentity( 2, 3 );
                Vector y3( 2 );
                sparse_matrix_multiply( D, Vector{1., 2., 3.}, y3 );
                expect_equal( y3, Vector{1., 2., 3.} );
            }
        }


        SECTION( "eckit::Vector [backend=" + sparse::current_backend().type() + "]" ) {
            Vector y( 3 );
            sparse_matrix_multiply( A, Vector{1., 2., 3.}, y );
            expect_equal( y, Vector{-7., 4., 6.} );
            // spmv of sparse matrix and vector of non-matching sizes should fail
            EXPECT_THROWS_AS( sparse_matrix_multiply( A, Vector( 2 ), y ), eckit::AssertionFailed );
        }

        SECTION( "View of atlas::Array [backend=" + backend + "]" ) {
            ArrayVector<double> x( Vector{1., 2., 3.} );
            ArrayVector<double> y( 3 );
            sparse_matrix_multiply( A, x.view(), y.view() );
            expect_equal( y.view(), Vector{-7., 4., 6.} );
            // sparse_matrix_multiply of sparse matrix and vector of non-matching sizes should fail
            {
                ArrayVector<double> x2( 2 );
                EXPECT_THROWS_AS( sparse_matrix_multiply( A, x2.view(), y.view() ), eckit::AssertionFailed );
            }
        }
    }
}

CASE( "sparse_matrix matrix multiply (spmm)" ) {
    // "square"
    // A =  2  . -3
    //      .  2  .
    //      .  .  2
    // x = 1 2 3
    // y = 1 2 3
    SparseMatrix A{3, 3, {{0, 0, 2.}, {0, 2, -3.}, {1, 1, 2.}, {2, 2, 2.}}};
    Matrix m{{1., 2.}, {3., 4.}, {5., 6.}};
    Matrix c_exp{{-13., -14.}, {6., 8.}, {10., 12.}};

    for ( std::string backend : {omp, eckit_linalg} ) {
        sparse::current_backend( backend );

        SECTION( "eckit::Matrix [backend=" + sparse::current_backend().type() + "]" ) {
            auto c = Matrix( 3, 2 );
            sparse_matrix_multiply( A, m, c );
            expect_equal( c, c_exp );
            // spmm of sparse matrix and matrix of non-matching sizes should fail
            EXPECT_THROWS_AS( sparse_matrix_multiply( A, Matrix( 2, 2 ), c ), eckit::AssertionFailed );
        }


        SECTION( "View of eckit::Matrix [backend=" + backend + "]" ) {
            auto c  = Matrix( 3, 2 );
            auto mv = atlas::linalg::make_view( m );  // convert Matrix to array::LocalView<double,2>
            auto cv = atlas::linalg::make_view( c );
            sparse_matrix_multiply( A, mv, cv, Indexing::layout_right );
            expect_equal( c, c_exp );
        }

        SECTION( "View of atlas::Array PointsRight [backend=" + sparse::current_backend().type() + "]" ) {
            ArrayMatrix<double, Indexing::layout_right> ma( m );
            ArrayMatrix<double, Indexing::layout_right> c( 3, 2 );
            sparse_matrix_multiply( A, ma.view(), c.view(), Indexing::layout_right );
            expect_equal( c.view(), c_exp );
        }
    }

    SECTION( "sparse_matrix_multiply [backend=omp]" ) {
        sparse::current_backend( eckit_linalg );  // expected to be ignored further
        auto backend = sparse::backend::omp();
        ArrayMatrix<float> ma( m );
        ArrayMatrix<float> c( 3, 2 );
        sparse_matrix_multiply( A, ma.view(), c.view(), backend );
        expect_equal( c.view(), ArrayMatrix<float>( c_exp ).view() );
    }

    SECTION( "SparseMatrixMultiply [backend=omp] 1" ) {
        sparse::current_backend( eckit_linalg );  // expected to be ignored
        auto spmm = SparseMatrixMultiply{sparse::backend::omp()};
        ArrayMatrix<float> ma( m );
        ArrayMatrix<float> c( 3, 2 );
        spmm( A, ma.view(), c.view() );
        expect_equal( c.view(), ArrayMatrix<float>( c_exp ).view() );
    }

    SECTION( "SparseMatrixMultiply [backend=omp] 2" ) {
        sparse::current_backend( eckit_linalg );  // expected to be ignored
        auto spmm = SparseMatrixMultiply{omp};
        ArrayMatrix<float> ma( m );
        ArrayMatrix<float> c( 3, 2 );
        spmm( A, ma.view(), c.view() );
        expect_equal( c.view(), ArrayMatrix<float>( c_exp ).view() );
    }
}

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

}  // namespace test
}  // namespace atlas

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