/*
 * (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 "eckit/geo/grid/reduced/HEALPix.h"

#include <algorithm>
#include <cctype>
#include <cmath>
#include <regex>
#include <utility>

#include "eckit/geo/Exceptions.h"
#include "eckit/geo/cache/MemoryCache.h"
#include "eckit/geo/iterator/Reduced.h"
#include "eckit/geo/iterator/Unstructured.h"
#include "eckit/geo/util.h"
#include "eckit/spec/Custom.h"
#include "eckit/utils/SafeCasts.h"


namespace eckit::geo::grid::reduced {


static const std::string HEALPIX_PATTERN = "h[rn][1-9][0-9]*|h[1-9][0-9]*[rn]?";


const std::vector<double>& healpix_longitudes(size_t Nside, size_t j) {
    using cache_t = cache::MemoryCacheT<std::pair<size_t, size_t>, std::vector<double>>;
    const cache_t::key_type key{Nside, j};

    static cache_t cache;
    if (cache.contains(key)) {
        return cache[key];
    }

    ASSERT(0 < Nside);

    const auto Nj = 4 * Nside - 1;
    ASSERT(j < Nj);

    const auto Ni = j < Nside ? 4 * (j + 1) : j < 3 * Nside ? 4 * Nside : 4 * (4 * Nside - 1 - j);
    ASSERT(0 < Ni);

    const auto step  = PointLonLat::FULL_ANGLE / static_cast<double>(Ni);
    const auto start = j < Nside || 3 * Nside - 1 < j || static_cast<bool>((j + Nside) % 2) ? step / 2. : 0.;

    std::vector<double> lons(Ni);
    std::generate_n(lons.begin(), Ni,
                    [start, step, n = 0ULL]() mutable { return start + static_cast<double>(n++) * step; });

    return (cache[key] = std::move(lons));
}


HEALPix::HEALPix(const Spec& spec) :
    HEALPix(into_unsigned(spec.get_long("Nside")), spec.get_string("order", order::HEALPix::order_default())) {}


HEALPix::HEALPix(size_t Nside, order_type order) : Nside_(Nside), healpix_(order) {
    boundingBox(new BoundingBox{});
    if (Nside_ == 0) {
        throw exception::GridError("HEALPix: Nside must be greater than zero", Here());
    }
}


Grid::iterator HEALPix::cbegin() const {
    if (healpix_.order() == order::HEALPix::RING) {
        return iterator{new geo::iterator::Reduced(*this, 0)};
    }

    // order=nested: cache full latitudes/longitudes for the unstructured iterator
    if (nested_latitudes_.empty()) {
        const auto ren = order::HEALPix{}.reorder(order(), Nside_);

        nested_latitudes_.resize(size());
        nested_longitudes_.resize(size());

        const auto& latvec = latitudes();
        for (size_t j = 0, k = 0; j < latvec.size(); ++j) {
            for (const auto lon : longitudes(j)) {
                nested_latitudes_[ren.at(k)]  = latvec[j];
                nested_longitudes_[ren.at(k)] = lon;
                ++k;
            }
        }
    }

    return iterator{new geo::iterator::Unstructured(*this, 0, nested_longitudes_, nested_latitudes_)};
}


Grid::iterator HEALPix::cend() const {
    if (healpix_.order() == order::HEALPix::RING) {
        return iterator{new geo::iterator::Reduced(*this, size())};
    }
    return iterator{new geo::iterator::Unstructured(*this)};
}


size_t HEALPix::nx(size_t j) const {
    ASSERT(j < ny());
    return j < Nside_ ? 4 * (j + 1) : j < 3 * Nside_ ? 4 * Nside_ : nx(ny() - 1 - j);
}


size_t HEALPix::ny() const {
    return 4 * Nside_ - 1;
}


Grid::Spec* HEALPix::spec(const std::string& name) {
    static const std::regex rex("[1-9][0-9]*");

    std::smatch match;
    ASSERT(std::regex_search(name, match, rex));

    auto Nside  = std::stoul(match.str());
    auto nested = (name.find("n") != std::string::npos || name.find("N") != std::string::npos);

    return new spec::Custom{
        {"type", "HEALPix"}, {"Nside", Nside}, {"order", nested ? order::HEALPix::NESTED : order::HEALPix::RING}};
}


size_t HEALPix::size_from_nside(size_t Nside) {
    return 12 * Nside * Nside;
}


size_t HEALPix::nside_from_size(size_t size) {
    if (auto Nside = static_cast<size_t>(std::sqrt(size / 12)); size == size_from_nside(Nside)) {
        return Nside;
    }

    throw exception::GridError("HEALPix: invalid size: " + std::to_string(size), Here());
}


size_t HEALPix::size() const {
    return size_from_nside(Nside_);
}


std::vector<Point> HEALPix::to_points() const {
    // This isn't very efficient
    std::vector<Point> points;
    points.reserve(size());

    if (!nested_latitudes_.empty() || !nested_longitudes_.empty()) {
        ASSERT(nested_latitudes_.size() == size() && nested_longitudes_.size() == size());
        for (size_t i = 0; i < size(); ++i) {
            points.emplace_back(PointLonLat{nested_longitudes_[i], nested_latitudes_[i]});
        }
        return points;
    }

    auto [lats, lons] = to_latlons();
    for (size_t i = 0; i < size(); ++i) {
        points.emplace_back(PointLonLat{lons[i], lats[i]});
    }

    return points;
}


const std::vector<double>& HEALPix::latitudes() const {
    const auto Nj = ny();

    if (healpix_latitudes_.empty()) {
        healpix_latitudes_.resize(Nj);

        auto i = healpix_latitudes_.begin();
        auto j = healpix_latitudes_.rbegin();
        for (size_t ring = 1; ring < 2 * Nside_; ++ring, ++i, ++j) {
            const auto f = ring < Nside_
                               ? 1. - static_cast<double>(ring * ring) / (3 * static_cast<double>(Nside_ * Nside_))
                               : 4. / 3. - 2 * static_cast<double>(ring) / (3 * static_cast<double>(Nside_));

            *i = PointLonLat::RIGHT_ANGLE - util::RADIAN_TO_DEGREE * std::acos(f);
            *j = -*i;
        }
        *i = 0.;
    }

    ASSERT(healpix_latitudes_.size() == Nj);
    return healpix_latitudes_;
}


const std::vector<double>& HEALPix::longitudes(size_t j) const {
    return healpix_longitudes(Nside_, j);
}


void HEALPix::fill_spec(spec::Custom& custom) const {
    custom.set("grid", "H" + std::to_string(Nside_));
    if (const auto o = order(); o != order::HEALPix::order_default()) {
        custom.set("order", o);
    }
}


const std::string& HEALPix::type() const {
    static const std::string type{"HEALPix"};
    return type;
}


static const GridRegisterType<HEALPix> GRIDTYPE1("HEALPix");
static const GridRegisterType<HEALPix> GRIDTYPE2("healpix");
static const auto GRIDNAME = GridRegisterName<HEALPix>(HEALPIX_PATTERN);


}  // namespace eckit::geo::grid::reduced
