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

// file deepcode ignore CppMemoryLeak: static pointers for global registry are OK and will be cleaned up at end

#include "atlas/grid/detail/partitioner/Partitioner.h"

#include <map>
#include <string>

#include "eckit/thread/AutoLock.h"
#include "eckit/thread/Mutex.h"

#include "atlas/grid/Distribution.h"
#include "atlas/grid/Partitioner.h"
#include "atlas/grid/detail/distribution/DistributionArray.h"
#include "atlas/grid/detail/partitioner/BandsPartitioner.h"
#include "atlas/grid/detail/partitioner/CheckerboardPartitioner.h"
#include "atlas/grid/detail/partitioner/CubedSpherePartitioner.h"
#include "atlas/grid/detail/partitioner/EqualBandsPartitioner.h"
#include "atlas/grid/detail/partitioner/EqualRegionsPartitioner.h"
#include "atlas/grid/detail/partitioner/MatchingFunctionSpacePartitionerLonLatPolygon.h"
#include "atlas/grid/detail/partitioner/MatchingMeshPartitioner.h"
#include "atlas/grid/detail/partitioner/MatchingMeshPartitionerBruteForce.h"
#include "atlas/grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.h"
#include "atlas/grid/detail/partitioner/MatchingMeshPartitionerLonLatPolygon.h"
#include "atlas/grid/detail/partitioner/MatchingMeshPartitionerSphericalPolygon.h"
#include "atlas/grid/detail/partitioner/RegularBandsPartitioner.h"
#include "atlas/grid/detail/partitioner/SerialPartitioner.h"
#include "atlas/library/config.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/runtime/Exception.h"
#include "atlas/runtime/Log.h"
#include "atlas/util/Config.h"

#if ATLAS_HAVE_TRANS
#include "atlas/grid/detail/partitioner/TransPartitioner.h"
#endif

namespace {

static eckit::Mutex* local_mutex                                                       = nullptr;
static std::map<std::string, atlas::grid::detail::partitioner::PartitionerFactory*>* m = nullptr;
static pthread_once_t once                                                             = PTHREAD_ONCE_INIT;

static void init() {
    local_mutex = new eckit::Mutex();
    m           = new std::map<std::string, atlas::grid::detail::partitioner::PartitionerFactory*>();
}
}  // namespace

namespace atlas {
namespace grid {
namespace detail {
namespace partitioner {

int Partitioner::extract_nb_partitions(const eckit::Parametrisation& config) {
    int N;
    if (config.has("mpi_comm")) {
        std::string mpi_comm;
        config.get("mpi_comm",mpi_comm);
        N = mpi::comm(mpi_comm).size();
    }
    else {
        N = mpi::size();
    }
    config.get("partitions",N);
    config.get("nb_partitions",N);
    config.get("nb_parts",N);
    return N;
}
std::string Partitioner::extract_mpi_comm(const eckit::Parametrisation& config) {
    if (config.has("mpi_comm")) {
        std::string mpi_comm;
        config.get("mpi_comm",mpi_comm);
        return mpi_comm;
    }
    return mpi::comm().name();
}


Partitioner::Partitioner():
    Partitioner(util::NoConfig()) {
}

Partitioner::Partitioner(const idx_t nb_partitions, const eckit::Parametrisation& config):
    nb_partitions_(nb_partitions),
    mpi_comm_(extract_mpi_comm(config)) {
}

Partitioner::Partitioner(const eckit::Parametrisation& config) :
    nb_partitions_(extract_nb_partitions(config)),
    mpi_comm_(extract_mpi_comm(config)) {
}

Partitioner::~Partitioner() = default;

idx_t Partitioner::nb_partitions() const {
    return nb_partitions_;
}

Distribution Partitioner::partition(const Grid& grid) const {
    return new distribution::DistributionArray{grid, atlas::grid::Partitioner(this)};
}

std::string Partitioner::mpi_comm() const { return mpi_comm_; }

namespace {

template <typename T>
void load_builder() {
    PartitionerBuilder<T>("tmp");
}

struct force_link {
    force_link() {
        load_builder<EqualRegionsPartitioner>();
        load_builder<CheckerboardPartitioner>();
        load_builder<CubedSpherePartitioner>();
        load_builder<BandsPartitioner>();
        load_builder<EqualBandsPartitioner>();
        load_builder<RegularBandsPartitioner>();
        load_builder<SerialPartitioner>();
#if ATLAS_HAVE_TRANS
        load_builder<TransPartitioner>();
#endif
    }
};

}  // namespace

PartitionerFactory::PartitionerFactory(const std::string& name): name_(name) {
    pthread_once(&once, init);

    eckit::AutoLock<eckit::Mutex> lock(local_mutex);

    if (m->find(name) != m->end()) {
        throw_Exception("Partitioner with name [" + name + "] is already registered.", Here());
    }
    (*m)[name] = this;
}

PartitionerFactory::~PartitionerFactory() {
    eckit::AutoLock<eckit::Mutex> lock(local_mutex);
    m->erase(name_);
}

void PartitionerFactory::list(std::ostream& out) {
    pthread_once(&once, init);

    eckit::AutoLock<eckit::Mutex> lock(local_mutex);

    static force_link static_linking;

    const char* sep = "";
    for (std::map<std::string, PartitionerFactory*>::const_iterator j = m->begin(); j != m->end(); ++j) {
        out << sep << (*j).first;
        sep = ", ";
    }
}

bool PartitionerFactory::has(const std::string& name) {
    pthread_once(&once, init);

    eckit::AutoLock<eckit::Mutex> lock(local_mutex);

    static force_link static_linking;

    return (m->find(name) != m->end());
}

Partitioner* PartitionerFactory::build(const std::string& name) {
    pthread_once(&once, init);

    eckit::AutoLock<eckit::Mutex> lock(local_mutex);

    static force_link static_linking;

    std::map<std::string, PartitionerFactory*>::const_iterator j = m->find(name);

    Log::debug() << "Looking for PartitionerFactory [" << name << "]" << std::endl;

    if (j == m->end()) {
        Log::error() << "No PartitionerFactory for [" << name << "]" << '\n';
        Log::error() << "PartitionerFactories are:" << '\n';
        for (j = m->begin(); j != m->end(); ++j) {
            Log::error() << "   " << (*j).first << '\n';
        }
        Log::error() << std::flush;
        throw_Exception(std::string("No PartitionerFactory called ") + name);
    }

    return (*j).second->make();
}

Partitioner* PartitionerFactory::build(const std::string& name, const idx_t nb_partitions) {
    return build(name, nb_partitions, util::NoConfig());
}

Partitioner* PartitionerFactory::build(const std::string& name, const idx_t nb_partitions,
                                       const eckit::Parametrisation& config) {
    pthread_once(&once, init);

    eckit::AutoLock<eckit::Mutex> lock(local_mutex);

    static force_link static_linking;

    std::map<std::string, PartitionerFactory*>::const_iterator j = m->find(name);

    Log::debug() << "Looking for PartitionerFactory [" << name << "]" << '\n';

    if (j == m->end()) {
        Log::error() << "No PartitionerFactory for [" << name << "]" << '\n';
        Log::error() << "PartitionerFactories are:" << '\n';
        for (j = m->begin(); j != m->end(); ++j) {
            Log::error() << "   " << (*j).first << '\n';
        }
        throw_Exception(std::string("No PartitionerFactory called ") + name);
    }

    return (*j).second->make(nb_partitions, config);
}

Partitioner* PartitionerFactory::build(const std::string& name,
                                       const eckit::Parametrisation& config) {
    pthread_once(&once, init);

    eckit::AutoLock<eckit::Mutex> lock(local_mutex);

    static force_link static_linking;

    std::map<std::string, PartitionerFactory*>::const_iterator j = m->find(name);

    Log::debug() << "Looking for PartitionerFactory [" << name << "]" << '\n';

    if (j == m->end()) {
        Log::error() << "No PartitionerFactory for [" << name << "]" << '\n';
        Log::error() << "PartitionerFactories are:" << '\n';
        for (j = m->begin(); j != m->end(); ++j) {
            Log::error() << "   " << (*j).first << '\n';
        }
        throw_Exception(std::string("No PartitionerFactory called ") + name);
    }

    return (*j).second->make(config);
}



Partitioner* MatchingPartitionerFactory::build(const std::string& type, const Mesh& partitioned, const eckit::Parametrisation& config) {
    if (type == MatchingMeshPartitionerSphericalPolygon::static_type()) {
        return new MatchingMeshPartitionerSphericalPolygon(partitioned, config);
    }
    else if (type == MatchingMeshPartitionerLonLatPolygon::static_type()) {
        return new MatchingMeshPartitionerLonLatPolygon(partitioned, config);
    }
    else if (type == MatchingMeshPartitionerBruteForce::static_type()) {
        return new MatchingMeshPartitionerBruteForce(partitioned, config);
    }
    else if (type == MatchingMeshPartitionerCubedSphere::static_type()) {
        return new MatchingMeshPartitionerCubedSphere(partitioned, config);
    }
    else {
        ATLAS_NOTIMPLEMENTED;
    }
}

Partitioner* MatchingPartitionerFactory::build(const std::string& type, const FunctionSpace& partitioned, const eckit::Parametrisation& config) {
    if (type == MatchingFunctionSpacePartitionerLonLatPolygon::static_type()) {
        return new MatchingFunctionSpacePartitionerLonLatPolygon(partitioned, config);
    }
    else {
        ATLAS_NOTIMPLEMENTED;
    }
}

}  // namespace partitioner
}  // namespace detail
}  // namespace grid
}  // namespace atlas
