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

#pragma once

#include <array>

#include "atlas/parallel/mpi/mpi.h"
#include "atlas/runtime/Trace.h"

#define ATLAS_TRACE_MPI(...)

#if ATLAS_HAVE_TRACE

#include "atlas/util/detail/BlackMagic.h"

#undef ATLAS_TRACE_MPI
#define ATLAS_TRACE_MPI(...) ATLAS_TRACE_MPI_(::atlas::mpi::Trace, Here(), __VA_ARGS__)
#define ATLAS_TRACE_MPI_(Type, location, operation, ...) \
    __ATLAS_TYPE_SCOPE(Type, location, __ATLAS_TRACE_MPI_ENUM(operation) __ATLAS_COMMA_ARGS(__VA_ARGS__))

#define __ATLAS_TRACE_MPI_ENUM(operation) ::atlas::mpi::Operation::__ATLAS_STRINGIFY(operation)

#endif

namespace atlas {
namespace mpi {

struct StatisticsTimerTraits {
    using Barriers = runtime::trace::NoBarriers;
    using Tracing  = runtime::trace::NoLogging;
};

enum class Operation
{
    BROADCAST,
    ALLREDUCE,
    ALLGATHER,
    ALLTOALL,
    REDUCE,
    GATHER,
    SCATTER,
    BARRIER,
    SENDRECEIVE,
    ISEND,
    IRECEIVE,
    WAIT,
    _COUNT_
};

static const std::string& name(Operation c) {
    static std::array<std::string, static_cast<size_t>(Operation::_COUNT_)> names{
        "mpi.broadcast", "mpi.allreduce", "mpi.allgather",   "mpi.alltoall", "mpi.reduce",   "mpi.gather",
        "mpi.scatter",   "mpi.barrier",   "mpi.sendreceive", "mpi.isend",    "mpi.ireceive", "mpi.wait"};
    return names[static_cast<size_t>(c)];
}

class Trace : public runtime::trace::TraceT<StatisticsTimerTraits> {
    using Base = runtime::trace::TraceT<StatisticsTimerTraits>;

public:
    Trace(const eckit::CodeLocation& loc, Operation c): Base(loc, name(c), make_labels(c)) {}
    Trace(const eckit::CodeLocation& loc, Operation c, const std::string& title): Base(loc, title, make_labels(c)) {}

private:
    static std::vector<std::string> make_labels(Operation c) { return {"mpi", name(c)}; }
};

}  // namespace mpi
}  // namespace atlas
