/*
 * (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 "atlas/functionspace/StructuredColumns.h"

#include <functional>
#include <iomanip>
#include <numeric>
#include <sstream>
#include <string>

#include "atlas/array/MakeView.h"
#include "atlas/field/FieldSet.h"
#include "atlas/grid/Distribution.h"
#include "atlas/grid/Partitioner.h"
#include "atlas/grid/StructuredGrid.h"
#include "atlas/library/Library.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/parallel/omp/omp.h"
#include "atlas/runtime/Exception.h"
#include "atlas/runtime/Trace.h"
#include "atlas/util/CoordinateEnums.h"

namespace atlas {
namespace functionspace {
namespace detail {

namespace {

struct GridPoint {
public:
    idx_t i, j;
    idx_t r;

    GridPoint() = default;
    GridPoint(idx_t _i, idx_t _j): i(_i), j(_j) {}
    GridPoint(idx_t _i, idx_t _j, idx_t _r): i(_i), j(_j), r(_r) {}

    /* No longer used:

    bool operator<( const GridPoint& other ) const {
        if ( j < other.j ) return true;
        if ( j == other.j ) return i < other.i;
        return false;
    }

    bool operator==( const GridPoint& other ) const { return ( j == other.j && i == other.i ); }

    */
};


struct GridPointSet {
public:
    GridPointSet() = default;

    atlas::vector<GridPoint> gp_;

    void reserve(size_t size) { gp_.reserve(size); }
    void resize(size_t size) { gp_.resize(size); }
    void set(idx_t i, idx_t j, idx_t r) {
        gp_[r].i = i;
        gp_[r].j = j;
        gp_[r].r = r;
    }

    idx_t size() const { return static_cast<idx_t>(gp_.size()); }

    const GridPoint& operator[](idx_t i) const { return gp_[i]; }

    // unused:
    //using const_iterator = decltype( gp_ )::const_iterator;
    //const_iterator begin() const { return gp_.begin(); }
    //const_iterator end() const { return gp_.end(); }
};

}  // namespace


void StructuredColumns::setup(const grid::Distribution& distribution, const eckit::Configuration& config) {
    config.get("periodic_points", periodic_points_);
    if (not(*grid_)) {
        throw_Exception("Grid is not a grid::Structured type", Here());
    }

    bool periodic_x = false, periodic_y = false;
    config.get("periodic_x", periodic_x);
    config.get("periodic_y", periodic_y);

    const double eps = 1.e-12;

    ny_                  = grid_->ny();
    north_pole_included_ = 90. - grid_->y(0) == 0.;
    south_pole_included_ = 90. + grid_->y(ny_ - 1) == 0;

    distribution_ = distribution.type();

    nb_partitions_ = distribution.nb_partitions();

    int mpi_rank = int(mpi::rank());
    int mpi_size = int(mpi::size());

    j_begin_ = std::numeric_limits<idx_t>::max();
    j_end_   = std::numeric_limits<idx_t>::min();
    i_begin_.resize(grid_->ny(), std::numeric_limits<idx_t>::max());
    i_end_.resize(grid_->ny(), std::numeric_limits<idx_t>::min());
    idx_t owned(0);

    ATLAS_TRACE_SCOPE("Compute bounds owned") {
        if (mpi_size == 1) {
            j_begin_ = 0;
            j_end_   = grid_->ny();
            for (idx_t j = j_begin_; j < j_end_; ++j) {
                i_begin_[j] = 0;
                i_end_[j]   = grid_->nx(j);
            }
            owned = grid_->size();
        }
        else {
            size_t num_threads = atlas_omp_get_max_threads();
            if (num_threads == 1) {
                idx_t c(0);
                for (idx_t j = 0; j < grid_->ny(); ++j) {
                    for (idx_t i = 0; i < grid_->nx(j); ++i, ++c) {
                        if (distribution.partition(c) == mpi_rank) {
                            j_begin_    = std::min<idx_t>(j_begin_, j);
                            j_end_      = std::max<idx_t>(j_end_, j + 1);
                            i_begin_[j] = std::min<idx_t>(i_begin_[j], i);
                            i_end_[j]   = std::max<idx_t>(i_end_[j], i + 1);
                            ++owned;
                        }
                    }
                }
            }
            else {
                std::vector<idx_t> thread_reduce_j_begin(num_threads, std::numeric_limits<idx_t>::max());
                std::vector<idx_t> thread_reduce_j_end(num_threads, std::numeric_limits<idx_t>::min());
                std::vector<std::vector<idx_t>> thread_reduce_i_begin(
                    grid_->ny(), std::vector<idx_t>(num_threads, std::numeric_limits<idx_t>::max()));
                std::vector<std::vector<idx_t>> thread_reduce_i_end(
                    grid_->ny(), std::vector<idx_t>(num_threads, std::numeric_limits<idx_t>::min()));
                std::vector<idx_t> thread_reduce_owned(num_threads, 0);
                atlas_omp_parallel {
                    const idx_t thread_num = atlas_omp_get_thread_num();
                    const idx_t begin =
                        static_cast<int>(size_t(thread_num) * size_t(grid_->size()) / size_t(num_threads));
                    const idx_t end =
                        static_cast<int>(size_t(thread_num + 1) * size_t(grid_->size()) / size_t(num_threads));
                    idx_t thread_j_begin = 0;
                    std::vector<idx_t> thread_i_begin(grid_->ny());
                    std::vector<idx_t> thread_i_end(grid_->ny());
                    idx_t n = 0;
                    for (idx_t j = 0; j < grid_->ny(); ++j) {
                        if (n + grid_->nx(j) > begin) {
                            thread_j_begin    = j;
                            thread_i_begin[j] = begin - n;
                            break;
                        }
                        n += grid_->nx(j);
                    }
                    idx_t thread_j_end{thread_j_begin};
                    for (idx_t j = thread_j_begin; j < grid_->ny(); ++j) {
                        idx_t i_end = end - n;
                        if (j > thread_j_begin) {
                            thread_i_begin[j] = 0;
                        }
                        if (i_end > grid_->nx(j)) {
                            thread_i_end[j] = grid_->nx(j);
                            n += grid_->nx(j);
                        }
                        else {
                            thread_i_end[j] = i_end;
                            thread_j_end    = j + 1;
                            break;
                        }
                    }
                    auto& __j_begin = thread_reduce_j_begin[thread_num];
                    auto& __j_end   = thread_reduce_j_end[thread_num];
                    auto& __owned   = thread_reduce_owned[thread_num];
                    idx_t c         = begin;

                    for (idx_t j = thread_j_begin; j < thread_j_end; ++j) {
                        auto& __i_begin = thread_reduce_i_begin[j][thread_num];
                        auto& __i_end   = thread_reduce_i_end[j][thread_num];
                        bool j_in_partition{false};
                        for (idx_t i = thread_i_begin[j]; i < thread_i_end[j]; ++i, ++c) {
                            if (distribution.partition(c) == mpi_rank) {
                                j_in_partition = true;
                                __i_begin      = std::min<idx_t>(__i_begin, i);
                                __i_end        = std::max<idx_t>(__i_end, i + 1);
                                ++__owned;
                            }
                        }
                        if (j_in_partition) {
                            __j_begin = std::min<idx_t>(__j_begin, j);
                            __j_end   = std::max<idx_t>(__j_end, j + 1);
                        }
                    }
                    ATLAS_ASSERT(c == end);
                }
                owned    = std::accumulate(thread_reduce_owned.begin(), thread_reduce_owned.end(), 0);
                j_begin_ = *std::min_element(thread_reduce_j_begin.begin(), thread_reduce_j_begin.end());
                j_end_   = *std::max_element(thread_reduce_j_end.begin(), thread_reduce_j_end.end());
                atlas_omp_parallel_for(idx_t j = j_begin_; j < j_end_; ++j) {
                    i_begin_[j] = *std::min_element(thread_reduce_i_begin[j].begin(), thread_reduce_i_begin[j].end());
                    i_end_[j]   = *std::max_element(thread_reduce_i_end[j].begin(), thread_reduce_i_end[j].end());
                }
            }
        }
    }

    size_owned_ = owned;

    int halo = config.getInt("halo", 0);
    halo_    = halo;

    j_begin_halo_ = j_begin_ - halo;
    j_end_halo_   = j_end_ + halo;
    i_begin_halo_.resize(-halo, grid_->ny() - 1 + halo);
    i_end_halo_.resize(-halo, grid_->ny() - 1 + halo);

    auto compute_i = [this](idx_t i, idx_t j) -> idx_t {
        const idx_t nx = grid_->nx(j);
        while (i >= nx) {
            i -= nx;
        }
        while (i < 0) {
            i += nx;
        }
        return i;
    };

    auto compute_i_fast = [](idx_t i, const idx_t nx) -> idx_t {
        while (i >= nx) {
            i -= nx;
        }
        while (i < 0) {
            i += nx;
        }
        return i;
    };

    std::function<idx_t(idx_t)> compute_j;
    compute_j = [this, &compute_j, &periodic_y](idx_t j) -> idx_t {
        if (periodic_y) {
            const idx_t ny = grid_->ny();
            while (j < 0) {
                j += ny;
            }
            while (j >= ny) {
                j -= ny;
            }
        }
        else {
            if (j < 0) {
                j = (grid_->y(0) == 90.) ? -j : -j - 1;
            }
            else if (j >= grid_->ny()) {
                idx_t jlast = grid_->ny() - 1;
                j           = (grid_->y(jlast) == -90.) ? jlast - 1 - (j - grid_->ny()) : jlast - (j - grid_->ny());
            }
            if (j < 0 or j >= grid_->ny()) {
                j = compute_j(j);
            }
        }
        return j;
    };

    auto compute_x = [this, &compute_i, &compute_j](idx_t i, idx_t j) -> double {
        const idx_t jj = compute_j(j);
        const idx_t ii = compute_i(i, jj);  // guaranteed between 0 and nx(jj)
        const idx_t nx = grid_->nx(jj);
        const double a = (ii - i) / nx;
        const double x = grid_->x(ii, jj) - a * grid_->x(nx, jj);
        return x;
    };

    auto compute_x_fast = [this, &compute_i_fast](idx_t i, idx_t jj, idx_t nx) -> double {
        const idx_t ii = compute_i_fast(i, nx);  // guaranteed between 0 and nx(jj)
        const double a = (ii - i) / nx;
        const double x = grid_->x(ii, jj) - a * (grid_->x(nx, jj) - grid_->x(0, jj));
        return x;
    };

    auto compute_i_less_equal_x = [this, &eps](const double& x, idx_t j) -> idx_t {
        const double dx = grid_->dx(j);
        idx_t i         = idx_t(std::floor((x + eps - grid_->xmin(j)) / dx));
        return i;
    };

    auto compute_y = [this, &compute_j, &periodic_y](idx_t j) -> double {
        double y;
        idx_t jj;
        jj = compute_j(j);
        if (periodic_y) {
            y = grid_->y(jj);
        }
        else {
            y = (j < 0) ? 90. + (90. - grid_->y(jj)) : (j >= grid_->ny()) ? -90. + (-90. - grid_->y(jj)) : grid_->y(jj);
        }
        return y;
    };

    std::vector<gidx_t> global_offsets(grid_->ny());
    idx_t grid_idx = 0;
    for (idx_t j = 0; j < grid_->ny(); ++j) {
        global_offsets[j] = grid_idx;
        grid_idx += grid_->nx(j);
    }

    auto compute_g = [this, &global_offsets, &compute_i, &compute_j, &periodic_y](idx_t i, idx_t j) -> gidx_t {
        idx_t ii, jj;
        gidx_t g;
        jj = compute_j(j);
        ii = compute_i(i, jj);
        if (!periodic_y) {
            if (jj != j) {
                ATLAS_ASSERT(grid_->nx(jj) % 2 == 0);  // assert even number of points
                ii = (ii < grid_->nx(jj) / 2)    ? ii + grid_->nx(jj) / 2
                     : (ii >= grid_->nx(jj) / 2) ? ii - grid_->nx(jj) / 2
                                                 : ii;
            }
        }
        g = global_offsets[jj] + ii + 1;
        return g;
    };

    auto compute_p = [&compute_g, &distribution](idx_t i, idx_t j) -> int {
        return distribution.partition(compute_g(i, j) - 1);
    };

    GridPointSet gridpoints;

    ATLAS_TRACE_SCOPE("Compute mapping") {
        idx_t imin = std::numeric_limits<idx_t>::max();
        idx_t imax = -std::numeric_limits<idx_t>::max();
        idx_t jmin = std::numeric_limits<idx_t>::max();
        idx_t jmax = -std::numeric_limits<idx_t>::max();

        ATLAS_TRACE_SCOPE("Compute bounds halo") {
            for (idx_t j = j_begin_halo_; j < j_end_halo_; ++j) {
                i_begin_halo_(j) = imin;
                i_end_halo_(j)   = imax;
            }

            // Following cannot be multithreaded in current form due to race-conditions related to index jj
            for (idx_t j = j_begin_; j < j_end_; ++j) {
                for (idx_t i : {i_begin_[j], i_end_[j] - 1}) {
                    // Following line only, increases periodic halo on the east side by 1
                    if (periodic_points_ && i == grid_->nx(j) - 1) {
                        ++i;
                    }

                    double x = grid_->x(i, j);

                    double x_next = grid_->x(i + 1, j);
                    double x_prev = grid_->x(i - 1, j);
                    for (idx_t jj = j - halo; jj <= j + halo; ++jj) {
                        idx_t jjj    = compute_j(jj);
                        idx_t nx_jjj = grid_->nx(jjj);
                        idx_t last   = grid_->nx(jjj) - 1;
                        if (i == grid_->nx(j)) {
                            ++last;
                        }

                        jmin = std::min(jmin, jj);
                        jmax = std::max(jmax, jj);
                        // Compute ii as index less-equal of x
                        //
                        //              x(i,j)
                        //    |-----|-----|-----|-----|
                        // ii-halo       ii
                        //
                        //                 x(i,j)
                        //    |-----|-----|--+--|-----|
                        // ii-halo       ii

                        // idx_t ii = -halo;
                        // while ( compute_x_fast( ii, jjj, nx_jjj ) < x - eps ) {
                        //     ii++;
                        // }
                        // Question: is following implementation reproducible with above original while loop?
                        idx_t ii = compute_i_less_equal_x(x, jjj);

                        // ATLAS-186 workaround
                        // This while should not have to be there, but is here because of
                        // the MatchingMeshDomainDecomposition algorithm. that may lead to points
                        // left of the point ii.
                        while (compute_x_fast(ii - 1, jjj, nx_jjj) > x_prev + eps) {
                            --ii;
                        }

                        idx_t i_minus_halo = ii - halo;

                        // Compute iii as index less-equal of x_next
                        //
                        //               x(i,j) x_next(i,j)
                        //   |-----|-----|-+---|-+---|-----|
                        // ii-halo            iii       iii+halo
                        //
                        idx_t iii = ii;
                        while (compute_x_fast(iii + 1, jjj, nx_jjj) < x_next - eps) {
                            ++iii;
                        }
                        iii               = std::min(iii, last);
                        idx_t i_plus_halo = iii + halo;

                        imin              = std::min(imin, i_minus_halo);
                        imax              = std::max(imax, i_plus_halo);
                        i_begin_halo_(jj) = std::min(i_begin_halo_(jj), i_minus_halo);
                        i_end_halo_(jj)   = std::max(i_end_halo_(jj), i_plus_halo + 1);
                    }
                }
            }
        }

        int extra_halo{0};
        for (idx_t j = j_begin_halo_; j < j_begin_; ++j) {
            extra_halo += i_end_halo_(j) - i_begin_halo_(j);
        }
        for (idx_t j = j_begin_; j < j_end_; ++j) {
            extra_halo += i_begin_[j] - i_begin_halo_(j);
            extra_halo += i_end_halo_(j) - i_end_[j];
        }
        for (idx_t j = j_end_; j < j_end_halo_; ++j) {
            extra_halo += i_end_halo_(j) - i_begin_halo_(j);
        }


        ATLAS_TRACE_SCOPE("Assemble gridpoints") {
            gridpoints.reserve(owned + extra_halo);
            gridpoints.resize(owned);

            if (atlas_omp_get_max_threads() == 1) {
                idx_t r = 0;
                for (idx_t j = j_begin_; j < j_end_; ++j) {
                    for (idx_t i = i_begin_[j]; i < i_end_[j]; ++i, ++r) {
                        gridpoints.set(i, j, r);
                    }
                }
                ATLAS_ASSERT(r == owned);
            }
            else {
                atlas_omp_parallel {
                    const size_t num_threads = atlas_omp_get_num_threads();
                    const size_t thread_num  = atlas_omp_get_thread_num();

                    auto chunksize = size_t(owned) / num_threads;
                    size_t begin   = chunksize * thread_num;
                    size_t end     = (thread_num == num_threads - 1) ? owned : begin + chunksize;

                    idx_t thread_j_begin = 0;
                    std::vector<idx_t> thread_i_begin(grid_->ny());
                    std::vector<idx_t> thread_i_end(grid_->ny());
                    size_t n = 0;
                    for (idx_t j = j_begin_; j < j_end_; ++j) {
                        size_t j_size = static_cast<size_t>(i_end_[j] - i_begin_[j]);
                        if (n + j_size > begin) {
                            thread_j_begin    = j;
                            thread_i_begin[j] = i_begin_[j] + static_cast<idx_t>(begin - n);
                            break;
                        }
                        n += j_size;
                    }
                    idx_t thread_j_end{thread_j_begin};
                    for (idx_t j = thread_j_begin; j < j_end_; ++j) {
                        if (j > thread_j_begin) {
                            thread_i_begin[j] = i_begin_[j];
                        }
                        idx_t j_size = (i_end_[j] - i_begin_[j]);

                        idx_t remaining = static_cast<idx_t>(end - n);
                        if (j_size <= remaining) {
                            thread_i_end[j] = i_end_[j];
                            n += j_size;
                            if (n == end) {
                                goto stop;
                            }
                        }
                        else {
                            thread_i_end[j] = i_begin_[j] + remaining;
                            goto stop;
                        }

                        continue;

                    stop:
                        thread_j_end = j + 1;
                        break;
                    }

                    idx_t r = idx_t(begin);
                    for (idx_t j = thread_j_begin; j < thread_j_end; ++j) {
                        for (idx_t i = thread_i_begin[j]; i < thread_i_end[j]; ++i, ++r) {
                            gridpoints.set(i, j, r);
                        }
                    }
                    if (r != idx_t(end)) {
                        ATLAS_DEBUG_VAR(thread_num);
                        ATLAS_DEBUG_VAR(begin);
                        ATLAS_DEBUG_VAR(end);
                        ATLAS_DEBUG_VAR(r);
                    }
                    ATLAS_ASSERT(r == idx_t(end));
                }
            }

            ATLAS_ASSERT(gridpoints.size() == owned);

            gridpoints.resize(owned + extra_halo);
            idx_t r = owned;
            for (idx_t j = j_begin_halo_; j < j_begin_; ++j) {
                for (idx_t i = i_begin_halo_(j); i < i_end_halo_(j); ++i, ++r) {
                    gridpoints.set(i, j, r);
                }
            }
            for (idx_t j = j_begin_; j < j_end_; ++j) {
                for (idx_t i = i_begin_halo_(j); i < i_begin_[j]; ++i, ++r) {
                    gridpoints.set(i, j, r);
                }
                for (idx_t i = i_end_[j]; i < i_end_halo_(j); ++i, ++r) {
                    gridpoints.set(i, j, r);
                }
            }
            for (idx_t j = j_end_; j < j_end_halo_; ++j) {
                for (idx_t i = i_begin_halo_(j); i < i_end_halo_(j); ++i, ++r) {
                    gridpoints.set(i, j, r);
                }
            }

            ATLAS_ASSERT(gridpoints.size() == owned + extra_halo);
        }

        ATLAS_TRACE_SCOPE("Fill in ij2gp ") {
            ij2gp_.resize({imin, imax}, {jmin, jmax});

            atlas_omp_parallel_for(idx_t n = 0; n < gridpoints.size(); ++n) {
                const GridPoint& gp = gridpoints[n];
                ij2gp_.set(gp.i, gp.j, gp.r);
            }
        }
    }

    ATLAS_TRACE_SCOPE("Create fields") {
        size_halo_          = gridpoints.size();
        field_partition_    = Field("partition", array::make_datatype<int>(), array::make_shape(size_halo_));
        field_ghost_        = Field("ghost", array::make_datatype<int>(), array::make_shape(size_halo_));
        field_global_index_ = Field("glb_idx", array::make_datatype<gidx_t>(), array::make_shape(size_halo_));
        field_index_i_      = Field("index_i", array::make_datatype<idx_t>(), array::make_shape(size_halo_));
        field_index_j_      = Field("index_j", array::make_datatype<idx_t>(), array::make_shape(size_halo_));
        field_xy_           = Field("xy", array::make_datatype<double>(), array::make_shape(size_halo_, 2));

        auto xy         = array::make_view<double, 2>(field_xy_);
        auto part       = array::make_view<int, 1>(field_partition_);
        auto ghost      = array::make_view<int, 1>(field_ghost_);
        auto global_idx = array::make_view<gidx_t, 1>(field_global_index_);
        auto index_i    = array::make_indexview<idx_t, 1>(field_index_i_);
        auto index_j    = array::make_indexview<idx_t, 1>(field_index_j_);

        atlas_omp_parallel_for(idx_t n = 0; n < gridpoints.size(); ++n) {
            const GridPoint& gp = gridpoints[n];
            if (gp.j >= 0 && gp.j < grid_->ny()) {
                xy(gp.r, XX) = grid_->x(gp.i, gp.j);
                xy(gp.r, YY) = grid_->y(gp.j);
            }
            else {
                xy(gp.r, XX) = compute_x(gp.i, gp.j);
                xy(gp.r, YY) = compute_y(gp.j);
            }

            bool in_domain(false);
            if (gp.j >= 0 && gp.j < grid_->ny()) {
                if (gp.i >= 0 && gp.i < grid_->nx(gp.j)) {
                    in_domain        = true;
                    gidx_t k         = global_offsets[gp.j] + gp.i;
                    part(gp.r)       = distribution.partition(k);
                    global_idx(gp.r) = k + 1;
                }
            }
            if (not in_domain) {
                global_idx(gp.r) = compute_g(gp.i, gp.j);
                part(gp.r)       = compute_p(gp.i, gp.j);
            }
            index_i(gp.r) = gp.i;
            index_j(gp.r) = gp.j;
            ghost(gp.r)   = 0;
        }

        // Following short loops are not parallelized with

        for (idx_t j = j_begin_halo_; j < j_begin_; ++j) {
            for (idx_t i = i_begin_halo_(j); i < i_end_halo_(j); ++i) {
                ghost(index(i, j)) = 1;
            }
        }
        atlas_omp_parallel_for(idx_t j = j_begin_; j < j_end_; ++j) {
            for (idx_t i = i_begin_halo_(j); i < i_begin_[j]; ++i) {
                ghost(index(i, j)) = 1;
            }
            for (idx_t i = i_end_[j]; i < i_end_halo_(j); ++i) {
                ghost(index(i, j)) = 1;
            }
        }
        for (idx_t j = j_end_; j < j_end_halo_; ++j) {
            for (idx_t i = i_begin_halo_(j); i < i_end_halo_(j); ++i) {
                ghost(index(i, j)) = 1;
            }
        }
    }
}

}  // namespace detail

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

}  // namespace functionspace
}  // namespace atlas
