// -*- 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_MathUtils_H
#define YODA_MathUtils_H

/// @todo Add SFINAE math type stuff (see Rivet) and add inrange() and inrange_closed_closed() etc. aliases cf. MCUtils

#include "YODA/Exceptions.h"
#include "YODA/Config/BuildConfig.h"

#include <algorithm>
#include <functional>
#include <numeric>
#include <cassert>
#include <cfloat>
#include <climits>
#include <cmath>
#include <functional>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <ostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace YODA {


  /// Pre-defined numeric type limits
  /// @deprecated Prefer the standard DBL/INT_MAX
  const static double MAXDOUBLE = DBL_MAX; // was std::numeric_limits<double>::max(); -- warns in GCC5
  const static double MAXINT = INT_MAX; // was std::numeric_limits<int>::max(); -- warns in GCC5

  /// A pre-defined value of \f$ \pi \f$.
  static const double PI = M_PI;

  /// A pre-defined value of \f$ 2\pi \f$.
  static const double TWOPI = 2*M_PI;

  /// A pre-defined value of \f$ \pi/2 \f$.
  static const double HALFPI = M_PI_2;

  /// Enum for signs of numbers.
  enum Sign { MINUS = -1, ZERO = 0, PLUS = 1 };


  /// @name Comparison functions for safe (floating point) equality tests
  //@{

  /// @brief Compare a number to zero
  ///
  /// This version for floating point types has a degree of fuzziness expressed
  /// by the absolute @a tolerance parameter, for floating point safety.
  template <typename NUM>
  inline typename std::enable_if_t<std::is_floating_point_v<NUM>, bool>
  isZero(NUM val, double tolerance=1e-5) {
    return fabs(val) < tolerance;
  }

  /// @brief Compare a number to zero
  ///
  /// SFINAE template specialisation for integers, since there is no FP
  /// precision issue.
  template <typename NUM>
  inline typename std::enable_if_t<std::is_integral_v<NUM>, bool>
  isZero(NUM val, double = 1e-5) {
    return val==0;
  }

  /// @brief Check if a number is NaN
  template <typename NUM>
  inline typename std::enable_if_t<std::is_floating_point_v<NUM>, bool>
  isNaN(NUM val) { return std::isnan(val); }

  /// @brief Check if a number is non-NaN
  template <typename NUM>
  inline typename std::enable_if_t<std::is_floating_point_v<NUM>, bool>
  notNaN(NUM val) { return !std::isnan(val); }

  /// @brief Compare two numbers for equality with a degree of fuzziness
  ///
  /// This version for floating point types (if any argument is FP) has a degree
  /// of fuzziness expressed by the fractional @a tolerance parameter, for
  /// floating point safety.
  template <typename N1, typename N2>
  inline typename std::enable_if_t<
    std::is_arithmetic_v<N1> && std::is_arithmetic_v<N2> &&
   (std::is_floating_point_v<N1> || std::is_floating_point_v<N2>), bool>
  fuzzyEquals(N1 a, N2 b, double tolerance=1e-5) {
    const double absavg = (std::abs(a) + std::abs(b))/2.0;
    const double absdiff = std::abs(a - b);
    const bool rtn = (isZero(a) && isZero(b)) || absdiff < tolerance*absavg;
    return rtn;
  }

  /// @brief Compare two numbers for equality with a degree of fuzziness
  ///
  /// Simpler SFINAE template specialisation for integers, since there is no FP
  /// precision issue.
  template <typename N1, typename N2>
  inline typename std::enable_if_t<
    std::is_integral_v<N1> && std::is_integral_v<N2>, bool>
    fuzzyEquals(N1 a, N2 b, double) { //< NB. unused tolerance parameter for ints, still needs a default value!
    return a == b;
  }

  /// @brief Comparator wrapper to use with STL algorithms, e.g. std::equal etc.
  static std::function<bool(const double, const double)> fuzzyEqComp =
    [](const double& lhs, const double& rhs) { return fuzzyEquals(lhs, rhs); };


  /// @brief Compare two numbers for >= with a degree of fuzziness
  ///
  /// The @a tolerance parameter on the equality test is as for @c fuzzyEquals.
  template <typename N1, typename N2>
  inline typename std::enable_if_t<
    std::is_arithmetic_v<N1> && std::is_arithmetic_v<N2>, bool>
  fuzzyGtrEquals(N1 a, N2 b, double tolerance=1e-5) {
    return a > b || fuzzyEquals(a, b, tolerance);
  }

  /// @brief Compare two floating point numbers for <= with a degree of fuzziness
  ///
  /// The @a tolerance parameter on the equality test is as for @c fuzzyEquals.
  template <typename N1, typename N2>
  inline typename std::enable_if_t<
    std::is_arithmetic_v<N1> && std::is_arithmetic_v<N2>, bool>
  fuzzyLessEquals(N1 a, N2 b, double tolerance=1e-5) {
    return a < b || fuzzyEquals(a, b, tolerance);
  }

  /// Returns a number floored at the nth decimal place.
  inline double approx(double a, int n = 5) {
    double roundTo = pow(10.0,n);
    a *= roundTo;
    a = floor(a);
    return a/roundTo;
  }
  //@}


  /// @name Ranges and intervals
  //@{

  /// Represents whether an interval is open (non-inclusive) or closed (inclusive).
  ///
  /// For example, the interval \f$ [0, \pi) \f$ is closed (an inclusive
  /// boundary) at 0, and open (a non-inclusive boundary) at \f$ \pi \f$.
  enum RangeBoundary { OPEN=0, SOFT=0, CLOSED=1, HARD=1 };


  /// @brief Determine if @a value is in the range @a low to @a high, for floating point numbers
  ///
  /// Interval boundary types are defined by @a lowbound and @a highbound.
  /// @todo Optimise to one-line at compile time?
  template<typename NUM>
  inline bool inRange(NUM value, NUM low, NUM high,
                      RangeBoundary lowbound=CLOSED, RangeBoundary highbound=OPEN) {
    if (lowbound == OPEN && highbound == OPEN) {
      return (value > low && value < high);
    } else if (lowbound == OPEN && highbound == CLOSED) {
      return (value > low && value <= high);
    } else if (lowbound == CLOSED && highbound == OPEN) {
      return (value >= low && value < high);
    } else { // if (lowbound == CLOSED && highbound == CLOSED) {
      return (value >= low && value <= high);
    }
  }

  /// Alternative version of inRange for doubles, which accepts a pair for the range arguments.
  template<typename NUM>
  inline bool inRange(NUM value, std::pair<NUM, NUM> lowhigh,
                      RangeBoundary lowbound=CLOSED, RangeBoundary highbound=OPEN) {
    return inRange(value, lowhigh.first, lowhigh.second, lowbound, highbound);
  }


  /// @brief Determine if @a value is in the range @a low to @a high, for integer types
  ///
  /// Interval boundary types are defined by @a lowbound and @a highbound.
  /// @todo Optimise to one-line at compile time?
  inline bool inRange(int value, int low, int high,
                      RangeBoundary lowbound=CLOSED, RangeBoundary highbound=CLOSED) {
    if (lowbound == OPEN && highbound == OPEN) {
      return (value > low && value < high);
    } else if (lowbound == OPEN && highbound == CLOSED) {
      return (value > low && value <= high);
    } else if (lowbound == CLOSED && highbound == OPEN) {
      return (value >= low && value < high);
    } else { // if (lowbound == CLOSED && highbound == CLOSED) {
      return (value >= low && value <= high);
    }
  }

  /// Alternative version of @c inRange for ints, which accepts a pair for the range arguments.
  inline bool inRange(int value, std::pair<int, int> lowhigh,
                      RangeBoundary lowbound=CLOSED, RangeBoundary highbound=OPEN) {
    return inRange(value, lowhigh.first, lowhigh.second, lowbound, highbound);
  }

  //@}


  /// @name Miscellaneous numerical helpers
  //@{

  /// Named number-type squaring operation.
  template <typename NUM>
  inline NUM sqr(NUM a) {
    return a*a;
  }

  /// Named number-type addition in quadrature operation.
  template <typename... Num>
  inline std::enable_if_t<std::conjunction_v<std::is_arithmetic<Num>...>, double>
  add_quad(Num ... vals) {
    return sqrt((0.0 + ... + sqr(vals)));
  }

  /// Find the sign of a number
  inline int sign(double val) {
    if (isZero(val)) return ZERO;
    const int valsign = (val > 0) ? PLUS : MINUS;
    return valsign;
  }

  /// Find the sign of a number
  inline int sign(int val) {
    if (val == 0) return ZERO;
    return (val > 0) ? PLUS : MINUS;
  }

  /// Find the sign of a number
  inline int sign(long val) {
    if (val == 0) return ZERO;
    return (val > 0) ? PLUS : MINUS;
  }

  /// Subtract two numbers with FP fuzziness
  inline double subtract(double a, double b, double tolerance = 1e-5) {
    if (fuzzyEquals(a,b,tolerance))  return 0.;
    return a - b;
  }

  /// Add two numbers with FP fuzziness
  inline double add(double a, double b, double tolerance = 1e-5) {
    return subtract(a,-b,tolerance);
  }

  //@}


  /// @name Binning helper functions
  //@{

  /// @brief Make a list of @a nbins + 1 values uniformly spaced between @a xmin and @a xmax inclusive.
  ///
  /// @note The arg ordering and the meaning of the nbins variable is "histogram-like",
  /// as opposed to the Numpy/Matlab version.
  inline std::vector<double> linspace(size_t nbins, double xmin, double xmax, bool include_end=true) {
    if (xmax < xmin)  throw RangeError("xmax should not be smaller than xmin!");
    if (nbins == 0)   throw RangeError("Requested number of bins is 0!");
    std::vector<double> rtn;
    const double interval = (xmax-xmin)/static_cast<double>(nbins);
    for (size_t i = 0; i < nbins; ++i) {
      rtn.push_back(xmin + i*interval);
    }
    assert(rtn.size() == nbins);
    if (include_end) rtn.push_back(xmax); // exact xmax, not result of n * interval
    return rtn;
  }


  /// @brief Make a list of @a nbins + 1 values uniformly spaced in log(x) between @a xmin and @a xmax inclusive.
  ///
  /// @note The arg ordering and the meaning of the nbins variable is "histogram-like",
  /// as opposed to the Numpy/Matlab version, and the xmin and xmax arguments are expressed
  /// in "normal" space, rather than as the logarithms of the xmin/xmax values as in Numpy/Matlab.
  inline std::vector<double> logspace(size_t nbins, double xmin, double xmax, bool include_end=true) {
    if (xmax < xmin)  throw RangeError("xmax should not be smaller than xmin!");
    if (xmin < 0)     throw RangeError("xmin should not be negative!");
    if (nbins == 0)   throw RangeError("Requested number of bins is 0!");
    const double logxmin = std::log(xmin);
    const double logxmax = std::log(xmax);
    const std::vector<double> logvals = linspace(nbins, logxmin, logxmax);
    assert(logvals.size() == nbins+1);
    std::vector<double> rtn; rtn.reserve(logvals.size());
    rtn.push_back(xmin);
    for (size_t i = 1; i < logvals.size()-1; ++i) {
      rtn.push_back(std::exp(logvals[i]));
    }
    assert(rtn.size() == nbins);
    if (include_end) rtn.push_back(xmax);
    return rtn;
  }


  /// @todo fspace() for uniform sampling from f(x); requires ability to invert fn... how, in general?
  //inline std::vector<double> fspace(size_t nbins, double xmin, double xmax, std::function<double(double)>& fn) {


  /// @brief Make a list of @a nbins + 1 values spaced with *density* ~ f(x) between @a xmin and @a end inclusive.
  ///
  /// The density function @a fn will be evaluated at @a nsample uniformly
  /// distributed points between @a xmin and @a xmax, its integral approximated
  /// via the Trapezium Rule and used to normalize the distribution, and @a
  /// nbins + 1 edges then selected to (approximately) divide into bins each
  /// containing fraction 1/@a nbins of the integral.
  ///
  /// @note The function @a fn does not need to be a normalized pdf, but it should be non-negative.
  /// Any negative values returned by @a fn(x) will be truncated to zero and contribute nothing to
  /// the estimated integral and binning density.
  ///
  /// @note The arg ordering and the meaning of the nbins variable is "histogram-like",
  /// as opposed to the Numpy/Matlab version, and the xmin and xmax arguments are expressed
  /// in "normal" space, rather than in the function-mapped space as with Numpy/Matlab.
  ///
  /// @note The naming of this function differs from the other, Matlab-inspired ones: the bin spacing is
  /// uniform in the CDF of the density function given, rather than in the function itself. For
  /// most use-cases this is more intuitive.
  inline std::vector<double> pdfspace(size_t nbins, double xmin, double xmax, std::function<double(double)>& fn, size_t nsample=10000) {
    const double dx = (xmax-xmin)/(double)nsample;
    const std::vector<double> xs = linspace(nsample, xmin, xmax);
    std::vector<double> ys(0, nsample);
    auto posfn = [&](double x){return std::max(fn(x), 0.0);};
    std::transform(xs.begin(), xs.end(), ys.begin(), posfn);
    std::vector<double> areas; areas.reserve(nsample);
    double areasum = 0;
    for (size_t i = 0; i < ys.size()-1; ++i) {
      const double area = (ys[i] + ys[i+1])*dx/2.0;
      areas[i] = area;
      areasum += area;
    }
    const double df = areasum/(double)nbins;
    std::vector<double> xedges{xmin}; xedges.reserve(nbins+1);
    double fsum = 0;
    for (size_t i = 0; i < nsample-1; ++i) {
      fsum += areas[i];
      if (fsum > df) {
        fsum = 0;
        xedges.push_back(xs[i+1]);
      }
    }
    xedges.push_back(xmax);
    assert(xedges.size() == nbins+1);
    return xedges;
  }


  /// @brief Return the bin index of the given value, @a val, given a vector of bin edges
  ///
  /// NB. The @a binedges vector must be sorted
  template <typename NUM>
  inline int index_between(const NUM& val, const std::vector<NUM>& binedges) {
    if (!inRange(val, binedges.front(), binedges.back())) return -1; //< Out of histo range
    int index = -1;
    for (size_t i = 1; i < binedges.size(); ++i) {
      if (val < binedges[i]) {
        index = i-1;
        break;
      }
    }
    assert(inRange(index, -1, binedges.size()-1));
    return index;
  }

  //@}


  /// @name Statistics functions
  //@{

  /// @brief Calculate the effective number of entries of a sample
  inline double effNumEntries(const double sumW, const double sumW2) {
    if (isZero(sumW2))  return 0;
    return sqr(sumW) / sumW2;
  }

  /// @brief Calculate the effective number of entries of a sample
  inline double effNumEntries(const std::vector<double>& weights) {
    double sumW = 0.0, sumW2 = 0.0;
    for (size_t i = 0; i < weights.size(); ++i) {
      sumW += weights[i];
      sumW2 += sqr(weights[i]);
    }
    return effNumEntries(sumW, sumW2);
  }

  /// @brief Calculate the mean of a sample
  inline double mean(const std::vector<int>& sample) {
    double mean = 0.0;
    for (size_t i=0; i<sample.size(); ++i) {
      mean += sample[i];
    }
    return mean/sample.size();
  }

  /// @brief Calculate the weighted mean of a sample
  inline double mean(const double sumWX, const double sumW) {
    return sumW? sumWX / sumW : std::numeric_limits<double>::quiet_NaN();
  }

  /// @brief Calculate the weighted mean of a sample
  inline double mean(const std::vector<double>& sample,
                     const std::vector<double>& weights) {
    if (sample.size() != weights.size())  throw RangeError("Inputs should have equal length!");
    double sumWX = 0., sumW = 0.;
    for (size_t i = 0; i < sample.size(); ++i) {
      sumW  += weights[i];
      sumWX += weights[i]*sample[i];
    }
    return mean(sumWX, sumW);
  }

  /// @brief Calculate the weighted variance of a sample
  ///
  /// Weighted variance defined as
  /// sig2 = ( sum(wx**2) * sum(w) - sum(wx)**2 ) / ( sum(w)**2 - sum(w**2) )
  /// see http://en.wikipedia.org/wiki/Weighted_mean
  inline double variance(const double sumWX, const double sumW,
                         const double sumWX2, const double sumW2) {
    const double num = subtract(sumWX2*sumW, sqr(sumWX));
    const double den = subtract(sqr(sumW), sumW2);
    /// @todo Isn't this sensitive to the overall scale of the weights?
    /// Shouldn't it check if den is bigger then num by a set number of
    /// orders of magnitude and vice versa?
    // if (fabs(num) < 1e-10 && fabs(den) < 1e-10) {
    //   return std::numeric_limits<double>::quiet_NaN();
    // }
    /// We take the modulus of the weighted variance
    /// since the ratio can be negative with weighted means
    /// @todo Is this the correct approach? There is no information
    /// online other than "weights are non-negative"...
    return den? fabs(num/den): std::numeric_limits<double>::quiet_NaN();
  }

  /// @brief Calculate the weighted variance of a sample
  inline double variance(const std::vector<double>& sample,
                         const std::vector<double>& weights) {
    if (sample.size() != weights.size())  throw RangeError("Inputs should have equal length!");
    if (fuzzyLessEquals(effNumEntries(weights), 1.0)) {
       //throw LowStatsError("Requested variance of a distribution with only one effective entry");
       return std::numeric_limits<double>::quiet_NaN();
    }
    double sumWX = 0., sumW = 0.;
    double sumWX2 = 0., sumW2 = 0.;
    for (size_t i = 0; i < sample.size(); ++i) {
      sumW   += weights[i];
      sumWX  += weights[i]*sample[i];
      sumW2  += sqr(weights[i]);
      sumWX2 += weights[i]*sqr(sample[i]);
    }
    return variance(sumWX, sumW, sumWX2, sumW2);
  }

  /// @brief Calculate the weighted standard deviation of a sample
  inline double stdDev(const double sumWX, const double sumW,
                       const double sumWX2, const double sumW2) {
    return std::sqrt(variance(sumWX, sumW, sumWX2, sumW2));
  }

  /// @brief Calculate the weighted variance of a sample
  inline double stdDev(const std::vector<double>& sample,
                       const std::vector<double>& weights) {
    return std::sqrt(variance(sample, weights));
  }

  /// @brief Calculate the weighted standard error of a sample
  inline double stdErr(const double sumWX, const double sumW,
                       const double sumWX2, const double sumW2) {
    const double effN = effNumEntries(sumW, sumW2);
    if (effN == 0)  return std::numeric_limits<double>::quiet_NaN();
    const double var = variance(sumWX, sumW, sumWX2, sumW2);
    return std::sqrt(var / effN);
  }

  /// @brief Calculate the weighted variance of a sample
  inline double stdErr(const std::vector<double>& sample,
                        const std::vector<double>& weights) {
    if (sample.size() != weights.size())  throw RangeError("Inputs should have equal length!");
    const double effN = effNumEntries(weights);
    if (effN == 0)  return std::numeric_limits<double>::quiet_NaN();
    const double var = variance(sample, weights);
    return std::sqrt(var / effN);
  }

  /// @brief Calculate the weighted RMS of a sample
  inline double RMS(const double sumWX2, const double sumW, const double sumW2) {
    // Weighted RMS defined as
    // rms = sqrt(sum{w x^2} / sum{w})
    const double effN = effNumEntries(sumW, sumW2);
    if (effN == 0)  return std::numeric_limits<double>::quiet_NaN();
    const double meanSq = sumWX2 / sumW;
    return std::sqrt(meanSq);
  }

  /// @brief Calculate the weighted RMS of a sample
  inline double RMS(const std::vector<double>& sample,
                    const std::vector<double>& weights) {
    if (sample.size() != weights.size())  throw RangeError("Inputs should have equal length!");
    double sumWX2 = 0., sumW = 0., sumW2 = 0.;
    for (size_t i = 0; i < sample.size(); ++i) {
      sumW   += weights[i];
      sumW2  += sqr(weights[i]);
      sumWX2 += weights[i]*sqr(sample[i]);
    }
    return RMS(sumWX2, sumW, sumW2);
  }

  /// @brief Calculate the covariance (variance) between two samples
  inline double covariance(const std::vector<int>& sample1, const std::vector<int>& sample2) {
    const double mean1 = mean(sample1);
    const double mean2 = mean(sample2);
    const size_t N = sample1.size();
    double cov = 0.0;
    for (size_t i = 0; i < N; i++) {
      const double cov_i = (sample1[i] - mean1)*(sample2[i] - mean2);
      cov += cov_i;
    }
    if (N > 1) return cov/(N-1);
    else return 0.0;
  }


  /// @brief Calculate the correlation strength between two samples
  inline double correlation(const std::vector<int>& sample1, const std::vector<int>& sample2) {
    const double cov = covariance(sample1, sample2);
    const double var1 = covariance(sample1, sample1);
    const double var2 = covariance(sample2, sample2);
    const double correlation = cov/sqrt(var1*var2);
    const double corr_strength = correlation*sqrt(var2/var1);
    return corr_strength;
  }


  /// @brief Calculate the error-weighted chi2 statistic between two samples
  ///
  /// @note This calculation is rather naive as it neglects the error breakdowns that
  /// may be available for the curves being compared. More sophisticated comparisons
  /// will require a separate analysis based on the Estimate objects utilising the full
  /// covariance information.
  inline double naiveChi2(const std::vector<double>& sample1, const std::vector<double>& sample2,
                     const std::vector<double>& s1errors = std::vector<double>{},
                     const std::vector<double>& s2errors = std::vector<double>{}) {
    if (sample1.size() != sample2.size()) {
      throw RangeError("Inputs should have equal length!");
    }
    if (s1errors.size() && sample1.size() != s1errors.size()) {
      throw RangeError("Inputs should have equal length!");
    }
    if (s2errors.size() && sample2.size() != s2errors.size()) {
      throw RangeError("Inputs should have equal length!");
    }
    const size_t N = sample1.size();
    double chi2 = 0.0;
    for (size_t i = 0; i < N; ++i) {
      double temp = sqr(sample1[i] - sample2[i]);
      if (s1errors.size()) {
        temp /= sqr(s1errors[i]) + sqr(s2errors[i]);
      }
      chi2 += temp;
    }
    return chi2;
  }

  /// @brief Calculate the error-weighted reduced chi2 statistic between two samples
  ///
  /// @note This calculation is rather naive as it neglects the error breakdowns that
  /// may be available for the curves being compared. More sophisticated comparisons
  /// will require a separate analysis based on the Estimate objects utilising the full
  /// covariance information.
  inline double naiveChi2reduced(const std::vector<double>& sample1, const std::vector<double>& sample2,
                        const std::vector<double>& s1errors = std::vector<double>{},
                        const std::vector<double>& s2errors = std::vector<double>{}) {
    if (sample1.empty()) throw RangeError("Inputs should not have 0 length!");
    return naiveChi2(sample1, sample2, s1errors, s2errors)/sample1.size();
  }

  //@}


}

#endif
