// -*- C++ -*-
//
// This file is part of YODA -- Yet more Objects for Data Analysis
// Copyright (C) 2008-2024 The YODA collaboration (see AUTHORS for details)
//
#ifndef YODA_FillableStorage_H
#define YODA_FillableStorage_H

#include "YODA/BinnedStorage.h"
#include "YODA/Dbn.h"

namespace YODA {

  /// @brief Type to adapter mapping used when user didn't provide type adapter.
  template <size_t FillDim, typename BinT>
  struct defaultAdapter; /// @note Error points here if there is no default adapter.


  /// @brief Default fill adapter for binned type double
  template <size_t FillDim, template<size_t, typename, typename> class BinT, typename BinningT, size_t N>
  struct defaultAdapter<FillDim, BinT<N, double, BinningT>> {

    using AdapterT = std::function<void(
      BinT<N, double, BinningT>&,
      typename BinningT::EdgeTypesTuple&&, double, double)>;

    AdapterT _adapter = [](auto& storedNumber, auto&& /* coords */,
                           double weight, double /* fraction */) {
      storedNumber = storedNumber + weight;
    };
  };

  namespace {
    /// @brief Aux method to create a tuple of doubles,
    /// used as additional Dbn coordinates in profiles
    template <size_t... Is>
    constexpr auto dblPadding(std::index_sequence<Is...>) {
      return std::tuple< std::decay_t<decltype((void)Is, std::declval<double>())>...>();
    }
  }

  /// @brief Default fill adapter for binned type Dbn<N>
  ///
  /// @note When non fundamental type is used, template parameter differ.
  template <size_t DbnN, template<size_t, typename, typename> class BinT, typename BinningT, size_t N>
  struct defaultAdapter<DbnN, BinT<N, Dbn<DbnN>, BinningT>> {
    static_assert(DbnN >= N, "Dimension of the Dbn needs to be at least as high as binning dimension!");

    /// The fill coordinates are a tuple of
    /// -) the bin edge types (in case of histograms)
    /// -) the bin edge types and additional doubles (in case of profiles)
    using FillCoords = decltype(std::tuple_cat(std::declval<typename BinningT::EdgeTypesTuple>(),
                     dblPadding(std::make_index_sequence<DbnN-BinningT::Dimension::value>{})));

    using AdapterT = std::function<void(BinT<N, Dbn<DbnN>, BinningT>&, FillCoords&&, double, double)>;

    /// @note This adapter nullifies discrete coordinates and transforms
    /// coordinate tuple to std::array<double, N> to pass it to the
    /// stored Dbn object. The coordinate is moved to avoid copying.
    /// In cases when continuous axis' edges are not of double type
    /// that may positively affect performance.
    AdapterT _adapter = [](auto& dbn, auto&& coords, double weight, double fraction) {
      using CoordsArrT = std::array<double, DbnN>;
      CoordsArrT dblCoords;
      size_t binIdx = dbn.index();
      auto nullifyDiscrete = [&dblCoords, &binIdx, &coords](auto I){
        using isArithmetic = typename BinningT::template is_Arithmetic<I>;
        dblCoords[I] = nullifyIfDiscCoord(std::move(std::get<I>(coords)),
                                          std::integral_constant<bool,
                                          isArithmetic::value>{}, binIdx);
      };
      MetaUtils::staticFor<BinningT::Dimension::value>(nullifyDiscrete);
      if constexpr (DbnN > BinningT::Dimension::value)
        dblCoords[N] = std::move(std::get<N>(coords));
      dbn.fill(std::move(dblCoords), weight, fraction);
    };
  };


  /// @brief FillableStorage, introduces FillAdapterT on top of BinnedStorage base class
  ///
  /// @note The additional abstraction layer is necessary to distinguish between binned objects
  /// that are "live" (i.e. fillable/incrementable, like a Histo1D) and those that are "dead"
  /// (e.g. a binned set of cross-section measurement points).
  template <size_t FillDim, typename BinContentT, typename... AxisT>
  class FillableStorage
        : public BinnedStorage<BinContentT, AxisT...> {

    static_assert((FillDim >= sizeof...(AxisT)),
                  "Fill dimension should be at least as large as the binning dimension.");

  protected:
    /// @brief Convenience alias to be used in constructor
    using BaseT = BinnedStorage<BinContentT, AxisT...>;
    using BinningT = typename BaseT::BinningT;
    using BinT = Bin<sizeof...(AxisT), BinContentT, BinningT>;
    using AdapterWrapperT = defaultAdapter<FillDim, BinT>;
    using FillableT = FillableStorage<FillDim, BinContentT, AxisT...>;
    using FillCoordsT = decltype(std::tuple_cat(std::declval<typename BinningT::EdgeTypesTuple>(),
                                dblPadding(std::make_index_sequence<FillDim-sizeof...(AxisT)>{})));

  public:

    /// @brief Type of the fill coordinates
    using FillType = FillCoordsT;

    /// @brief Fill dimension
    using FillDimension = std::integral_constant<size_t, FillDim>;

    /// @brief Adapter type (type of lambda used to access stored object).
    ///
    /// BinT is a stored object type;
    /// FillCoordsT is a coordinates init list (or tuple) type;
    /// Doubles are for weight and fraction correspondingly.
    using FillAdapterT = std::function<void(BinT&, FillCoordsT&&, double, double)>;

    /// @name Constructors
    // @{

    /// @brief Nullary constructor for unique pointers etc.
    FillableStorage(FillAdapterT adapter = AdapterWrapperT()._adapter)
        : BaseT(), _fillAdapter(adapter), _nancount(0), _nansumw(0.), _nansumw2(0.) { }

    /// @brief Constructs FillableStorage from Binning.
    FillableStorage(const BinningT& binning, FillAdapterT adapter = AdapterWrapperT()._adapter)
        : BaseT(binning), _fillAdapter(adapter), _nancount(0), _nansumw(0.), _nansumw2(0.) { }

    /// @brief Constructs FillableStorage from Binning. Rvalue.
    FillableStorage(BinningT&& binning, FillAdapterT adapter = AdapterWrapperT()._adapter)
        : BaseT(std::move(binning)), _fillAdapter(adapter), _nancount(0), _nansumw(0.), _nansumw2(0.) { }

    /// @brief Constructs binning from an adapter and vectors of axes' edges
    FillableStorage(const std::vector<AxisT>&... edges, FillAdapterT adapter = AdapterWrapperT()._adapter)
        : BaseT(edges...), _fillAdapter(adapter), _nancount(0), _nansumw(0.), _nansumw2(0.) { }

    /// @brief Constructs binning from an adapter and Rvalue vectors of axes' edges
    FillableStorage(std::vector<AxisT>&&... edges, FillAdapterT adapter = AdapterWrapperT()._adapter)
        : BaseT(std::move(edges)...), _fillAdapter(adapter), _nancount(0), _nansumw(0.), _nansumw2(0.) { }

    /// @brief Constructs binning from an adapter and a sequence of axes
    FillableStorage(const Axis<AxisT>&... axes, FillAdapterT adapter = AdapterWrapperT()._adapter)
        : BaseT(axes...), _fillAdapter(adapter), _nancount(0), _nansumw(0.), _nansumw2(0.) { }

    /// @brief Constructs binning from an adapter and a sequence of Rvalue axes
    FillableStorage(Axis<AxisT>&&... axes, FillAdapterT adapter = AdapterWrapperT()._adapter)
        : BaseT(std::move(axes)...), _fillAdapter(adapter), _nancount(0), _nansumw(0.), _nansumw2(0.) { }

    /// @brief Copy constructor.
    FillableStorage(const FillableStorage& other)
        : BaseT(other), _fillAdapter(other._fillAdapter),
          _nancount(other._nancount), _nansumw(other._nansumw), _nansumw2(other._nansumw2) { }

    /// @brief Move constructor.
    FillableStorage(FillableStorage&& other)
        : BaseT(std::move(other)), _fillAdapter(std::move(other._fillAdapter)),
          _nancount(std::move(other._nancount)), _nansumw(std::move(other._nansumw)),
          _nansumw2(std::move(other._nansumw2)) { }

    // @}

    /// @name Methods
    // @{

    /// @brief Triggers fill adapter on the bin corresponding to coords
    ///
    /// @note Accepts coordinates only as rvalue tuple. The tuple members
    /// are then moved (bringing tuple member to unspecified state) later in adapters.
    template <size_t... Is>
    int fill(FillCoordsT&& coords, std::index_sequence<Is...>,
             const double weight = 1.0, const double fraction = 1.0) noexcept {

      // make sure the user isn't trying to fill with NaN ...
      // include all fill coordinates here
      if (containsNan(coords)) {
        _nancount += 1;
        _nansumw += weight*fraction;
        _nansumw2 += sqr(weight*fraction);
        return -1;
      }
      // select binned coordinates (possibly a subset of fill coordinates)
      auto binCoords = std::tuple<AxisT...>(std::get<Is>(coords)...);
      const size_t binIdx = FillableT::_binning.globalIndexAt(binCoords);
      _fillAdapter(BaseT::bin(binIdx), std::move(coords), weight, fraction);
      return int(binIdx);
    }

    /// @brief Triggers fill adapter on the bin corresponding to coords
    int fill(FillCoordsT&& coords, const double weight = 1.0, const double fraction = 1.0) noexcept {
      return fill(std::move(coords), std::make_index_sequence<sizeof...(AxisT)>{}, weight, fraction);
    }


    /// @name Utilities
    // @{

    /// @brief Returns the dimension of the filling tuple
    size_t fillDim() const { return FillDim; }

    size_t nanCount() const { return _nancount; }

    double nanSumW() const { return _nansumw; }

    double nanSumW2() const { return _nansumw2; }

    void setNanLog(size_t count, double sumw, double sumw2) {
      _nancount = count;
      _nansumw  = sumw;
      _nansumw2 = sumw2;
    }

    /// @brief Reset the Fillable.
    ///
    /// Keep the binning but set all bin contents and related quantities to zero
    void reset() noexcept {
      _nancount = 0;
      _nansumw = _nansumw2 = 0.;
      BaseT::clearBins();
    }

    // @}

    /// @name Operators
    // @{

    /// @brief Copy assignment
    FillableStorage& operator = (const FillableStorage& other) noexcept {
      if (this != &other) {
        _fillAdapter = other._fillAdapter;
        _nancount    = other._nancount;
        _nansumw     = other._nansumw;
        _nansumw2    = other._nansumw2;
        BaseT::operator=(other);
      }
      return *this;
    }

    /// @brief Move assignment
    FillableStorage& operator = (FillableStorage&& other) noexcept {
      if (this != &other) {
        _fillAdapter = std::move(other._fillAdapter);
        _nancount    = std::move(other._nancount);
        _nansumw     = std::move(other._nansumw);
        _nansumw2    = std::move(other._nansumw2);
        BaseT::operator=(std::move(other));
      }
      return *this;
    }

    /// @brief Add another BinnedStorage to this one.
    FillableStorage& operator += (const FillableStorage& other) {
      if (*this != other)
        throw std::logic_error("YODA::BinnedStorage<" + std::to_string(sizeof...(AxisT)) +\
                               ">: Cannot add BinnedStorages with different binnings.");
      size_t i = 0;
      for (auto& bin : FillableT::bins(true)) {
        bin += other.bin(i++);
      }

      return *this;
    }

    /// @brief Subtract another BinnedStorage from this one.
    FillableStorage& operator -= (const FillableStorage& other) {
      if (*this != other)
        throw std::logic_error("YODA::FillableStorage<" + std::to_string(sizeof...(AxisT)) +\
                               ">: Cannot substract FillableStorages with different binnings.");

      size_t i = 0;
      for (auto& bin : FillableT::bins(true)) {
        bin -= other.bin(i++);
      }

      return *this;
    }

    // @}

    private:

      /// @brief Adapter used to access stored objects
      FillAdapterT _fillAdapter;

      size_t _nancount;

      double _nansumw, _nansumw2;

  };

} // namespace YODA

#endif
