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

#include "YODA/AnalysisObject.h"
#include "YODA/Counter.h"
#include "YODA/Estimate0D.h"
#include "YODA/Scatter.h"
#include "YODA/Histo.h"
#include "YODA/Profile.h"
#include "YODA/BinnedEstimate.h"

#ifdef HAVE_HDF5
#include <YODA/Config/BuildConfig.h>
#include "YODA/Utils/H5Utils.h"
#ifdef WITH_HIGHFIVE
#include <YODA/highfive/H5File.hpp>
#else
#include <highfive/H5File.hpp>
#define YODA_H5 HighFive
#endif
#endif

#include <iostream>
#include <locale>
#include <cstring>
#include <regex>
#include <cmath>

#ifdef WITH_OSX
#include <xlocale.h>
#endif

using std::map;
using std::string;
using std::vector;

namespace YODA {

  /// Anonymous namespace to limit visibility
  namespace {

    static const std::regex regex_string_pat("([\"\'])(?:(?=(\\\\?))\\2.)*?\\1");

  }


  class AOReaderBase {

    /// Fast ASCII tokenizer, extended from FastIStringStream by Gavin Salam.
    class aistringstream {
    public:
      // Constructor from char*
      aistringstream(const char* line=0) {
        reset(line);
        _set_locale();
      }
      // Constructor from std::string
      aistringstream(const string& line) {
        reset(line);
        _set_locale();
      }
      ~aistringstream() {
        _reset_locale();
      }

      bool peek(const std::string& s) const {
        return s == std::string(_next, s.size());
      }

      // Re-init to new line as char*
      void reset(const char* line=0) {
        _next = const_cast<char*>(line);
        _new_next = _next;
        _error = false;
      }
      // Re-init to new line as std::string
      void reset(const string& line) { reset(line.c_str()); }

      // Tokenizing stream operator (forwards to specialisations)
      template<class T>
      aistringstream& operator >> (T& value) {
        _get(value);
        if (_new_next == _next) _error = true; // handy error condition behaviour!
        _next = _new_next;
        return *this;
      }

      // Allow use of operator>> in a while loop
      operator bool() const { return !_error; }


    private:

      // Changes the thread-local locale to interpret numbers in the "C" locale
      void _set_locale() {
        _locale_set = newlocale(LC_NUMERIC_MASK, "C", NULL);
        _locale_prev = uselocale(_locale_set);
        if (!_locale_prev) {
          throw ReadError(std::string("Error setting locale: ") + strerror(errno));
        }
      }
      void _reset_locale() {
        if (!uselocale(_locale_prev)) {
          throw ReadError(std::string("Error setting locale: ") + strerror(errno));
        }
        freelocale(_locale_set);
      }

      void _get(double& x) { x = std::strtod(_next, &_new_next); }
      void _get(float& x) { x = std::strtof(_next, &_new_next); }
      void _get(int& i) { i = std::strtol(_next, &_new_next, 10); } // force base 10!
      void _get(long& i) { i = std::strtol(_next, &_new_next, 10); } // force base 10!
      void _get(unsigned int& i) { i = std::strtoul(_next, &_new_next, 10); } // force base 10!
      void _get(long unsigned int& i) { i = std::strtoul(_next, &_new_next, 10); } // force base 10!
      void _get(string& x) {
        while (std::isspace(*_next) && _next[0] != '\0') _next += 1;
        _new_next = _next;
        while (!std::isspace(*_new_next) && _new_next[0] != '\0') _new_next += 1;
        x = string(_next, _new_next-_next);
      }

      locale_t _locale_set, _locale_prev;
      char *_next, *_new_next;
      bool _error;
    };


  public:

    /// Default constructor
    AOReaderBase() { }

    /// Default destructor
    virtual ~AOReaderBase() { }

    virtual void parse(const string& line) = 0;

    virtual AnalysisObject* assemble(const string& path = "") = 0;

    #ifdef HAVE_HDF5

    virtual AnalysisObject* mkFromH5(H5FileManager&) = 0;

    virtual void skip(H5FileManager& h5file) {
      h5file.skipCommon();
    }

    #endif

    template<typename T>
    void extractVector(const std::string& line, std::vector<T>& vec) {
      if constexpr (std::is_same<T, std::string>::value) {
        string::const_iterator initpos( line.cbegin() );
        const string::const_iterator finpos( line.cend() );
        std::smatch m;
        while ( std::regex_search(initpos, finpos, m, regex_string_pat) ) {
          string label;
          std::stringstream ss(m[0].str());
          ss >> std::quoted(label); // removes outer quotes and de-escapes inner quotes
          vec.push_back(label);
          initpos = m.suffix().first;
        }
      }
      else {
         std::string content = line.substr(line.find(": [")+3);
         content.pop_back(); // remove the "]" at the end
         for (const std::string& item : Utils::split(content, ",")) {
           aiss.reset(item);
           T tmp;
           aiss >> tmp;
           vec.push_back(std::move(tmp));
         }
      }
    }

  protected:

    aistringstream aiss;

  };




  template<class T>
  class AOReader;

  template<>
  class AOReader<Counter> : public AOReaderBase {

    Dbn0D dbn;

    public:

    void parse(const string& line) {
      aiss.reset(line);
      double sumw(0), sumw2(0), n(0);
      aiss >> sumw >> sumw2 >> n;
      dbn = Dbn0D(n, sumw, sumw2);
    }

    AnalysisObject* assemble(const string& path = "") {
      auto* ao = new Counter(path);
      ao->setDbn(dbn);
      dbn = Dbn0D();
      return ao;
    }

    #ifdef HAVE_HDF5

    AnalysisObject* mkFromH5(H5FileManager& h5file) {
      vector<string> annos = h5file.loadAnnotations();
      auto* ao = new Counter(h5file.path(), annos.back());
      annos.pop_back();
      ao->deserializeMeta(annos);
      ao->deserializeContent(h5file.loadContent());
      return ao;
    };

    #endif
  };


  template<>
  class AOReader<Estimate0D> : public AOReaderBase {

    Estimate0D est;
    vector<string> sources;

    void readErrors(std::map<string,std::pair<double,double>>& errors) {
      string eDn, eUp;
      for (size_t i = 0; i < sources.size(); ++i) {
        aiss >> eDn >> eUp;
        if (eDn != "---" && eUp != "---") {
          errors[sources[i]] = { Utils::toDbl(eDn), Utils::toDbl(eUp) };
        }
      }
    }

    public:

    void parse(const string& line) {
      if (!line.rfind("ErrorLabels: ", 0)) { // parse error labels
        extractVector<std::string>(line, sources);
        return;
      }
      // parse content
      aiss.reset(line);
      double val(0);
      aiss >> val;
      std::map<string,std::pair<double,double>> errors;
      readErrors(errors);
      est = Estimate0D(val, errors);
    }

    AnalysisObject* assemble(const string& path = "") {

      auto* ao = new Estimate0D(est, path);
      est = Estimate0D();
      sources.clear();
      return ao;
    }

    #ifdef HAVE_HDF5
    AnalysisObject* mkFromH5(H5FileManager& h5file) {
      vector<string> annos = h5file.loadAnnotations();
      auto* ao = new Estimate0D(h5file.path(), annos.back());
      annos.pop_back();
      ao->deserializeMeta(annos);
      ao->deserializeContent(h5file.loadContent());
      ao->deserializeSources(h5file.loadSources());
      return ao;
    }

    void skip(H5FileManager& h5file) {
      h5file.skipCommon();
      h5file.skipSources();
    }

    #endif
  };


  template <size_t N>
  class AOReader<ScatterND<N>> : public AOReaderBase {

    vector<PointND<N>> points;

    template<size_t I>
    void readCoords(vector<double>& vals, vector<double>& errm, vector<double>& errp) {
      if constexpr(I < N) {
        double v(0), em(0), ep(0);
        aiss >> v >> em >> ep;
        vals[I] = v;
        errm[I] = fabs(em);
        errp[I] = fabs(ep);
        readCoords<I+1>(vals, errm, errp);
      }
    }

    public:

    void parse(const string& line) {
      aiss.reset(line);
      vector<double> vals(N), errm(N), errp(N);
      readCoords<0>(vals, errm, errp);
      points.push_back(PointND<N>(vals, errm, errp));
    }

    AnalysisObject* assemble(const string& path = "") {
      auto* ao = new ScatterND<N>();
      ao->setPath(path);
      ao->addPoints(points);
      points.clear();
      return ao;
    }

    #ifdef HAVE_HDF5

    AnalysisObject* mkFromH5(H5FileManager& h5file) {
      vector<string> annos = h5file.loadAnnotations();
      auto* ao = new ScatterND<N>(h5file.path(), annos.back());
      annos.pop_back();
      ao->deserializeMeta(annos);
      ao->deserializeContent(h5file.loadContent());
      return ao;
    }

    #endif
  };


  template <size_t DbnN, typename... AxisT>
  class AOReader<BinnedDbn<DbnN, AxisT...>> : public AOReaderBase {

    using BaseT = BinnedDbn<DbnN, AxisT...>;

    template <size_t I>
    using is_CAxis = typename std::is_floating_point<typename std::tuple_element_t<I, std::tuple<AxisT...>>>;

    std::tuple<vector<AxisT> ...> edges;
    Dbn<DbnN> yoda1Overflow;
    vector<Dbn<DbnN>> dbns;
    vector<size_t> maskedBins;
    std::array<double,DbnN*(DbnN-1)/2> crossTerms;
    bool isYODA1 = false;
    size_t axisCheck = 0;


    template<size_t I>
    void readEdges() { // YODA1 version for backwards compatibility
      if constexpr(I < sizeof...(AxisT)) {
        using EdgeT = std::tuple_element_t<I, std::tuple<AxisT...>>;
        if constexpr (is_CAxis<I>::value) { // continuous case
          EdgeT lo, hi;
          aiss >> lo >> hi;
          auto& curr_edges = std::get<I>(edges);
          if (!std::isinf(lo)) {
            if (curr_edges.empty())  curr_edges.push_back(lo);
          }
          if (!std::isinf(hi)) {
            if (curr_edges.size() && curr_edges[ curr_edges.size() - 1 ] != hi) {
              curr_edges.push_back(hi);
            }
          }
        }
        else { // discrete case
          throw BinningError("Discrete axes are not supported in this YODA1-style legacy format.");
        }
        readEdges<I+1>();
      }
    }

    template<size_t I>
    void readEdges(const std::string& line) { // YODA2 version
      if constexpr(I < sizeof...(AxisT)) {
        if (I == axisCheck) {
          using EdgeT = std::tuple_element_t<I, std::tuple<AxisT...>>;
          auto& curr_edges = std::get<I>(edges);
          extractVector<EdgeT>(line, curr_edges);
        }
        readEdges<I+1>(line);
      }
    }

    #ifdef HAVE_HDF5

    template <size_t I>
    void loadEdges(H5FileManager& h5file) {
      if constexpr(I < sizeof...(AxisT)) {
        using EdgeT = std::tuple_element_t<I, std::tuple<AxisT...>>;
        auto& curr_edges = std::get<I>(edges);
        curr_edges = h5file.loadEdges<EdgeT>();
        loadEdges<I+1>(h5file);
      }
    }

    template <size_t I>
    void skipEdges(H5FileManager& h5file) {
      if constexpr(I < sizeof...(AxisT)) {
        using EdgeT = std::tuple_element_t<I, std::tuple<AxisT...>>;
        h5file.skipEdges<EdgeT>();
        skipEdges<I+1>(h5file);
      }
    }

    #endif

    template<size_t I>
    void readDbn(std::array<double,DbnN+1>& sumW, std::array<double,DbnN+1>& sumW2) {
      if constexpr(I <= DbnN) {
        double w(0), w2(0);
        aiss >> w >> w2;
        sumW[I] = w;
        sumW2[I] = w2;
        readDbn<I+1>(sumW, sumW2);
      }
    }

    template <class tupleT, size_t... Is>
    BaseT* make_from_tuple(tupleT&& tuple, std::index_sequence<Is...> ) {
      BaseT* rtn = new BaseT{std::get<Is>(std::forward<tupleT>(tuple))...};
      rtn->maskBins(maskedBins);
      return rtn;
    }

    template <class tupleT>
    BaseT* make_from_tuple(tupleT&& tuple) {
      return make_from_tuple(std::forward<tupleT>(tuple),
                             std::make_index_sequence<sizeof...(AxisT)+1>{});
    }

    template<size_t I>
    void clearEdges() {
      if constexpr(I < sizeof...(AxisT)) {
        std::get<I>(edges).clear();
        clearEdges<I+1>();
      }
    }

    public:

    void parse(const string& line) {
      if (line.find("Total") != string::npos) {
        isYODA1 = true;
        return; // YODA1 backwards compatibility
      }
      if (!line.rfind("Edges(A", 0)) { // parse binning
        readEdges<0>(line);
        ++axisCheck;
        return;
      }
      if (!line.rfind("MaskedBins: ", 0)) { // parse indices of masked bins
        extractVector<size_t>(line, maskedBins);
        return;
      }
      aiss.reset(line);
      if (line.find("Underflow") != string::npos || line.find("Overflow") != string::npos) {
        // This must be the YODA1-style format ...
        if constexpr (sizeof...(AxisT) == 1) {
          string tmp1, tmp2;
          aiss >> tmp1 >> tmp2; // not needed
        }
      }
      else if (isYODA1)  readEdges<0>();
      std::array<double,DbnN+1> sumW, sumW2;
      readDbn<0>(sumW, sumW2);
      for (size_t i = 0; i < crossTerms.size(); ++i) {
        double tmp(0.);
        aiss >> tmp;
        crossTerms.at(i) = tmp;
      }
      double numEntries(0);
      aiss >> numEntries;
      if (line.find("Overflow") != string::npos) {
        if constexpr (sizeof...(AxisT) == 1) {
          if constexpr (DbnN < 2)
            yoda1Overflow = Dbn<DbnN>(numEntries, sumW, sumW2);
          else
            yoda1Overflow = Dbn<DbnN>(numEntries, sumW, sumW2, crossTerms);
        }
      }
      else {
        if constexpr (DbnN < 2) {
          dbns.emplace_back(numEntries, sumW, sumW2);
        }
        else {
          dbns.emplace_back(numEntries, sumW, sumW2, crossTerms);
        }
      }
    }

    AnalysisObject* assemble(const string& path = "") {

      auto args = std::tuple_cat(edges, std::make_tuple(path));
      BaseT* ao = make_from_tuple(std::move(args));

      size_t global_index = 0;
      if constexpr (sizeof...(AxisT) == 2) {
        if (isYODA1) { // 2D objects had no under-/overflows in Y1
          for (size_t ix = 1; ix < ao->numBinsAt(0)+1; ++ix) { //< visible bins only
            for (size_t iy = 1; iy < ao->numBinsAt(1)+1; ++iy) { //< visible bins only
              ao->bin(ix,iy).set(std::move(dbns[global_index++]));
            }
          }
        }
      }
      if ( !(isYODA1 && sizeof...(AxisT) == 2) ) { //< still works for Y1-style 1D
        for (auto&& d : dbns) {
          ao->bin(global_index++).set(std::move(d));
        }
      }

      if constexpr (sizeof...(AxisT) == 1) { // YODA1-style overflows
        if (isYODA1)  ao->bin(global_index).set(yoda1Overflow);
        yoda1Overflow = Dbn<DbnN>();
      }

      crossTerms.fill(0);
      maskedBins.clear();
      isYODA1 = false;
      clearEdges<0>();
      dbns.clear();
      axisCheck = 0;
      return ao;
    }

    #ifdef HAVE_HDF5

    AnalysisObject* mkFromH5(H5FileManager& h5file) {
      loadEdges<0>(h5file);
      maskedBins = h5file.loadMasks();
      vector<string> annos = h5file.loadAnnotations();
      auto args = std::tuple_cat(edges, std::make_tuple(h5file.path(), annos.back()));
      BaseT* ao = make_from_tuple(std::move(args));
      annos.pop_back();
      ao->deserializeMeta(annos);
      ao->deserializeContent(h5file.loadContent());
      maskedBins.clear();
      clearEdges<0>();
      return ao;
    }

    void skip(H5FileManager& h5file) {
      skipEdges<0>(h5file);
      h5file.skipMasks();
      h5file.skipCommon();
    }

    #endif

  };


  template <typename... AxisT>
  class AOReader<BinnedEstimate<AxisT...>> : public AOReaderBase {

    using BaseT = BinnedEstimate<AxisT...>;

    std::tuple<vector<AxisT> ...> edges;
    vector<Estimate> estimates;
    vector<size_t> maskedBins;
    vector<string> sources;
    size_t axisCheck = 0;


    template<size_t I>
    void readEdges(const std::string& line) {
      if constexpr(I < sizeof...(AxisT)) {
        if (I == axisCheck) {
          using EdgeT = std::tuple_element_t<I, std::tuple<AxisT...>>;
          auto& curr_edges = std::get<I>(edges);
          extractVector<EdgeT>(line, curr_edges);
        }
        readEdges<I+1>(line);
      }
    }

    void readErrors(std::map<string,std::pair<double,double>>& errors) {
      string eDn, eUp;
      for (const std::string& src : sources) {
        aiss >> eDn >> eUp;
        if (eDn != "---" && eUp != "---") {
          errors[src] = { Utils::toDbl(eDn), Utils::toDbl(eUp) };
        }
      }
    }

    #ifdef HAVE_HDF5

    template <size_t I>
    void loadEdges(H5FileManager& h5file) {
      if constexpr(I < sizeof...(AxisT)) {
        using EdgeT = std::tuple_element_t<I, std::tuple<AxisT...>>;
        auto& curr_edges = std::get<I>(edges);
        curr_edges = h5file.loadEdges<EdgeT>();
        loadEdges<I+1>(h5file);
      }
    }

    template <size_t I>
    void skipEdges(H5FileManager& h5file) {
      if constexpr(I < sizeof...(AxisT)) {
        using EdgeT = std::tuple_element_t<I, std::tuple<AxisT...>>;
        h5file.skipEdges<EdgeT>();
        skipEdges<I+1>(h5file);
      }
    }

    #endif

    template <class tupleT, size_t... Is>
    BaseT* make_from_tuple(tupleT&& tuple, std::index_sequence<Is...> ) {
      BaseT* rtn = new BaseT{std::get<Is>(std::forward<tupleT>(tuple))...};
      rtn->maskBins(maskedBins);
      return rtn;
    }

    template <class tupleT>
    BaseT* make_from_tuple(tupleT&& tuple) {
      return make_from_tuple(std::forward<tupleT>(tuple),
                             std::make_index_sequence<sizeof...(AxisT)+1>{});
    }

    template<size_t I>
    void clearEdges() {
      if constexpr(I < sizeof...(AxisT)) {
        std::get<I>(edges).clear();
        clearEdges<I+1>();
      }
    }

    public:

    void parse(const string& line) {
      if (!line.rfind("Edges(A", 0)) { // parse binning
        readEdges<0>(line);
        ++axisCheck;
        return;
      }
      if (!line.rfind("MaskedBins: ", 0)) { // parse indices of masked bins
        extractVector<size_t>(line, maskedBins);
        return;
      }
      if (!line.rfind("ErrorLabels: ", 0)) { // parse error labels
        extractVector<std::string>(line, sources);
        return;
      }
      // parse bin content
      aiss.reset(line);
      double val(0);
      aiss >> val;
      std::map<string,std::pair<double,double>> errors;
      readErrors(errors);
      estimates.emplace_back(val, errors);
    }

    AnalysisObject* assemble(const string& path = "") {

      auto args = std::tuple_cat(edges, std::make_tuple(path));
      BaseT* ao = make_from_tuple(std::move(args));

      size_t global_index = 0;
      for (auto&& e : estimates) {
        ao->bin(global_index++) = std::move(e);
      }

      clearEdges<0>();
      sources.clear();
      estimates.clear();
      maskedBins.clear();
      axisCheck = 0;
      return ao;
    }

    #ifdef HAVE_HDF5

    AnalysisObject* mkFromH5(H5FileManager& h5file) {
      loadEdges<0>(h5file);
      maskedBins = h5file.loadMasks();

      vector<string> annos = h5file.loadAnnotations();
      auto args = std::tuple_cat(edges, std::make_tuple(h5file.path(), annos.back()));
      BaseT* ao = make_from_tuple(std::move(args));
      annos.pop_back();
      ao->deserializeMeta(annos);
      ao->deserializeContent(h5file.loadContent());

      for (auto& b : ao->bins(true, true)) {
        b.deserializeSources(h5file.loadSources());
      }

      maskedBins.clear();
      clearEdges<0>();
      return ao;
    }

    void skip(H5FileManager& h5file) {
      skipEdges<0>(h5file);
      h5file.skipMasks();
      h5file.skipSources();
      h5file.skipCommon();
    }

    #endif

  };


}

#endif
