#include "CondTools/BTau/interface/BTagCalibrationReader.h"
#include "FWCore/Utilities/interface/Exception.h"

class BTagCalibrationReader::BTagCalibrationReaderImpl {
  friend class BTagCalibrationReader;

public:
  struct TmpEntry {
    float etaMin;
    float etaMax;
    float ptMin;
    float ptMax;
    float discrMin;
    float discrMax;
    TF1 func;
  };

private:
  BTagCalibrationReaderImpl(BTagEntry::OperatingPoint op,
                            const std::string &sysType,
                            const std::vector<std::string> &otherSysTypes = {});

  void load(const BTagCalibration &c, BTagEntry::JetFlavor jf, std::string measurementType);

  double eval(BTagEntry::JetFlavor jf, float eta, float pt, float discr) const;

  double eval_auto_bounds(const std::string &sys, BTagEntry::JetFlavor jf, float eta, float pt, float discr) const;

  std::pair<float, float> min_max_pt(BTagEntry::JetFlavor jf, float eta, float discr) const;

  std::pair<float, float> min_max_eta(BTagEntry::JetFlavor jf, float discr) const;

  BTagEntry::OperatingPoint op_;
  std::string sysType_;
  std::vector<std::vector<TmpEntry>> tmpData_;  // first index: jetFlavor
  std::vector<bool> useAbsEta_;                 // first index: jetFlavor
  std::map<std::string, std::shared_ptr<BTagCalibrationReaderImpl>> otherSysTypeReaders_;
};

BTagCalibrationReader::BTagCalibrationReaderImpl::BTagCalibrationReaderImpl(
    BTagEntry::OperatingPoint op, const std::string &sysType, const std::vector<std::string> &otherSysTypes)
    : op_(op), sysType_(sysType), tmpData_(3), useAbsEta_(3, true) {
  for (const std::string &ost : otherSysTypes) {
    if (otherSysTypeReaders_.count(ost)) {
      throw cms::Exception("BTagCalibrationReader")
          << "Every otherSysType should only be given once. Duplicate: " << ost;
    }
    otherSysTypeReaders_[ost] = std::unique_ptr<BTagCalibrationReaderImpl>(new BTagCalibrationReaderImpl(op, ost));
  }
}

void BTagCalibrationReader::BTagCalibrationReaderImpl::load(const BTagCalibration &c,
                                                            BTagEntry::JetFlavor jf,
                                                            std::string measurementType) {
  if (!tmpData_[jf].empty()) {
    throw cms::Exception("BTagCalibrationReader") << "Data for this jet-flavor is already loaded: " << jf;
  }

  BTagEntry::Parameters params(op_, measurementType, sysType_);
  const std::vector<BTagEntry> &entries = c.getEntries(params);

  for (const auto &be : entries) {
    if (be.params.jetFlavor != jf) {
      continue;
    }

    TmpEntry te;
    te.etaMin = be.params.etaMin;
    te.etaMax = be.params.etaMax;
    te.ptMin = be.params.ptMin;
    te.ptMax = be.params.ptMax;
    te.discrMin = be.params.discrMin;
    te.discrMax = be.params.discrMax;

    if (op_ == BTagEntry::OP_RESHAPING) {
      te.func = TF1("", be.formula.c_str(), be.params.discrMin, be.params.discrMax);
    } else {
      te.func = TF1("", be.formula.c_str(), be.params.ptMin, be.params.ptMax);
    }

    tmpData_[be.params.jetFlavor].push_back(te);
    if (te.etaMin < 0) {
      useAbsEta_[be.params.jetFlavor] = false;
    }
  }

  for (auto &p : otherSysTypeReaders_) {
    p.second->load(c, jf, measurementType);
  }
}

double BTagCalibrationReader::BTagCalibrationReaderImpl::eval(BTagEntry::JetFlavor jf,
                                                              float eta,
                                                              float pt,
                                                              float discr) const {
  bool use_discr = (op_ == BTagEntry::OP_RESHAPING);
  if (useAbsEta_[jf] && eta < 0) {
    eta = -eta;
  }

  // search linearly through eta, pt and discr ranges and eval
  // future: find some clever data structure based on intervals
  const auto &entries = tmpData_.at(jf);
  for (unsigned i = 0; i < entries.size(); ++i) {
    const auto &e = entries.at(i);
    if (e.etaMin <= eta && eta <= e.etaMax  // find eta
        && e.ptMin < pt && pt <= e.ptMax    // check pt
    ) {
      if (use_discr) {                                    // discr. reshaping?
        if (e.discrMin <= discr && discr < e.discrMax) {  // check discr
          return e.func.Eval(discr);
        }
      } else {
        return e.func.Eval(pt);
      }
    }
  }

  return 0.;  // default value
}

double BTagCalibrationReader::BTagCalibrationReaderImpl::eval_auto_bounds(
    const std::string &sys, BTagEntry::JetFlavor jf, float eta, float pt, float discr) const {
  auto sf_bounds_eta = min_max_eta(jf, discr);
  bool eta_is_out_of_bounds = false;

  if (sf_bounds_eta.first < 0)
    sf_bounds_eta.first = -sf_bounds_eta.second;

  if (useAbsEta_[jf] && eta < 0) {
    eta = -eta;
  }

  if (eta <= sf_bounds_eta.first || eta > sf_bounds_eta.second) {
    eta_is_out_of_bounds = true;
  }

  if (eta_is_out_of_bounds) {
    return 1.;
  }

  auto sf_bounds = min_max_pt(jf, eta, discr);
  float pt_for_eval = pt;
  bool is_out_of_bounds = false;

  if (pt <= sf_bounds.first) {
    pt_for_eval = sf_bounds.first + .0001;
    is_out_of_bounds = true;
  } else if (pt > sf_bounds.second) {
    pt_for_eval = sf_bounds.second - .0001;
    is_out_of_bounds = true;
  }

  // get central SF (and maybe return)
  double sf = eval(jf, eta, pt_for_eval, discr);
  if (sys == sysType_) {
    return sf;
  }

  // get sys SF (and maybe return)
  if (!otherSysTypeReaders_.count(sys)) {
    throw cms::Exception("BTagCalibrationReader") << "sysType not available (maybe not loaded?): " << sys;
  }
  double sf_err = otherSysTypeReaders_.at(sys)->eval(jf, eta, pt_for_eval, discr);
  if (!is_out_of_bounds) {
    return sf_err;
  }

  // double uncertainty on out-of-bounds and return
  sf_err = sf + 2 * (sf_err - sf);
  return sf_err;
}

std::pair<float, float> BTagCalibrationReader::BTagCalibrationReaderImpl::min_max_pt(BTagEntry::JetFlavor jf,
                                                                                     float eta,
                                                                                     float discr) const {
  bool use_discr = (op_ == BTagEntry::OP_RESHAPING);
  if (useAbsEta_[jf] && eta < 0) {
    eta = -eta;
  }

  const auto &entries = tmpData_.at(jf);
  float min_pt = -1., max_pt = -1.;
  for (const auto &e : entries) {
    if (e.etaMin <= eta && eta <= e.etaMax  // find eta
    ) {
      if (min_pt < 0.) {  // init
        min_pt = e.ptMin;
        max_pt = e.ptMax;
        continue;
      }

      if (use_discr) {                                    // discr. reshaping?
        if (e.discrMin <= discr && discr < e.discrMax) {  // check discr
          min_pt = min_pt < e.ptMin ? min_pt : e.ptMin;
          max_pt = max_pt > e.ptMax ? max_pt : e.ptMax;
        }
      } else {
        min_pt = min_pt < e.ptMin ? min_pt : e.ptMin;
        max_pt = max_pt > e.ptMax ? max_pt : e.ptMax;
      }
    }
  }

  return std::make_pair(min_pt, max_pt);
}

std::pair<float, float> BTagCalibrationReader::BTagCalibrationReaderImpl::min_max_eta(BTagEntry::JetFlavor jf,
                                                                                      float discr) const {
  bool use_discr = (op_ == BTagEntry::OP_RESHAPING);

  const auto &entries = tmpData_.at(jf);
  float min_eta = 0., max_eta = 0.;
  for (const auto &e : entries) {
    if (use_discr) {                                    // discr. reshaping?
      if (e.discrMin <= discr && discr < e.discrMax) {  // check discr
        min_eta = min_eta < e.etaMin ? min_eta : e.etaMin;
        max_eta = max_eta > e.etaMax ? max_eta : e.etaMax;
      }
    } else {
      min_eta = min_eta < e.etaMin ? min_eta : e.etaMin;
      max_eta = max_eta > e.etaMax ? max_eta : e.etaMax;
    }
  }

  return std::make_pair(min_eta, max_eta);
}

BTagCalibrationReader::BTagCalibrationReader(BTagEntry::OperatingPoint op,
                                             const std::string &sysType,
                                             const std::vector<std::string> &otherSysTypes)
    : pimpl(new BTagCalibrationReaderImpl(op, sysType, otherSysTypes)) {}

void BTagCalibrationReader::load(const BTagCalibration &c,
                                 BTagEntry::JetFlavor jf,
                                 const std::string &measurementType) {
  pimpl->load(c, jf, measurementType);
}

double BTagCalibrationReader::eval(BTagEntry::JetFlavor jf, float eta, float pt, float discr) const {
  return pimpl->eval(jf, eta, pt, discr);
}

double BTagCalibrationReader::eval_auto_bounds(
    const std::string &sys, BTagEntry::JetFlavor jf, float eta, float pt, float discr) const {
  return pimpl->eval_auto_bounds(sys, jf, eta, pt, discr);
}

std::pair<float, float> BTagCalibrationReader::min_max_pt(BTagEntry::JetFlavor jf, float eta, float discr) const {
  return pimpl->min_max_pt(jf, eta, discr);
}