#include <string>
#include <iostream>
#include <iomanip>
#include <map>
#include <sstream>

#include "TROOT.h"
#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TProfile.h"

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

#include "JAAnet/JHead.hh"
#include "JAAnet/JHeadToolkit.hh"
#include "JAAnet/JAAnetToolkit.hh"

#include "JDetector/JDetector.hh"
#include "JDetector/JDetectorToolkit.hh"
#include "JDetector/JPMTRouter.hh"
#include "JGeometry3D/JCylinder3D.hh"

#include "JSupport/JMultipleFileScanner.hh"
#include "JSupport/JMonteCarloFileSupportkit.hh"
#include "JSupport/JSupport.hh"

#include "JROOT/JManager.hh"
#include "JGizmo/JGizmoToolkit.hh"
#include "JTools/JWeight.hh"
#include "JROOT/JRootToolkit.hh"
#include "JPhysics/JPDFToolkit.hh"

#include "JLang/JException.hh"
#include "JLang/JPredicate.hh"
#include "JLang/JLangToolkit.hh"

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


namespace {

  /**
   * Get weight for hits from muon.
   *
   * \param  E             muon energy [GeV]
   * \param  npe           number of photo-electrons
   * \return               weight
   */
  inline double getMuonWeight(const double E,
			      const double npe)
  {
    using namespace JPP;

    return npe / (1.0 + gWater.getB() * E * geanc());
  }


  /**
   * Get weight for hits from muon.
   *
   * \param  E             muon energy [GeV]
   * \param  R             minimal distance of approach [m]
   * \param  npe           number of photo-electrons
   * \return               weight
   */
  inline double getMuonWeight(const double E,
			      const double R,
			      const double npe)
  {
    using namespace JPP;

    static const double l_att = 45.0 * getSinThetaC();   // [m]

    return R * exp(+R/l_att) * getMuonWeight(E, npe);
  }


  /**
   * Get weight for hits from vertex.
   *
   * \param  E             neutrino energy [GeV]
   * \param  npe           number of photo-electrons
   * \return               weight
   */
  inline double getNeutrinoWeight(const double E,
				  const double npe)
  {
    return npe / E;
  }


  /**
   * Get weight for hits from vertex.
   *
   * \param  E             neutrino energy [GeV]
   * \param  D             distance [m]
   * \param  npe           number of photo-electrons
   * \return               weight
   */
  inline double getNeutrinoWeight(const double E,
				  const double D,
				  const double npe)
  {
    static const double l_att = 45.0;                    // [m]

    return D*D * exp(+D/l_att) * getNeutrinoWeight(E, npe);
  }


  /**
   * Function that returns true if second element of first pair is larger than second element of second pair.
   *
   * Implemented to provide sort functionality for function `makeSortedH1`
   *
   * \param first    first  instance of a pair
   * \param second   second instance of a pair
   * \return         first.second > second.second
   */
  inline bool compare(std::pair<Double_t, Double_t> const& first, std::pair<Int_t, Double_t> const& second)
  {
    return first.second > second.second;
  }
  

  /**
   * Function to sort input histogram in order of decreasing bin contents.
   *
   * \param  in     pointer to input histogram
   * \param  name   name of sorted histogram
   * \return        pointer to sorted histogram
   */
  inline TH1D* makeSortedH1(TH1D* in, const TString& name)
  {
    using namespace std;
    using namespace JPP;
   
    // read bin contents to a vector

    vector< pair<Double_t, Double_t> > data;

    for (Int_t ix = 1; ix <= in->GetXaxis()->GetNbins(); ++ix) {

      const Int_t    x = (Int_t) in->GetXaxis()->GetBinCenter(ix);
      const Double_t y = in->GetBinContent(ix);

      if (y > 0.0) {	
	data.push_back(make_pair(x, y));
      }
    }

    // sort

    sort(data.begin(), data.end(), compare);

    // create a histogram

    TH1D* out = new TH1D(name, NULL, data.size(), -0.5, data.size() - 0.5);

    for (Int_t i = 0; i < (Int_t) data.size(); i++) {

      const Int_t ix = i + 1;

      out->GetXaxis()->SetBinLabel(ix, MAKE_CSTRING(fixed << setprecision(0) << data[i].first));
      out->SetBinContent(ix, data[i].second);
    }

    return out;    
  }
}


/**
 * \file
 *
 * Example program to verify Monte Carlo data.
 * \author mdejong
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;

  //-------------------------------------------------------------------------
  // command-line arguments
  //-------------------------------------------------------------------------
  
  JMultipleFileScanner<Evt> inputFile;
  JLimit_t&      numberOfEvents = inputFile.getLimit();
  string         outputFile;
  string         detectorFile;
  set<int>       hit_types;
  int            debug;

  try { 

    JParser<> zap("Example program to verify Monte Carlo data.");
    
    zap['f'] = make_field(inputFile);
    zap['n'] = make_field(numberOfEvents)      = JLimit::max();
    zap['o'] = make_field(outputFile)          = "domino.root";
    zap['a'] = make_field(detectorFile)        = "";
    zap['T'] = make_field(hit_types)           = JPARSER::initialised();
    zap['d'] = make_field(debug)               = 2;
    
    zap(argc, argv);
  }
  catch(const exception &error) {
    FATAL(error.what() << endl);
  }


  //-------------------------------------------------------------------------
  // load detector file
  //-------------------------------------------------------------------------
  
  JDetector detector;

  if (detectorFile != "") {

    try {
      load(detectorFile, detector);
    }
    catch(const JException& error) {
      FATAL(error);
    }
  }

  const JPMTRouter router(detector);

  const JHead header = getHeader(inputFile);


  //-------------------------------------------------------------------------
  // try to determine coordinate origin
  //-------------------------------------------------------------------------

  const Vec offset = getOffset(header);
  
  NOTICE("Apply detector offset " << offset << endl); 

  detector -= getPosition(offset);


  //-------------------------------------------------------------------------
  // calculate instrumented volume
  //-------------------------------------------------------------------------
  
  JCylinder3D inst_vol(detector.begin(), detector.end());

  NOTICE("JDomino: Instrumented volume dimensions (x, y, r, zmin, zmax): " << inst_vol << endl);


  //-------------------------------------------------------------------------
  // histogram binning
  //-------------------------------------------------------------------------

  typedef map<int, JWeight> map_type;

  double xmin = 2.0;
  double xmax = 8.0;

  // energy range is taken from -in order of priority- neutrino or muon data

  if        (header.is_valid(&JHead::cut_nu)) {
    xmin = log10(header.cut_nu.E.getLowerLimit());
    xmax = log10(header.cut_nu.E.getUpperLimit());
  } else if (header.is_valid(&JHead::cut_in)) {
    xmin = log10(header.cut_in.E.getLowerLimit());
    xmax = log10(header.cut_in.E.getUpperLimit());
  } 
  const int nx   = (int) ((xmax - xmin) / 0.1);

  const Double_t T[] = { -50.0, -20.0, -10.0, -5.0, -2.0, 0.0, +2.0, +5.0, +10.0, +15.0, +20.0, +30.0, +40.0, +50.0,
			 +75.0, +100.0, +150.0, +200.0, +250.0, +300.0, +400.0, +500.0, 1000.0 };  // [ns]
  
  //-------------------------------------------------------------------------
  // histograms for events
  //-------------------------------------------------------------------------
  
  TH1D hits("hits", NULL,    100,     0.0,      8.0);                // number of mc_hits per event
  TH1D trks("trks", NULL,  10001, -5000.5,   5000.5);                // number of track_in type's per event
  TH1D job ("job" , NULL,  10001, -5000.5,   5000.5);                // number of photo-electrons by track_in type

  TProfile hits_per_E_in ("hits_per_E_in",  "average number of hits per E_{#nu} inside  instrumented volume", nx, xmin, xmax);
  TProfile hits_per_E_out("hits_per_E_out", "average number of hits per E_{#nu} outside instrumented volume", nx, xmin, xmax);

  typedef JManager<int, TH1D>  JManager_t;

  JManager_t npe_per_type(new TH1D("npe[%]", NULL, 5000, 0.5, 5000.5), '%', ios::fmtflags(ios::showpos));

  TH1D       npe_per_pmt ("pmt", NULL, 100001, -0.5, 100000.5);      // number of photo-electrons per PMT
    

  //-------------------------------------------------------------------------
  // histograms for neutrino
  //-------------------------------------------------------------------------

  TH2D  nuExD("nuExD", NULL, nx, xmin,  xmax, 100,  0.0, 1000.0);    // Enu versus distance between vertex and PMT
  TH2D* nuExD_tmp = (TH2D*) nuExD.Clone("nuExD.tmp");                // normalisation

  TH2D  nuExc("nuExc", NULL, nx, xmin,  xmax, 100, -1.0,   +1.0);    // Enu versus emission angle
  TH2D* nuExc_tmp = (TH2D*) nuExc.Clone("nuExc.tmp");                // normalisation

  TH2D  nuDxc("nuDxc", NULL, 50, 0.0, 1000.0, 100, -1.0,   +1.0);    // distance between vertex and PMT versus emission  angle
  TH2D* nuDxc_tmp = (TH2D*) nuDxc.Clone("nuDxc.tmp");                // normalisation

  TH2D  nuDxU("nuDxU", NULL, 50, 0.0, 1000.0, 100, -1.0,   +1.0);    // distance between vertex and PMT versus incidence angle
  TH2D* nuDxU_tmp = (TH2D*) nuDxU.Clone("nuDxU.tmp");                // normalisation

  TH2D  nuDxT("nuDxT", NULL, 50, 0.0, 1000.0, getSize(T) - 1, T);    // distance between vertex and PMT versus time residual
  TH2D* nuDxT_tmp = (TH2D*) nuDxT.Clone("nuDxT.tmp");                // normalisation

  nuExD.Sumw2();
  nuExc.Sumw2();
  nuDxc.Sumw2();
  nuDxU.Sumw2();
  nuDxT.Sumw2();


  //-------------------------------------------------------------------------
  // histograms for muon
  //-------------------------------------------------------------------------  

  TH2D  muExR("muExR", NULL, nx, xmin,  xmax,  30,  0.0,  300.0);    // Emu versus distance between muon and PMT
  TH2D* muExR_tmp = (TH2D*) muExR.Clone("muExR.tmp");                // normalisation

  TH2D  muRxU("muRxU", NULL, 15, 0.0,  300.0, 100, -1.0,   +1.0);    // distance between muon and PMT versus incidence angle
  TH2D* muRxU_tmp = (TH2D*) muRxU.Clone("muRxU.tmp");                // normalisation

  TH2D  muRxT("muRxT", NULL, 15, 0.0,  300.0, getSize(T) - 1, T);    // distance between muon and PMT versus time residual
  TH2D* muRxT_tmp = (TH2D*) muRxT.Clone("muRxT.tmp");                // normalisation

  muExR.Sumw2();
  muRxU.Sumw2();
  muRxT.Sumw2();


  //-------------------------------------------------------------------------
  // loop over events
  //-------------------------------------------------------------------------
  
  while (inputFile.hasNext()) {

    STATUS("event: " << setw(10) << inputFile.getCounter() << '\r'); DEBUG(endl);

    const Evt* event = inputFile.next();

    double NPE = 0.0;

    for (vector<Hit>::const_iterator hit = event->mc_hits.begin(); hit != event->mc_hits.end(); ++hit) {
      NPE += getNPE(*hit);
    }

    hits.Fill(log10((Double_t) event->mc_hits.size()));
        
    for (vector<Trk>::const_iterator track = event->mc_trks.begin(); track != event->mc_trks.end(); ++track) {
      trks.Fill(track->type);
    }

    JTrack3E neutrino;

    if (has_neutrino(*event)) {
      neutrino = getTrack(get_neutrino(*event));
    }

    const JVertex3D vertex = neutrino.getVertex();

    map_type npe_pmt;   // temporary buffer to accumulate total npe's by individual PMTs
    map_type npe_type;  // temporary buffer to accumulate total npe's by hit origin track


    //-------------------------------------------------------------------------
    // loop over hits
    //-------------------------------------------------------------------------
    
    for (vector<Hit>::const_iterator hit = event->mc_hits.begin(); hit != event->mc_hits.end(); ++hit) {

      const int    type = hit->type;
      const double t1   = getTime(*hit);
      const double npe  = getNPE (*hit);
	      
      npe_pmt[hit->pmt_id].put(npe);

      npe_type[0]   .put(npe);
      npe_type[type].put(npe);	  

      if (hit_types.empty() || hit_types.count(type) != 0) {

	vector<Trk>::const_iterator track = find_if(event->mc_trks.begin(),
						    event->mc_trks.end(),
						    make_predicate(&Trk::id, hit->origin));
	
	if (track == event->mc_trks.end()) {
	  ERROR("Hit " << *hit << " has no origin." << endl);
	  continue;
	}
	
	if (count_if(event->mc_trks.begin(),
		     event->mc_trks.end(),
		     make_predicate(&Trk::id, hit->origin)) > 1) {
	  ERROR("Hit " << *hit << " has ambiguous origin." << endl);
	  continue;
	}
      
	job.Fill((double) track->type, npe);

	if (router.hasPMT(hit->pmt_id)) {
	      
	  const JPMT&  pmt  = router.getPMT(hit->pmt_id);

	  if (is_muon(*track)) {
	  
	    const JTrack3E muon = getTrack(*track);
	    
	    const double E   = muon.getE();
	    const double x   = log10(E);
	    const double R   = muon.getDistance(pmt);
	    const double t0  = muon.getT       (pmt);
	    const double ct1 = muon.getDot     (pmt);

	    muExR.Fill(x, R,       getMuonWeight(E, npe));
	    muRxU.Fill(R, ct1,     getMuonWeight(E, R, npe));
	    muRxT.Fill(R, t1 - t0, getMuonWeight(E, R, npe));

	  } else if (has_neutrino(*event)) {

	    const double E   = neutrino.getE();
	    const double x   = log10(E);
	    const double D   = vertex.getDistance(pmt);
	    const double t0  = vertex.getT(pmt);
	    const double ct0 = neutrino.getDot(pmt.getPosition());
	    const double ct1 = vertex.getDot(pmt);

	    nuExD.Fill(x, D,       getNeutrinoWeight(E, npe));
	    nuExc.Fill(x, ct0,     getNeutrinoWeight(E, npe));
	    nuDxc.Fill(D, ct0,     getNeutrinoWeight(E, D, npe));
	    nuDxU.Fill(D, ct1,     getNeutrinoWeight(E, D, npe));
	    nuDxT.Fill(D, t1 - t0, getNeutrinoWeight(E, D, npe));
	  }
	}	
      } // end if PMT of hit found in PMT router
    } // end loop over hits


    //-------------------------------------------------------------------------
    // continue with logic for event, fill histograms common to all input files
    //-------------------------------------------------------------------------

    for (map_type::const_iterator i = npe_pmt.begin(); i != npe_pmt.end(); ++i) {
      npe_per_pmt.Fill(i->second.getTotal());                     // npe / PMT
    }

    for (map_type::const_iterator i = npe_type.begin(); i != npe_type.end(); ++i) {
      npe_per_type[i->first]->Fill(i->second.getTotal());         // npe / type
    }


    //-------------------------------------------------------------------------
    // normalisation of muon histograms
    //-------------------------------------------------------------------------

    for (vector<Trk>::const_iterator track = event->mc_trks.begin(); track != event->mc_trks.end(); ++track) {

      if (is_muon(*track)) {
	
	const JTrack3E muon = getTrack(*track);

	const double E   = muon.getE();
	const double x   = log10(E);

	for (JDetector::const_iterator module = detector.begin(); module != detector.end(); ++module) {

	  const double R   = muon.getDistance(*module);

	  muExR_tmp->Fill(x, R, module->size());
   
	  if (R < muRxU.GetXaxis()->GetXmax()) {
	    for (JModule::const_iterator pmt = module->begin(); pmt != module->end(); ++pmt) {
	      muRxU_tmp->Fill(R, muon.getDot(*pmt));
	    }
	  }

	  if (R < muRxT.GetXaxis()->GetXmax()) {
	    for (Int_t iy = 1; iy <= muRxT_tmp->GetYaxis()->GetNbins(); ++iy) {
	      muRxT_tmp->Fill(R, muRxT_tmp->GetYaxis()->GetBinCenter(iy), muRxT_tmp->GetYaxis()->GetBinWidth(iy));
	    }
	  }
	}
      }
    }


    //-------------------------------------------------------------------------
    // normalisation of neutrino histograms
    //-------------------------------------------------------------------------
    
    if (has_neutrino(*event)) {

      const double E   = neutrino.getE();
      const double x   = log10(E);

      if (inst_vol.is_inside(vertex))
	hits_per_E_in .Fill(x, NPE);
      else
	hits_per_E_out.Fill(x, NPE);

      for (JDetector::const_iterator module = detector.begin(); module != detector.end(); ++module) {

	const double D   = vertex.getDistance(*module);
	const double ct0 = neutrino.getDot(module->getPosition());
	  
	nuExD_tmp->Fill(x, D,   module->size());
	nuExc_tmp->Fill(x, ct0, module->size());
	nuDxc_tmp->Fill(D, ct0, module->size());

	if (D < nuDxU.GetXaxis()->GetXmax()) {
	  for (JModule::const_iterator pmt = module->begin(); pmt != module->end(); ++pmt) {
	    nuDxU_tmp->Fill(D, neutrino.getDot(*pmt));
	  }
	}
	
	if (D < nuDxT.GetXaxis()->GetXmax()) {
	  for (Int_t iy = 1; iy <= nuDxT_tmp->GetYaxis()->GetNbins(); ++iy) {
	    nuDxT_tmp->Fill(D, nuDxT_tmp->GetYaxis()->GetBinCenter(iy), nuDxT_tmp->GetYaxis()->GetBinWidth(iy));
	  }
	}
      }
    }    
  } // end loop over events

  STATUS(endl); 


  //-------------------------------------------------------------------------
  // normalisation of histograms
  //-------------------------------------------------------------------------

  TH1D *job_sorted  = makeSortedH1(&job,  "hits_by_pdg"); // hits by track_in particles (PDG codes) per event, sorted
  TH1D *trks_sorted = makeSortedH1(&trks, "trks_sorted"); // track_in particles (PDG codes) per event, sorted
  
  if (inputFile.getCounter() != 0) {

    const Double_t W = 1.0 / ((Double_t) inputFile.getCounter());
    
    for (TH1D* p : { &hits, &npe_per_pmt, &job, &trks, job_sorted, trks_sorted }) {
      p->Scale(W);
    }

    for (JManager_t::iterator i = npe_per_type.begin(); i != npe_per_type.end(); ++i) {
      i->second->Scale(W);
    }

    nuExD.Divide(nuExD_tmp);
    nuExc.Divide(nuExc_tmp);
    nuDxc.Divide(nuDxc_tmp);
    nuDxU.Divide(nuDxU_tmp);
    nuDxT.Divide(nuDxT_tmp);

    muExR.Divide(muExR_tmp);
    muRxU.Divide(muRxU_tmp);
    muRxT.Divide(muRxT_tmp);
  }


  //-------------------------------------------------------------------------
  // output
  //-------------------------------------------------------------------------
  
  TFile out(outputFile.c_str(), "recreate");

  out << hits << job << *job_sorted << trks << *trks_sorted << hits_per_E_in << hits_per_E_out << npe_per_type << npe_per_pmt;

  out << nuExD << nuExc << nuDxc << nuDxU << nuDxT;
  out << muExR << muRxU << muRxT;

  out.Write();
  out.Close();
}