// ****************************************************************************
//
//          Aevol - An in silico experimental evolution platform
//
// ****************************************************************************
//
// Copyright: See the AUTHORS file provided with the package or <www.aevol.fr>
// Web: http://www.aevol.fr/
// E-mail: See <http://www.aevol.fr/contact/>
// Original Authors : Guillaume Beslon, Carole Knibbe, David Parsons
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
//
// ****************************************************************************

#ifndef AEVOL_MEANSTATSHANDLER_H_
#define AEVOL_MEANSTATSHANDLER_H_

#include "AbstractStatsHandler.h"

#include <filesystem>
#include <fstream>
#include <numeric>
#include <ranges>
#include <utility>

#include "AeTime.h"
#include "MultiIndividualDataCollector.h"
#include "StatsDataConcepts.h"

namespace aevol {


template <GlobalStatsData Head, MultiIndividualStatsData... Tail>
class MeanStatsHandler : public AbstractStatsHandler {
 protected:
  MeanStatsHandler()                                   = default;
  MeanStatsHandler(const MeanStatsHandler&)            = delete;
  MeanStatsHandler(MeanStatsHandler&&)                 = delete;
  MeanStatsHandler& operator=(const MeanStatsHandler&) = delete;
  MeanStatsHandler& operator=(MeanStatsHandler&&)      = delete;

 public:
  static auto make_empty(std::filesystem::path path) -> std::unique_ptr<MeanStatsHandler>;
  static auto make_resume(std::filesystem::path path, time_type generation) -> std::unique_ptr<MeanStatsHandler>;

  virtual ~MeanStatsHandler() = default;

  void write_header() override;
  void write_data(time_type generation,
                  size_t pop_size,
                  indiv_id_type indiv_idx,
                  std::ranges::input_range auto&& indiv_range);
  void write_data(time_type generation,
                  size_t pop_size,
                  indiv_id_type indiv_idx,
                  std::input_iterator auto first,
                  std::sentinel_for<decltype(first)> auto end);
};

template <GlobalStatsData Head, MultiIndividualStatsData... Tail>
auto MeanStatsHandler<Head, Tail...>::make_empty(std::filesystem::path path)
    -> std::unique_ptr<MeanStatsHandler> {
  auto new_stats_handler =
      std::unique_ptr<MeanStatsHandler>(new MeanStatsHandler<Head, Tail...>());
  // Protected ctor -> can't use make_unique

  new_stats_handler->prepare_empty(path);

  return new_stats_handler;
}

template <GlobalStatsData Head, MultiIndividualStatsData... Tail>
auto MeanStatsHandler<Head, Tail...>::make_resume(std::filesystem::path path, time_type generation)
    -> std::unique_ptr<MeanStatsHandler> {
  auto new_stats_handler = std::unique_ptr<MeanStatsHandler>(new MeanStatsHandler());
  // Protected ctor -> can't use make_unique

  new_stats_handler->prepare_for_resume(path, generation);

  return new_stats_handler;
}

template <GlobalStatsData Head, MultiIndividualStatsData... Tail>
void MeanStatsHandler<Head, Tail...>::write_header() {
  statfile_ << Head::header();
  (..., (statfile_ << "," << Tail::header())); // Left fold; write comma separated headers for Tail...
  statfile_ << std::endl;
}

template <GlobalStatsData Head, MultiIndividualStatsData... Tail>
void MeanStatsHandler<Head, Tail...>::write_data(time_type generation,
                                                 size_t pop_size,
                                                 indiv_id_type indiv_idx,
                                                 std::ranges::input_range auto&& indiv_range) {
  statfile_ << Head::collect_data(generation, pop_size, indiv_idx);
  std::vector<std::vector<std::vector<double>>> data;
  (..., data.push_back(collect_data<Tail>(indiv_range)));  // Left fold; collect data for each element in Tail...
  for (const auto& chunk : data) { // Each chunk of data corresponds to the data collected for an element in Tail
    for (const auto& metric : chunk) { // metric contains the value of one given metric for each indiv in the range
      auto nb_nan = 0;
      statfile_ << ',' << std::accumulate(metric.cbegin(), metric.cend(), 0., [&nb_nan](const double &a, const double &b) {
                            if (std::isnan(b)) {
                              ++nb_nan;
                              return std::move(a);
                            }
                            return std::move(a) + b;
                          }) / (metric.size() - nb_nan);
    }
  }
  statfile_ << std::endl;
}

template <GlobalStatsData Head, MultiIndividualStatsData... Tail>
void MeanStatsHandler<Head, Tail...>::write_data(time_type generation,
                                                 size_t pop_size,
                                                 indiv_id_type indiv_idx,
                                                 std::input_iterator auto first,
                                                 std::sentinel_for<decltype(first)> auto end) {
  write_data(generation, pop_size, indiv_idx, std::ranges::subrange(first, end));
}

}  // namespace aevol

#endif  // AEVOL_MEANSTATSHANDLER_H_
