#include <iostream>
#include <iomanip>
#include <string>
#include <vector>
#include <set>
#include <stdio.h>

#include "km3net-dataformat/offline/Head.hh"
#include "km3net-dataformat/offline/Evt.hh"
#include "km3net-dataformat/definitions/trkmembers.hh"

#include "Jeep/JPrint.hh"
#include "Jeep/JParser.hh"
#include "Jeep/JMessage.hh"
#include "Jeep/JeepToolkit.hh"

#include "JLang/JException.hh"

#include "JTools/JRange.hh"

#include "JSupport/JMeta.hh"
#include "JSupport/JSupport.hh"
#include "JSupport/JFileRecorder.hh"
#include "JSupport/JMultipleFileScanner.hh"
#include "JSupport/JEvtWeightFileScanner.hh"
#include "JSupport/JEvtWeightFileScannerSet.hh"

#include "JAAnet/JHead.hh"


namespace {

  using JLANG::JException;
  using JLANG::JValueOutOfRange;

  using JTOOLS::JRange;

  using JAANET::cut;
  
  using JSUPPORT::JFileRecorder;
  using JSUPPORT::JMultipleFileScanner;
  using JSUPPORT::JEvtWeightFileScanner;

  
  /**
   * Auxiliary class containing methods for separating a collection of Monte-Carlo files\n
   * into output files which are disjoint with respect to a specifiable JAANET::cut header-field.
   */
  struct JSeparateCuts
  {
    typedef typename JTYPELIST<Head, JMeta, Evt>::typelist   typelist;
    typedef JRange<double>                                   JRange_t;

    /**
     * Default constructor.
     */
    JSeparateCuts() :
      filenameTemplate("cut%.root"),
      wildcard        ('%')
    {
      pos = filenameTemplate.find(wildcard);
    }
    
    
    /**
     * Constructor.
     *
     * \param  templateName                  file-name template; must contain wildcard
     * \param  wc                            wildcard character                       
     * \param  metaInfo                      meta-information
     */
    JSeparateCuts(const std::string& templateName,
		  const char         wc,
		  const JMeta&       metaInfo = JMeta()) :
      meta            (metaInfo),
      filenameTemplate(templateName),
      wildcard        (wc)
    {
      if (filenameTemplate.find(wildcard) == std::string::npos) {
	THROW(JException,
	      "JSeparateCuts::JSeparateCuts(): Given file-name template \"" << filenameTemplate <<
	      "\" does not contain specified wildcard \'"                   << wildcard << "\'.");
      }
      
      pos = filenameTemplate.find(wildcard);
    }


    /*
     * Separate two cuts into a vector of fully disjoint, non-overlapping cuts.
     *
     * \param  first           first  cut
     * \param  second          second cut
     * \return                 vector of disjoint cuts
     */
    std::vector<cut> separate(const cut& first, const cut& second) const {

      using namespace std;
      using namespace JPP;

      vector<cut> out;

      const set<double> Ebounds   { first .E.getLowerLimit(), first .E.getUpperLimit(),
	                            second.E.getLowerLimit(), second.E.getUpperLimit() };

      const set<double> cosTbounds{ first .cosT.getLowerLimit(), first .cosT.getUpperLimit(),
	                            second.cosT.getLowerLimit(), second.cosT.getUpperLimit() };

      for (set<double>::const_iterator Ebound = next(Ebounds.cbegin()); Ebound != Ebounds.cend(); ++Ebound) {

	const JRange_t E(*prev(Ebound), *Ebound);

	// Check if given energy range appears in only one or in both cuts

	const bool overlap1 = overlap(E, first .E);
	const bool overlap2 = overlap(E, second.E);

	if (overlap1 && overlap2) { // Loop through all disjoint cosine zenith angle ranges

	  for (set<double>::const_iterator cosTbound = next(cosTbounds.cbegin()); cosTbound != cosTbounds.cend(); ++cosTbound) {
	    out.push_back(cut(*prev(Ebound), *Ebound, *prev(cosTbound), *cosTbound));
	  }
	
	} else if (overlap1) {      // Use first  cosine zenith angle range
	  
	  out.push_back(cut(*prev(Ebound), *Ebound, first.cosT.getLowerLimit(), first.cosT.getUpperLimit()));
			
	} else if (overlap2) {      // Use second cosine zenith angle range
	  
	  out.push_back(cut(*prev(Ebound), *Ebound, second.cosT.getLowerLimit(), second.cosT.getUpperLimit()));
	}
      }
    
      return out;
    }

  
    /**
     * Extract all events in the given file corresponding to a given cut\n
     * and append them to the given file-recorder.
     *
     * \param  recorder             file-recorder
     * \param  scanner              file-scanner
     * \param  cut                  JAANET::cut object
     * \return                      number of extracted events
     */
    size_t extract(JFileRecorder<typelist>& recorder,
		   JEvtWeightFileScanner<>& scanner,
		   const cut&               cut) const
    {
      using namespace std;
      
      const bool open = recorder.is_open();
      if (!open) { recorder.open(); }
      scanner.rewind();

      size_t n = 0;
      
      while (scanner.hasNext()) {

	const Evt* event = scanner.next();

	if (event != NULL && !event->mc_trks.empty()) {

	  double E          = 0.0;
	  double costh      = 0.0;
	  double Nprimaries = 0.0;

	  for (vector<Trk>::const_iterator trk = event->mc_trks.cbegin(); trk != event->mc_trks.cend(); ++trk) {
	  
	    if (trk->status == TRK_ST_PRIMARYNEUTRINO ||
		trk->status == TRK_ST_PRIMARYCOSMIC) {

	      E     += trk->E;
	      costh += trk->dir.z / trk->dir.len();
	      
	      ++Nprimaries;
	    }
	  }
	
	  costh /= Nprimaries;
	
	  if (cut.E(E) && cut.cosT(costh)) {
	    recorder.put(*event);
	    ++n;
	  }
	}
      }
    
      if (!open) { recorder.close(); }
      scanner.rewind();

      return n;
    }
    

    /**
     * Separate two file-scanners into disjoint output files which do not overlap in the given cut data member.\n
     * The separation will only proceed if the given cut data member is the only other mismatching header-field\n
     * other than JHead::spectrum.
     *
     * \param  first                first  file-scanner
     * \param  second               second file-scanner
     * \param  pd                   pointer to JAANET::JHead data member derived from JAANET::cut
     * \return                      number of extracted disjoing output files
     */
    template<class T>
    size_t separate(JEvtWeightFileScanner<>& first,
		    JEvtWeightFileScanner<>& second,
		    T JHead::*               pd)
    {
      using namespace  std;
    
      typedef typename JTYPELIST<Head, JMeta, Evt>::typelist    typelist;

      const JHead& header0 = first .getHeader();
      const JHead& header1 = second.getHeader();

      // Check if given data member and the spectrum data member correspond to the only mismatching header-fields

      const size_t N = JHead::getMaximumNumberOfMatches() - (header0.spectrum.match(header1.spectrum) ? 1 : 2);
      
      if (!((header0.*pd).match(header1.*pd)) && header0.getNumberOfMatches(header1) == N) {
	
	vector<cut> cuts = separate(header0.*pd, header1.*pd);

	for (vector<cut>::iterator cut = cuts.begin(); cut != cuts.end(); ++cut) {
	  
	  // Open new file-recorder
	  
	  const string outputFilename = string(filenameTemplate).replace(pos, 1, MAKE_CSTRING(distance(cuts.begin(), cut)));
      
	  JFileRecorder<typelist> recorder(outputFilename.c_str());

	  // Write header and meta-information
	  
	  JHead header2(header0);
    
	  (header2.*pd).E    = cut->E;
	  (header2.*pd).cosT = cut->cosT;

	  header2.genvol.numberOfEvents += header1.genvol.numberOfEvents;

	  Head head;
	  copy(header2, head);
    
	  recorder.open();
	  recorder.put(head);
	  recorder.put(meta);

	  for (JMultipleFileScanner<JMeta> in(first.getFilelist()); in.hasNext(); ) {
	    recorder.put(*in.next());
	  }

	  for (JMultipleFileScanner<JMeta> in(second.getFilelist()); in.hasNext(); ) {
	    recorder.put(*in.next());
	  }

	  // Write events that are compatible with the cut

	  extract(recorder, first,  *cut);
	  extract(recorder, second, *cut);

	  recorder.close();
	}

	return cuts.size();
	
      } else {
	
	return 0;
      }
    }


    /**
     * Set filename template.
     *
     * \param  _filenameTemplate                 
     */
    void setFilenameTemplate(const std::string& _filenameTemplate)
    {
      std::size_t p = _filenameTemplate.find(wildcard);
      
      if (p != std::string::npos) {
	
	this->filenameTemplate = _filenameTemplate;
	this->pos              = p;
	
      } else {
	
	THROW(JValueOutOfRange, "JSeparateCuts::setFilenameTemplate(): Given filename template " << filenameTemplate << " does not contain wildcard " << wildcard);
      }
    }


    /**
     * Set wildcard character.
     *
     * \param  wildcard             wildcard character
     */
    void setWildcard(const char wc)
    {
      std::size_t p = this->filenameTemplate.find(wc);

      if (p != std::string::npos) {
	
	this->wildcard = wc;
	this->pos      = p;
	
      } else {
	
	THROW(JValueOutOfRange, "JSeparateCuts::setWildcard(): Given wildcard character " << wildcard << " is not contained in current filename template " << this->filenameTemplate);
      }
    }


    /**
     * Get filename template.
     *
     * \return                      filename template
     */
    const std::string& getFilenameTemplate()
    {
      return filenameTemplate;
    }


    /**
     * Get wildcard character.
     *
     * \return                       wildcard character
     */
    const char getWildcard()
    {
      return wildcard;
    }
    
    
    JMeta       meta;             ///!< Meta-data
    
  private:
    
    std::string filenameTemplate; ///!< file-name template
    char        wildcard;         ///!< wildcard character
    size_t      pos;              ///!< position of wildcard in file-name template
  };
}


/**
 * \file
 * Program for extracting disjoint output files without overlapping cuts\n
 * from a list of Monte-Carlo files with overlapping cuts.
 *
 * Note: Already existing disjoint output files which have been extracted previously will be overwritten.
 *
 * \author bjung
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;
  
  JMultipleFileScanner_t  inputFiles;
  string                  outputFile;
  char                    wildcard;
  bool                    replace;

  int                     debug;

  try {
    
    JParser<> zap;
    
    zap['f'] = make_field(inputFiles);
    zap['o'] = make_field(outputFile)         = "%.root";
    zap['w'] = make_field(wildcard)           = '%';
    zap['r'] = make_field(replace);
    zap['d'] = make_field(debug)              = 1;

    zap(argc, argv);
  }
  catch(const exception& error) {
    FATAL(error.what() << endl);
  }

  const size_t pos = outputFile.find(wildcard);

  if (pos == string::npos) {
    FATAL("Valid wildcard must be specified (<" << outputFile << "> does not contain \'" << wildcard << "\').");
  }

  
  // Create set of files ordered based on header-info
  
  JEvtWeightFileScannerSet<> scanners(inputFiles);

  // Split files with overlapping cuts into separate files with disjoint cuts.

  set<string>   buffer;
  JSeparateCuts separator(outputFile, wildcard, JMeta(argc, argv));

  for (JEvtWeightFileScannerSet<>::iterator scanner1 = next(scanners.begin()); scanner1 != scanners.end(); ++scanner1) {

    JEvtWeightFileScannerSet<>::iterator scanner0 = prev(scanner1);
    
    if (scanner0->hasNext() && scanner1->hasNext()) {      
      
      const string identifier = scanners.getUniqueIdentifier(scanner0) + MAKE_STRING(".cut" << wildcard);
      const string name       = string(outputFile).replace(pos, 1, identifier.c_str());
      
      separator.setFilenameTemplate(name);

      const size_t N = (separator.separate(*scanner0, *scanner1, &JHead::cut_primary) +
			separator.separate(*scanner0, *scanner1, &JHead::cut_seamuon) +
			separator.separate(*scanner0, *scanner1, &JHead::cut_in)      +
			separator.separate(*scanner0, *scanner1, &JHead::cut_nu));

      STATUS("\rExtracted "        << N << " output files with"            <<
	     " disjoint cuts for " << getFilename(scanner0->getFilename()) <<
	     " and "               << getFilename(scanner1->getFilename())); DEBUG(endl);

      if (N > 0) {
	
	buffer.insert(scanner0->getFilename());
	buffer.insert(scanner1->getFilename());
      }
    }
  }
  
  for (set<string>::const_iterator i = buffer.begin(); replace && i != buffer.end(); ++i) {
      
    if (remove(i->c_str()) != 0) {
      FATAL(endl << "Error removing file " << *i);
    } else {
      NOTICE(endl << "Removing file " << *i);
    }
  }

  return 0;
}