/*
 * (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 <algorithm>
#include <iomanip>
#include <sstream>
#include <vector>

#include "atlas/array.h"
#include "atlas/field.h"
#include "atlas/functionspace.h"
#include "atlas/grid.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/util/Config.h"


#include <stdio.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>


#include "tests/AtlasTestEnvironment.h"

using Grid   = atlas::Grid;
using Config = atlas::util::Config;

namespace atlas {
namespace test {

atlas::FieldSet getIJ(const atlas::functionspace::StructuredColumns& fs) {
    atlas::FieldSet ij;

    auto vi0 = atlas::array::make_view<atlas::idx_t, 1>(fs.index_i());
    auto vj0 = atlas::array::make_view<atlas::idx_t, 1>(fs.index_j());

    auto fi = fs.createField<atlas::idx_t>();
    auto fj = fs.createField<atlas::idx_t>();

    auto vi1 = atlas::array::make_view<atlas::idx_t, 1>(fi);
    auto vj1 = atlas::array::make_view<atlas::idx_t, 1>(fj);

    for (int i = 0; i < fs.size(); i++) {
        vi1(i) = vi0(i);
        vj1(i) = vj0(i);
    }

    ij.add(fi);
    ij.add(fj);

    return ij;
}


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

CASE("Two haloexchanges for StructuredColumns differing only by 'periodic_points'") {
    Grid grid("L400x200");

    grid::Distribution dist(grid, grid::Partitioner("checkerboard"));

    functionspace::StructuredColumns fs1(grid, dist, Config("halo", 1) | Config("periodic_points", true));
    functionspace::StructuredColumns fs2(grid, dist, Config("halo", 1));

    if (mpi::size() == 1) {
        EXPECT_EQ(fs1.sizeOwned(), 80000);
        EXPECT_EQ(fs2.sizeOwned(), 80000);
        EXPECT_EQ(fs1.sizeHalo(), 81406);
        EXPECT_EQ(fs2.sizeHalo(), 81204);
    }
    if (mpi::size() == 4) {
        EXPECT_EQ(fs1.sizeOwned(), 20000);
        EXPECT_EQ(fs2.sizeOwned(), 20000);
        EXPECT_EQ(fs1.sizeHalo(), (std::vector<int>{20604, 20604, 20604, 20806}[mpi::rank()]));
        EXPECT_EQ(fs2.sizeHalo(), (std::vector<int>{20604, 20604, 20604, 20604}[mpi::rank()]));
    }

    auto ij1 = getIJ(fs1);
    auto ij2 = getIJ(fs2);

    EXPECT_NO_THROW(fs1.haloExchange(ij1));
    EXPECT_NO_THROW(fs2.haloExchange(ij2));
}

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

CASE("Two haloexchanges for StructuredColumns differing only by 'distribution'") {
    Grid grid("L400x201");

    grid::Distribution dist1(grid, grid::Partitioner("checkerboard"));
    grid::Distribution dist2(grid, grid::Partitioner("regular_bands"));

    functionspace::StructuredColumns fs1(grid, dist1, Config("halo", 1));
    functionspace::StructuredColumns fs2(grid, dist2, Config("halo", 1));

    if (mpi::size() == 1) {
        EXPECT_EQ(fs1.sizeOwned(), 80400);
        EXPECT_EQ(fs2.sizeOwned(), 80400);
        EXPECT_EQ(fs1.sizeHalo(), 81606);
        EXPECT_EQ(fs2.sizeHalo(), 81606);
    }
    if (mpi::size() == 4) {
        EXPECT_EQ(fs1.sizeOwned(), 20100);
        EXPECT_EQ(fs2.sizeOwned(), (std::vector<int>{20400, 20000, 20000, 20000}[mpi::rank()]));
        EXPECT_EQ(fs1.sizeHalo(), (std::vector<int>{20706, 20706, 20706, 20706}[mpi::rank()]));
        EXPECT_EQ(fs2.sizeHalo(), (std::vector<int>{21306, 20904, 20904, 20904}[mpi::rank()]));
    }

    auto ij1 = getIJ(fs1);
    auto ij2 = getIJ(fs2);

    EXPECT_NO_THROW(fs1.haloExchange(ij1));
    EXPECT_NO_THROW(fs2.haloExchange(ij2));
}

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

}  // namespace test
}  // namespace atlas

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