// -*- C++ -*-
#include "Rivet/Run.hh"
#include "Rivet/AnalysisHandler.hh"
#include "Rivet/Math/MathUtils.hh"
#include "Rivet/Tools/RivetPaths.hh"
#include "Rivet/Tools/RivetHepMC.hh"
#include "HepMC3/Version.h"
#if HEPMC3_VERSION_CODE < 3003000
#ifdef HAVE_LIBZ
#include "zstr/zstr.hpp"
#endif
#endif
#include <limits>
#include <iostream>

using std::cout;
using std::endl;

namespace Rivet {


  /// Byte/number conversion via union, for HepMC file inspection
  union magic_t {
    uint8_t bytes[2];
    uint16_t number;
  };


  Run::Run(AnalysisHandler& ah)
    : _ah(ah)
  { }


  Run::~Run() { }


  Run& Run::setCrossSection(double xs) {
    _xs = xs;
    return *this;
  }


  Run& Run::setListAnalyses(bool dolist) {
    _listAnalyses = dolist;
    return *this;
  }


  Log& Run::getLog() const {
    return Log::getLog("Rivet.Run");
  }


  // Fill event and check for a bad read state
  bool Run::readEvent() {
    /// @todo Clear rather than new the GenEvent object per-event?
    _evt.reset(new GenEvent());
    if (!HepMCUtils::readEvent(_hepmcReader, _evt)) {
      MSG_DEBUG("Read failed. End of file?");
      _ah.notifyEndOfFile();
      return false;
    }
    // increment counter when event number changes
    if (_evt->event_number() != _evtnumber) {
      _evtnumber = _evt->event_number();
      _evtcount++;
    }
    // Rescale event weights by file-level weight, if scaling is non-trivial
    if (_fileweight != 1.0) {
      for (size_t i = 0; i < (size_t) _evt->weights().size(); ++i) {
        _evt->weights()[i] *= _fileweight;
      }
    }
    return true;
  }



  bool Run::openFile(const std::string& evtfile, double weight) {
    // Set current weight-scaling member
    _fileweight = weight;

    if (evtfile == "-") {
      #if HEPMC3_VERSION_CODE < 3003000
        // Turn off the buffering to make IO faster and make ungetc work on cin
        std::basic_ios<char>::sync_with_stdio(false);
        #ifdef HAVE_LIBZ
        _istr = make_shared<zstr::istream>(std::ref(std::cin));
        #else
        _istr = std::shared_ptr<std::istream>(&std::cin, [](auto*){ /* no deletion */ });
        #endif
        _hepmcReader = RivetHepMC::deduce_reader(*_istr);
      #else
        _hepmcReader = RivetHepMC::deduce_reader(std::cin);
      #endif
    } else {
      // Use standard HepMC3 deduction on file
      _hepmcReader = RivetHepMC::deduce_reader(evtfile);
      #if HEPMC3_VERSION_CODE < 3003000
      // Check if the file is compressed, if the deduction fails
      /// @todo Can we move this into the RivetHepMC.hh header? This is a *lot* of HepMC-specific noise for the Run manager class
      if (!_hepmcReader) {
        MSG_INFO("No success with deduction of file type. Test if the file is compressed");
        std::ifstream file_test(evtfile);
        magic_t my_magic = {{0x1f, 0x8b}};
        magic_t file_magic;
        file_test.read((char *) file_magic.bytes, sizeof(file_magic));
        if (file_magic.number == my_magic.number) {
          MSG_INFO("File is compressed");
          #ifdef HAVE_LIBZ
          _istr = make_shared<zstr::ifstream>(evtfile);
          _hepmcReader = RivetHepMC::deduce_reader(*_istr);
          #else
          MSG_INFO("No zlib support.");
          #endif
        } else {
          // File is not compressed. Open stream and let the code below to handle it
          MSG_INFO("File is not compressed. No success with deduction of file type.");
          _istr = make_shared<std::ifstream>(evtfile);
        }
      }
    }

    // If the deduction has failed, check custom formats
    /// @todo Move this into the RivetHepMC.hh header: this is a *lot* of HepMC-specific noise for the Run manager class
    if (!_hepmcReader) {
      std::vector<std::string> head;
      head.push_back("");
      size_t back=0;
      size_t backnonempty=0;
      while (back < 200 && backnonempty < 100 && _istr) { ///< @todo Document this
        char c = _istr->get();
        back++;
        if (c == '\n') {
          if (head.back().length() != 0)
            head.push_back("");
        } else {
          head.back() += c;
          backnonempty++;
        }
      }
      if (!_istr) MSG_INFO("Info in deduce_reader: input stream is too short or invalid.");
      for (size_t i = 0; i < back; ++i) _istr->unget();
    #endif
    }

    if (_hepmcReader == nullptr) {
      MSG_ERROR("Read error in file '" << evtfile);
      return false;
    }
    return true;
  }


  bool Run::init(const std::string& evtfile, double weight) {
    try {
      if (!openFile(evtfile, weight)) return false;

      // Read first event to define run conditions
      bool ok = readEvent();
      if (!ok) return false;
      if(HepMCUtils::particles(_evt).size() == 0) {
        MSG_ERROR("Empty first event.");
        return false;
      }
      _evtnumber = _evt->event_number(); // set first event number
      _evtcount = 1; // first event used for initialisation

      // Initialise AnalysisHandler with beam information from first event
      _ah.init(*_evt);

      // Set cross-section from command line
      if (notNaN(_xs)) {
        MSG_DEBUG("Setting user cross-section = " << _xs << " pb");
        _ah.setCrossSection(make_pair(_xs, 0.0), true);
      }

      // List the chosen & compatible analyses if requested
      if (_listAnalyses) {
        for (const std::string& ana : _ah.analysisNames()) {
          cout << ana << endl;
        }
      }

      return true;
    }
    catch (const Error &err){
      std::cerr << "Error in run initialisation: " << err.what();
      exit(1);
    }
  }


  bool Run::processEvent() {
    try {
      _ah.analyze(*_evt);
      return true;
    }
    catch (const Error &err){
      std::cerr << "Error in processing event: " << err.what();
      exit(1);
    }
  }


  bool Run::finalize() {
    try {
      _evt.reset();
      _ah.finalize();
      return true;
    }
    catch (const Error &err){
      std::cerr << "Error in run finalize: " << err.what();
      exit(1);
    }
  }


}
