#include <string>
#include <iostream>
#include <fstream>
#include <iomanip>
#include <stdexcept>

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

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

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

#include "JPhysics/JPDFToolkit.hh"
#include "JPhysics/JDispersion.hh"
#include "JPhysics/Antares.hh"
#include "JPhysics/KM3NeT.hh"
#include "JIO/JFileStreamIO.hh"

#include "JGeometry3D/JGeometry3DToolkit.hh"
#include "JGeometry3D/JPosition3D.hh"
#include "JGeometry3D/JDirection3D.hh"
#include "JGeometry3D/JAxis3D.hh"

#include "JPhysics/JConstants.hh"
#include "JTools/JHistogram1D_t.hh"
#include "JTools/JHistogramMap_t.hh"
#include "JTools/JTransformableMultiHistogram.hh"
#include "JTools/JFunction1D_t.hh"
#include "JTools/JFunctionalMap_t.hh"

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

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

#include "JPhysics/JPDFTransformer.hh"
#include "JPhysics/JPDF.hh"

#include "JAAnet/JParticleTypes.hh"

/**
 * \file
 *
 * Program to histogram event-by-event and track-by-track for making PDFs of lights of individual particles.
 * \author jseneca
 */
inline std::vector<Int_t> getHitRemapping(const std::vector<Trk>* tracklist)
{
  using namespace std;
  using namespace JPP;

  vector<Int_t> hit_remapping;

  for (vector<Trk>::const_iterator i = tracklist->begin(); i != tracklist->end(); ++i) {

    //Skip pathological neutrino that is not the primary incoming neutrino

    if(i->type == TRACK_TYPE_NUMU && i != tracklist->begin()) {
      continue;
    }

    hit_remapping.push_back(i->id);
  }

  return hit_remapping;
}

int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;

  JMultipleFileScanner<Evt> inputFile;
  //JLimit_t&   numberOfEvents = inputFile.getLimit();
  string      outputFile = "J%s.HDE"; //'s' for single (particle), or simulation
  string      detectorFile;
  vector<int> particles;
  uint        fix_bug;
  int         debug;

  try {

    JParser<> zap("Program to histogram event-by-event data of shower light for making PDFs.");

    zap['f'] = make_field(inputFile);
    zap['o'] = make_field(outputFile);
    //zap['n'] = make_field(numberOfEvents)         = JLimit::max();
    zap['a'] = make_field(detectorFile);
    zap['p'] = make_field(particles);
    zap['b'] = make_field(fix_bug);
    zap['d'] = make_field(debug)                  = 1;
    
    if (zap.read(argc, argv) != 0)
      return 1;
  }
  catch(const exception &error) {
    FATAL(error.what() << endl);
  }

  JDetector detector;

  if(outputFile.find('%') != string::npos){
    sort(particles.begin(), particles.end());
    string prefix_t = "";
    for(vector<int>::iterator ptype = particles.begin(); ptype != particles.end(); ptype++){
      prefix_t += to_string((long long int)*ptype) + "_";
    }
    prefix_t.erase(prefix_t.size() - 1);
    string::size_type pos_1 = outputFile.find('%');
    outputFile.replace(pos_1, 1, prefix_t);
  }
  
  try {
    load(detectorFile, detector);
  }
  catch(const JException& error) {
    FATAL(error);
  }

  const JPosition3D center = get<JPosition3D>(getHeader(inputFile));

  NOTICE("Apply detector offset " << center << endl);

  detector -= center;

  const JPMTRouter pmtRouter(detector);
  
  const double P_atm = NAMESPACE::getAmbientPressure();
  const double wmin  = getMinimalWavelength();
  const double wmax  = getMaximalWavelength();

  const JDispersion dispersion(P_atm);

  const double ng[] = { dispersion.getIndexOfRefractionGroup(wmax),
			dispersion.getIndexOfRefractionGroup(wmin) };

  typedef JHistogram1D_t::abscissa_type                            abscissa_type;

  typedef JTransformableMultiHistogram<JHistogram1D_t,
                               JMAPLIST<JHistogramMap_t,
                                        JHistogramMap_t,
                                        JHistogramMap_t,
                                    JHistogramGridMap_t,
                                    JHistogramGridMap_t>::maplist> JMultiHistogram_t;
  
  typedef JPDFTransformer<5, abscissa_type>  JFunction5DTransformer_t;

  JMultiHistogram_t h0;   // occurrence rate of PMT (used for normalisation)
  JMultiHistogram_t h1;   // light from track cascade

  h1.transformer.reset(new JFunction5DTransformer_t(21.5, 2, ng[0], 0.0, JGeant(JGeanx(1.00, -2.2)), 1e-2, NAMESPACE::getAngularAcceptance, 0.05));
  
  set<double> C;  // cosine emission angle

  JQuadrature qeant(-1.0, 1.0, 30, JGeanx(1.00, -2.2));

  for (JQuadrature::const_iterator i = qeant.begin(); i != qeant.end(); ++i) {
    C.insert(i->getX());
  }

  C.insert(-1.01);
  C.insert(-1.00);
  C.insert( 1.00);
  C.insert( 1.01);

  const double E_b[] = {0.01, 0.2, 0.4, 0.6, 0.8, 1.0, 1.5, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0, 20.0, 25.0, 30.0, 40.0, 55.0, 70.0, 85.0, 100.0};

  const double R[]   = {0.0, 1.0, 2.5, 5.0, 7.5, 10.0, 12.5, 15.0, 17.5, 20.0, 25.0, 30.0, 35.0, 40.0, 50.0, 65.0, 80.0};
  const double Dmax_m = R[sizeof(R)/sizeof(R[0]) - 1]; //Last element of R

  for (int iE = 0; iE != sizeof(E_b)/sizeof(E_b[0]); ++iE) {

    const double E_m = E_b[iE];

    for (int i = 0; i != sizeof(R)/sizeof(R[0]); ++i) {

      const double R_m = R[i];

      for (set<double>::const_iterator c = C.begin(); c != C.end(); ++c) {

	const double cd = *c;
	
	const double grid  = 10.0 +  0.0 * R_m/100.0;                    // [deg]
	const double alpha = 2.0 * sqrt(1.0 - cos(grid * PI / 180.0));   // azimuth angle unit step size

	const int    number_of_theta_points = (max(2, (int) (180.0/(1.4 * grid))));
	const double theta_step             = PI / (number_of_theta_points + 1);
	
	for (double theta = -0.5*theta_step; theta < PI + theta_step; theta += theta_step) {	  
	  const int    number_of_phi_points = (max(2, (int) (PI * sin(theta) / alpha)));
	  const double phi_step             = PI / (number_of_phi_points + 1);
	  
	  for (double phi = -0.5*phi_step; phi < PI + phi_step; phi += phi_step) {

	    for (JMultiHistogram_t* p : { &h0, &h1 }) {
	      (*p)[E_m][R_m][cd][theta][phi];
	    }
	  }
	}
      }
    }
  }
  
  double buffer[] = {-15.0,-10.0,-7.5,-5,-4,-3,-2,-1.5,-1.0,-0.5,-0.1, 0.0,0.1, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 5.0, 7.5, 10.0, 15.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 85.0, 100.0 };
  for (JMultiHistogram_t::super_iterator
         i1 = h1.super_begin(); i1 != h1.super_end(); ++i1) {
    for (int j = 0; j != sizeof(buffer)/sizeof(buffer[0]); ++j) {
      i1.getValue()[buffer[j]];
    }
  }//Buffer defines the binning of the inner histograms (dt)
  
  
  long long int h0_fillcount = 0;
  long long int h1_fillcount = 0;
  
  while (inputFile.hasNext()) {

    //NOTICE("event: " << setw(10) << inputFile.getCounter() << '\r'); STATUS(endl);
    //I suspect the printing above slows down the program considerably

    const Evt* event = inputFile.next();
    const vector<Trk>* mc_tracks = &(event->mc_trks);
    
    vector<int> hit_remap;

    //*************************************************
    //Fixes 2016 atmospheric neutrino simulation hit-making neutrino bug
    //Remove after fix.
    if(fix_bug == 0){
      for (vector<Trk>::const_iterator i = mc_tracks->begin(); i != mc_tracks->end(); ++i) {
	hit_remap.push_back(i->id);
      }
    } else if(fix_bug == 1){
      hit_remap = getHitRemapping(mc_tracks);
    } else if(fix_bug == 2){
      bool skipevent = false;
      for (vector<Trk>::const_iterator i = mc_tracks->begin(); i != mc_tracks->end(); ++i) {
	if(i->type == 14){
	  skipevent = true;
	  break;
	}
	hit_remap.push_back(i->id);
      }
      if(skipevent) continue;
    }
    //*************************************************

    double t0 = (*mc_tracks)[1].t;
    
    for (vector<Hit>::const_iterator hit = event->mc_hits.begin(); hit != event->mc_hits.end(); ++hit)  {
      
      try {
	
	if(hit->origin >= (int)(*mc_tracks).size()){
	  std::out_of_range err("Hit origin not in tracklist. Avoided segfault");
	  throw err;
	}

	Int_t corr_origin  = hit_remap[hit->origin];
	Trk curr_track     = (*mc_tracks)[corr_origin]; //Get track from corrected origin
	JDirection3D dir = getDirection(curr_track);    const JRotation3D R(dir);
	JPosition3D pos = getPosition(curr_track);      	   pos.rotate(R);
	JAxis3D axis = pmtRouter.getPMT(hit->pmt_id);		
	axis.transform(R, pos);
	
	const double E     = curr_track.E;
	const double t1    = t0 + axis.getLength() * getInverseSpeedOfLight() * getIndexOfRefraction();
	const double D_m   = axis.getLength();
	const double cd    = axis.getZ()/D_m;
	const double theta = axis.getTheta();
	const double phi   = fabs(axis.getPhi());
	const double dt    = getTime(*hit) - t1;
	const double npe   = getNPE (*hit);

	if(D_m < Dmax_m){
	  h1.fill(E, D_m, cd, theta, phi, dt, npe);
	  h1_fillcount += 1;
	}	
      }
      catch(const exception& error) {
	std::cout << "Fatal error for event: " << inputFile.getCounter() << std::endl;
	FATAL(error.what() << endl);
      }
    }
    
    //                               ignore primary neutrino ->    ^
    for (vector<Trk>::const_iterator tr = event->mc_trks.begin() + 1; tr != event->mc_trks.end(); ++tr) {
      
      if(find(particles.begin(), particles.end(), tr->type) == particles.end()){
	continue;
      }
      
      //cout << "Particle accepted: " << tr->type << endl;


      JDirection3D dir = getDirection(*tr);    const JRotation3D R(dir);
      JPosition3D pos = getPosition(*tr);	pos.rotate(R);

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

	JPosition3D P = module->getPosition();

	P.rotate(R);

	P.sub(pos);

	if (P.getLength() < Dmax_m) {

	  for (JModule::const_iterator pmt = module->begin(); pmt != module->end(); ++pmt) {
	    
	    JAxis3D axis = *pmt;	    
	    axis.transform(R, pos);
	    
	    h0.fill(tr->E, axis.getLength(), axis.getZ()/axis.getLength(), axis.getTheta(), fabs(axis.getPhi()), 0.0, 1.0);
	    h0_fillcount += 1;
	  }
	}
      }
    }
    
    if(h1_fillcount > h0_fillcount){
      std::cout << "WARNING: recorded hits in normalization histogram that were not recorded in normalization histogram. This should not happen." << std::endl;
      std::cout << "h1_fillcount: " << h1_fillcount << ", h0_fillcount: " << h0_fillcount << std::endl;
    }
  };
  
  double integral = 0;
  for (JMultiHistogram_t::super_iterator i = h0.super_begin(); i != h0.super_end(); ++i) {
    integral+=i.getValue().getIntegral();
  }
  DEBUG("Integral:\t" << integral << endl);

  // output

  JFileStreamWriter out(outputFile.c_str());
  
  NOTICE("Storing " << outputFile << ", " << flush);

  for (const JMultiHistogram_t* p : { &h0, &h1 }) {
    out.store(*p);
  }

  out.close();
  NOTICE("JHistHDE done. " << endl);
  return 1;
}