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

#include "eckit/log/Bytes.h"

#include "atlas/array.h"
#include "atlas/functionspace.h"
#include "atlas/functionspace/PointCloud.h"
#include "atlas/grid.h"
#include "atlas/interpolation.h"
#include "atlas/linalg/sparse.h"
#include "atlas/mesh.h"
#include "atlas/meshgenerator.h"
#include "atlas/util/CoordinateEnums.h"

#include "tests/AtlasTestEnvironment.h"

using namespace eckit;
using namespace atlas::functionspace;
using namespace atlas::util;

namespace atlas {
namespace test {

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

void set_field(Field& field, Grid& grid, std::function<double(double, double)> func) {
    auto view = array::make_view<double, 1>(field);
    idx_t j{0};
    for (auto& p : grid.lonlat()) {
        view(j) = func(p.lon(), p.lat());
        ++j;
    }
}

void set_field(Field& field, double v) {
    auto view = array::make_view<double, 1>(field);
    for (idx_t j = 0; j < view.shape(0); ++j) {
        view(j) = v;
    }
}


void check_field(Field& field, Grid& grid, std::function<double(double, double)> func, double tolerance) {
    auto view = array::make_view<double, 1>(field);
    idx_t j{0};
    for (auto& p : grid.lonlat()) {
        EXPECT_APPROX_EQ(view(j), func(p.lon(), p.lat()), tolerance);
        ++j;
    }
}

interpolation::MatrixCache get_or_create_cache(Grid& source, Grid& target) {
    ATLAS_TRACE();
    static std::map<std::string, interpolation::Cache> map;
    std::string key = source.uid() + target.uid();
    if (map.find(key) == map.end()) {
        // Create interpolator, extract matrix, and store it in map
        map[key] = interpolation::MatrixCache(Interpolation(option::type("finite-element"), source, target));
    }
    return map[key];
}

auto func = [](double x, double) -> double { return std::sin(x * M_PI / 180.); };

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

CASE("extract cache and use") {
    Grid grid_source("F32");
    Grid grid_target("F16");

    Field field_source("source", array::make_datatype<double>(), array::make_shape(grid_source.size()));
    Field field_target("target", array::make_datatype<double>(), array::make_shape(grid_target.size()));

    set_field(field_source, grid_source, func);

    interpolation::Cache cache = get_or_create_cache(grid_source, grid_target);

    Log::info() << "Cache contains " << eckit::Bytes(cache.footprint()) << std::endl;

    ATLAS_TRACE_SCOPE("Interpolate with cache") {
        Interpolation interpolation_using_cache(option::type("finite-element"), grid_source, grid_target, cache);
        interpolation_using_cache.execute(field_source, field_target);
    }

    check_field(field_target, grid_target, func, 1.e-4);
}

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

CASE("extract cache, copy it, and move it for use") {
    Grid grid_source("F32");
    Grid grid_target("F16");

    Field field_source("source", array::make_datatype<double>(), array::make_shape(grid_source.size()));
    Field field_target("target", array::make_datatype<double>(), array::make_shape(grid_target.size()));

    set_field(field_source, grid_source, func);

    interpolation::MatrixCache created_cache = get_or_create_cache(grid_source, grid_target);
    eckit::linalg::SparseMatrix eckit_matrix_view = atlas::linalg::make_non_owning_eckit_sparse_matrix(created_cache.matrix());
    eckit::linalg::SparseMatrix eckit_matrix_copy = atlas::linalg::make_eckit_sparse_matrix(created_cache.matrix());

    EXPECT(not created_cache.matrix().empty());
    EXPECT(not eckit_matrix_view.empty());
    EXPECT(not eckit_matrix_copy.empty());

    auto cache = interpolation::MatrixCache(atlas::linalg::make_sparse_matrix_storage(std::move(eckit_matrix_copy)));

    EXPECT(not eckit_matrix_view.empty()); // We didn't touch the view
    EXPECT(eckit_matrix_copy.empty());     // This has been moved into the new 'cache'

    ATLAS_TRACE_SCOPE("Interpolate with cache") {
        Interpolation interpolation_using_cache(option::type("finite-element"), grid_source, grid_target, cache);
        interpolation_using_cache.execute(field_source, field_target);
    }

    check_field(field_target, grid_target, func, 1.e-4);
}

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

CASE("extract cache, copy it, and pass non-owning pointer") {
    Grid grid_source("F32");
    Grid grid_target("F16");

    Field field_source("source", array::make_datatype<double>(), array::make_shape(grid_source.size()));
    Field field_target("target", array::make_datatype<double>(), array::make_shape(grid_target.size()));

    set_field(field_source, grid_source, func);

    interpolation::MatrixCache created_cache = get_or_create_cache(grid_source, grid_target);
    const auto& matrix_storage = created_cache.matrix();
    atlas::linalg::SparseMatrixStorage matrix_storage_copy(matrix_storage);
    auto matrix = atlas::linalg::make_host_view<eckit::linalg::Scalar,eckit::linalg::Index>(matrix_storage_copy);

    EXPECT(not matrix.empty());

    ATLAS_TRACE_SCOPE("Interpolate with [eckit_linalg] sparse_matrix_multiply of ArrayView") {
        atlas::linalg::sparse::current_backend("eckit_linalg");
        atlas::linalg::sparse::default_backend("eckit_linalg").set("backend", "generic");
        auto src = array::make_view<double, 1>(field_source);
        auto tgt = array::make_view<double, 1>(field_target);
        atlas::linalg::sparse_matrix_multiply(matrix, src, tgt);
    }
    check_field(field_target, grid_target, func, 1.e-4);
    set_field(field_target, 0.);
    ATLAS_TRACE_SCOPE("Interpolate with [openmp] sparse_matrix_multiply of eckit::linalg::Vector") {
        atlas::linalg::sparse::current_backend("openmp");
        eckit::linalg::Vector src{array::make_view<double, 1>(field_source).data(), field_source.size()};
        eckit::linalg::Vector tgt{array::make_view<double, 1>(field_target).data(), field_target.size()};
        atlas::linalg::sparse_matrix_multiply(matrix, src, tgt);
    }
    check_field(field_target, grid_target, func, 1.e-4);
    set_field(field_target, 0.);

    auto cache = interpolation::MatrixCache( &matrix_storage_copy ); 

    EXPECT(not matrix.empty());

    ATLAS_TRACE_SCOPE("Interpolate with cache") {
        set_field(field_target, 0.);
        Interpolation interpolation_using_cache(option::type("finite-element"), grid_source, grid_target, cache);
        interpolation_using_cache.execute(field_source, field_target);
    }
    check_field(field_target, grid_target, func, 1.e-4);
    set_field(field_target, 0.);
}

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

}  // namespace test
}  // namespace atlas

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