/*******************************************************************************
  Markov path reader
 * \file Reads and analyzes photon path ensembles
  
*******************************************************************************/
// C++ standard library
#include<iostream>
#include<sstream>
#include<iomanip>
#include<vector>
#include<cmath>
#include<cstdlib>

// ROOT
#include "TRandom3.h"
#include "TFile.h"
#include "TPolyLine3D.h"
#include "TPolyMarker3D.h"
#include "TAxis3D.h"
#include "TView.h"
#include "TView3D.h"
#include "TCanvas.h"
#include "TPad.h"
#include "TH1F.h"
#include "TH2F.h"

// JPP
#include "Jeep/JParser.hh"
#include "JMarkov/JPhotonPath.hh"
#include "JMarkov/JPhotonPathReader.hh"
#include "JMarkov/JScatteringModel.hh"

// namespaces
using namespace std ;

// this little line removes the make_field macro defined in JParser.hh
// so that we can call the
// JPARSER::make_field(T,string) function ourselves
#undef make_field
using namespace JPP;

int main( int argc, char** argv ) {
  cout << "MarkovPathReader" << endl 
       << "Written by Martijn Jongen" << endl 
       << endl ; 

  vector<string> ifnames ;
  string ifname_model ;
  string ofname ;
  vector<double> Pscat ;

  try {
    JParser<string> zap ; // this argument parser can handle strings    
    zap["f"] = make_field(ifnames,"input file names (binary files containing JPhotonPaths)") ;
    zap["m"] = make_field(ifname_model,"input file name (root file containing JScatteringModel)") ;
    zap["o"] = make_field(ofname,"file name for root output") ;
    zap["P"] = make_field(Pscat,"list of Pscat for each nscat") ;
    // NOTE: multiple files can be passed by calling e.g.
    // ./JMarkovPathreader [other arguments] -f "file1.paths
    // file2.paths
    // file3.paths"

    if (zap.read(argc, argv) != 0) {
      return 1 ;
    }
  }
  catch(const exception &error) {
    cout << error.what() << endl ;
    cout << "Type '" << argv[0] << " -h!' to display the command-line options." << endl ;
    exit(1) ;
  }


  // counter for number of paths of each length
  //map<int,int> npaths ;

  // these will contain the minimal and maximal values of the vertex
  // coordinates
  // they are accessed as e.g. hX[nscat][vertex_index]
  /*vector< vector<double> > Xmin ; // minimum x position of vertices
  vector< vector<double> > Xmax ; // maximum x position of vertices
  vector< vector<double> > Ymin ; // minimum y position of vertices
  vector< vector<double> > Ymax ; // maximum y position of vertices
  vector< vector<double> > Zmin ; // minimum z position of vertices
  vector< vector<double> > Zmax ; // maximum z position of vertices
  vector<double> Lmax ; // maximum path length
  Xmin.resize(nmax) ;
  Xmax.resize(nmax) ;
  Ymin.resize(nmax) ;
  Ymax.resize(nmax) ;
  Zmin.resize(nmax) ;
  Zmax.resize(nmax) ;
  Lmax.resize(nmax) ;
  for( int n=0 ; n<nmax ; ++n ) {
    Xmin[n].resize(n+2) ;
    Xmax[n].resize(n+2) ;
    Ymin[n].resize(n+2) ;
    Ymax[n].resize(n+2) ;
    Zmin[n].resize(n+2) ;
    Zmax[n].resize(n+2) ;
    Lmax[n] = -1.0/0.0 ;
    for( int i=0 ; i<n+2 ; ++i ) {
      Xmin[n][i] = 1.0/0.0 ;
      Ymin[n][i] = 1.0/0.0 ;
      Zmin[n][i] = 1.0/0.0 ;
      Xmax[n][i] = -1.0/0.0 ;
      Ymax[n][i] = -1.0/0.0 ;
      Zmax[n][i] = -1.0/0.0 ;
    }
    }*/

  // read in the scattering model
  JMARKOV::JScatteringModel* sm = NULL ;
  TFile* f = new TFile(ifname_model.c_str(),"read") ;
  if( f->IsZombie() ) {
    cerr << "Could not open file '" << ifname_model << "'." << endl ;
    exit(1) ;
  }
  sm = (JMARKOV::JScatteringModel*) f->Get("JMARKOV::JScatteringModel") ;
  if( sm == NULL ) {
    cerr << "Could not read JScatteringModel from file!" << endl ;
    exit(1) ;
  }
  cout << "JScatteringModel loaded" << endl
       << "  absorption length = " << sm->getAbsorptionLength() << " m" << endl 
       << "  scattering length = " << sm->getScatteringLength() << " m" << endl
       << endl ;

  // read the photon paths into memory
  // we keep track of them in a vector of vectors (they are ordered by nscat)
  vector< vector<JMARKOV::JPhotonPath> > paths ;

  // read the file once to determine the minimal and maximal values
  // of the vertex coordinates
  cout << "Reading files." << endl ;
  for( vector<string>::iterator it=ifnames.begin() ; it!=ifnames.end() ; ++it ) {
    cout << "Reading '" << *it << "'." << endl ;
      // read the paths from the file
    JMARKOV::JPhotonPathReader reader ;
    reader.open(it->c_str()) ;
    if( !reader.is_open() ) {
      cerr << "FATAL ERROR: unable to open input file '" << *it << "'." << endl ;
      exit(1) ;
    }
    int nread = 0 ;
    JMARKOV::JPhotonPath* p = NULL ;
    while( reader.hasNext() ) {
      p = reader.next() ;
      int nscat = p->n-2 ; //p->size()-1 ;
      if( nscat >= (int)paths.size() ) paths.resize(nscat+1) ;
      paths[nscat].push_back(*p) ;
      ++nread ;
    }
    cout << "--> read " << nread << " paths." << endl ;
    reader.close() ;
  }
  cout << endl ;
  cout << "Done reading paths." << endl ;
  cout << endl ;

  const int nscatmax = paths.size() ; // maximal number of scatterings


  Pscat.resize( nscatmax ) ;
  vector<double> weight(nscatmax) ;
  for( int nscat=0 ; nscat<nscatmax ; ++nscat ) {
    weight[nscat] = 0 ;
    if( Pscat[nscat]>0 && paths[nscat].size()>0 ) {
      weight[nscat] = Pscat[nscat] / paths[nscat].size() ;
    }
  }


  cout << setw(5) << "nscat" << setw(15) << "npaths" << setw(15) << "Pscat" << endl ;
  for( int nscat=0; nscat<nscatmax; ++nscat ) {
    if( paths[nscat].size() > 0 ) {
      cout << setw(5)  << nscat 
	   << setw(15) << paths[nscat].size() 
	   << setw(15) << Pscat[nscat] 
	   << endl ;
    }
  }

  cout << "Determining histogram ranges." << endl ;
  vector<JMARKOV::JPhotonPath> posmin(nscatmax) ;
  vector<JMARKOV::JPhotonPath> posmax(nscatmax) ;
  vector<double> Lmin(nscatmax) ; // minimum path length
  vector<double> Lmax(nscatmax) ; // maximum path length
  // initialize to +- infinity
  double INF = 1.0/0.0 ;
  double MININF = -1.0/0.0 ;
  for( int nscat=0 ; nscat<nscatmax ; ++nscat ) {
    posmin[nscat] = JMARKOV::JPhotonPath(nscat) ;
    posmax[nscat] = JMARKOV::JPhotonPath(nscat) ;
    Lmax[nscat] = 0 ;
    Lmin[nscat] = INF ;
    int nvert = nscat + 2 ;
    for( int j=0 ; j<nvert ; ++j ) {
      posmin[nscat][j] = JGEOMETRY3D::JPosition3D(    INF,    INF,    INF ) ;
      posmax[nscat][j] = JGEOMETRY3D::JPosition3D( MININF, MININF, MININF ) ;
    }
  }

  // loop over the paths to find the histogram ranges
  for( int nscat=0 ; nscat<(int)paths.size() ; ++nscat ) {
    // loop over paths with nscat scatterings
    for( vector<JMARKOV::JPhotonPath>::iterator p=paths[nscat].begin() ; p!=paths[nscat].end() ; ++p ) {
      double length = p->getLength() ;
      if( length > Lmax[nscat] ) Lmax[nscat] = length ;
      if( length < Lmin[nscat] ) Lmin[nscat] = length ;
      // loop over the vertices
      int nvert = nscat + 2 ;
      for( int j=0 ; j<nvert ; ++j ) {
	posmin[nscat][j] = JGEOMETRY3D::JPosition3D( min(posmin[nscat][j].getX(),p->at(j).getX()),
						     min(posmin[nscat][j].getY(),p->at(j).getY()), 
						     min(posmin[nscat][j].getZ(),p->at(j).getZ())  ) ;  
	posmax[nscat][j] = JGEOMETRY3D::JPosition3D( max(posmax[nscat][j].getX(),p->at(j).getX()),
						     max(posmax[nscat][j].getY(),p->at(j).getY()), 
						     max(posmax[nscat][j].getZ(),p->at(j).getZ())  ) ;  
      }
    }
  }
  cout << endl ;

  cout << "Allocating histograms." << endl ;
  // create pointers to the histograms
  vector< vector<TH1F*> > hX(nscatmax) ; // x position of vertices
  vector< vector<TH1F*> > hY(nscatmax) ; // y position of vertices
  vector< vector<TH1F*> > hZ(nscatmax) ; // z position of vertices
  vector< vector<TH2F*> > hXY(nscatmax) ; // xy position of vertices
  vector< vector<TH2F*> > hXZ(nscatmax) ; // xz position of vertices
  vector< vector<TH2F*> > hYZ(nscatmax) ; // yz position of vertices
  vector<TH1F*> hL(nscatmax) ; // path length
  vector<TH1F*> hL_zoom(nscatmax) ; // path length
  vector<TH1F*> hCosThetaSource(nscatmax) ; // source angle
  vector<TH1F*> hCosThetaTarget(nscatmax) ; // target angle
  vector<TH2F*> hCosThetaTarget_L(nscatmax) ; // target angle versus path length
  // number of bins 
  int nbins_L = 10000 ;
  int nbins_ct = 10000 ;
  int nbins_X = 1000 ;
  int nbins_Y = 1000 ;
  int nbins_Z = 1000 ;
  // allocate the histograms
  for( int nscat=0 ; nscat<nscatmax ; ++nscat ) {
    if( paths[nscat].size() == 0 ) continue ;

    char hname[200] ;
    char htitle[200] ;

    sprintf( hname, "hL_nscat_%i", nscat ) ;
    sprintf( htitle, "path length distribution of paths with %i scatterings.", nscat ) ;
    hL[nscat] = new TH1F(hname,htitle,nbins_L,Lmin[nscat],1.01*Lmax[nscat]) ;

    sprintf( hname, "hL_zoom_nscat_%i", nscat ) ;
    sprintf( htitle, "path length distribution of paths with %i scatterings.", nscat ) ;
    hL_zoom[nscat] = new TH1F(hname,htitle,nbins_L,Lmin[nscat],Lmin[nscat]+10) ;
    
    double ctmax = 1+1e-8 ;
    sprintf( hname, "hCosThetaSource_nscat_%i", nscat ) ;
    sprintf( htitle, "Source angle with %i scatterings;cos(#theta_{S});a.u.", nscat ) ;
    hCosThetaSource[nscat] = new TH1F(hname,htitle,nbins_ct,-ctmax,ctmax) ;

    sprintf( hname, "hCosThetaTarget_nscat_%i", nscat ) ;
    sprintf( htitle, "Target angle with %i scatterings;cos(#theta_{T});a.u.", nscat ) ;
    hCosThetaTarget[nscat] = new TH1F(hname,htitle,nbins_ct,-ctmax,ctmax) ;

    sprintf( hname, "hCosThetaTarget_L_nscat_%i", nscat ) ;
    sprintf( htitle, "%i scatterings;cos(#theta_{T});path length L", nscat ) ;
    hCosThetaTarget_L[nscat] = new TH2F(hname,htitle,20,-ctmax,ctmax,200,Lmin[nscat],1.01*Lmax[nscat]) ;
    hCosThetaTarget_L[nscat]->SetOption("colz") ;

    int nvert = nscat + 2 ;
    hX[nscat].resize(nvert) ;
    hY[nscat].resize(nvert) ;
    hZ[nscat].resize(nvert) ;
    hXY[nscat].resize(nvert) ;
    hXZ[nscat].resize(nvert) ;
    hYZ[nscat].resize(nvert) ;
    for( int j=0 ; j<nvert; ++j ) {
      sprintf(hname, "hX_nscat_%i_vertex_%i", nscat, j ) ;
      sprintf(htitle, "X of vertex %i in paths with %i scatterings.", j, nscat ) ;
      hX[nscat][j] = new TH1F(hname,htitle,nbins_X,1.01*posmin[nscat][j].getX(),1.01*posmax[nscat][j].getX()) ;

      sprintf(hname, "hY_nscat_%i_vertex_%i", nscat, j ) ;
      sprintf(htitle, "Y of vertex %i in paths with %i scatterings.", j, nscat ) ;
      hY[nscat][j] = new TH1F(hname,htitle,nbins_Y,1.01*posmin[nscat][j].getY(),1.01*posmax[nscat][j].getY()) ;

      sprintf(hname, "hZ_nscat_%i_vertex_%i", nscat, j ) ;
      sprintf(htitle, "Z of vertex %i in paths with %i scatterings.", j, nscat ) ;
      hZ[nscat][j] = new TH1F(hname,htitle,nbins_Z,1.01*posmin[nscat][j].getZ(),1.01*posmax[nscat][j].getZ()) ;

      sprintf(hname, "hXY_nscat_%i_vertex_%i", nscat, j ) ;
      sprintf(htitle, "XY of vertex %i in paths with %i scatterings.", j, nscat ) ;
      hXY[nscat][j] = new TH2F(hname,htitle,
			       nbins_X,1.01*posmin[nscat][j].getX(),1.01*posmax[nscat][j].getX(),
			       nbins_Y,1.01*posmin[nscat][j].getY(),1.01*posmax[nscat][j].getY() ) ;
      hXY[nscat][j]->SetOption("colz") ;

      sprintf(hname, "hXZ_nscat_%i_vertex_%i", nscat, j ) ;
      sprintf(htitle, "XZ of vertex %i in paths with %i scatterings.", j, nscat ) ;
      hXZ[nscat][j] = new TH2F(hname,htitle,
			       nbins_X,1.01*posmin[nscat][j].getX(),1.01*posmax[nscat][j].getX(),
			       nbins_Z,1.01*posmin[nscat][j].getZ(),1.01*posmax[nscat][j].getZ() ) ;
      hXZ[nscat][j]->SetOption("colz") ;

      sprintf(hname, "hYZ_nscat_%i_vertex_%i", nscat, j ) ;
      sprintf(htitle, "YZ of vertex %i in paths with %i scatterings.", j, nscat ) ;
      hYZ[nscat][j] = new TH2F(hname,htitle,
			       nbins_Y,1.01*posmin[nscat][j].getY(),1.01*posmax[nscat][j].getY(),
			       nbins_Z,1.01*posmin[nscat][j].getZ(),1.01*posmax[nscat][j].getZ() ) ;
      hYZ[nscat][j]->SetOption("colz") ;
    }
  }
  cout << endl ;

  // loop over the paths again, filling the histograms
  cout << "Filling histograms." << endl ;
  for( int nscat=0 ; nscat<(int)paths.size() ; ++nscat ) {
    int nvert = nscat + 2 ;
    // loop over paths with nscat scatterings
    for( vector<JMARKOV::JPhotonPath>::iterator p=paths[nscat].begin() ; p!=paths[nscat].end() ; ++p ) {
      double L = p->getLength() ;
      hL[nscat]->Fill(L) ;
      hL_zoom[nscat]->Fill(L) ;
      // source angle
      {
	JGEOMETRY3D::JVersor3D seg(p->at(1)-p->at(0)) ;
	hCosThetaSource[nscat]->Fill(seg.getDZ()) ;
      }
      // target angle
      {
	JGEOMETRY3D::JVersor3D seg(p->at(nvert-1)-p->at(nvert-2)) ;
	hCosThetaTarget[nscat]->Fill(seg.getDZ()) ;
	hCosThetaTarget_L[nscat]->Fill(seg.getDZ(),L) ;
      }
      // loop over vertices
      for( int j=0; j<nvert; ++j ) {
	double x = p->at(j).getX() ;
	double y = p->at(j).getY() ;
	double z = p->at(j).getZ() ;
	hX[nscat][j]->Fill(x) ;
	hY[nscat][j]->Fill(y) ;
	hZ[nscat][j]->Fill(z) ;
	hXY[nscat][j]->Fill(x,y) ;
	hXZ[nscat][j]->Fill(x,z) ;
	hYZ[nscat][j]->Fill(y,z) ;
      }
    }
  }
  cout << endl ;

  // normalize the histograms to 1
  cout << "Normalizing histograms." << endl ;
  for( int nscat=0 ; nscat<nscatmax ; ++nscat ) {
    if( paths[nscat].size() == 0 ) continue ;
    int nvert = nscat + 2 ;
    hL[nscat]->Scale( weight[nscat] ) ;
    hL_zoom[nscat]->Scale( weight[nscat] ) ;
    hCosThetaSource[nscat]->Scale( weight[nscat] ) ;
    hCosThetaTarget[nscat]->Scale( weight[nscat] ) ;
    for( int j=0 ; j<nvert; ++j ) {
      hX[nscat][j]->Scale( weight[nscat] ) ;
      hY[nscat][j]->Scale( weight[nscat] ) ;
      hZ[nscat][j]->Scale( weight[nscat] ) ;
      hXY[nscat][j]->Scale( weight[nscat] ) ;
      hXZ[nscat][j]->Scale( weight[nscat] ) ;
      hYZ[nscat][j]->Scale( weight[nscat] ) ;
    }
  }

  // write results to a file
  cout << "Writing results to '" << ofname << "'." << endl ;
  TFile* fout = new TFile(ofname.c_str(),"recreate") ;
  for( int nscat=0; nscat<nscatmax; ++nscat ) {
    if( paths[nscat].size() == 0 ) continue ;
    int nvert = nscat + 2 ;

    char dname[200] ;
    sprintf( dname, "nscat%i", nscat ) ;
    gDirectory->mkdir(dname)->cd() ;
    hL[nscat]->Write() ;
    hL_zoom[nscat]->Write() ;
    hCosThetaSource[nscat]->Write() ;
    hCosThetaTarget[nscat]->Write() ;
    hCosThetaTarget_L[nscat]->Write() ;
    for( int j=0; j<nvert ; ++j ) {
      hX[nscat][j]->Write() ;
      hY[nscat][j]->Write() ;
      hZ[nscat][j]->Write() ;
      hXY[nscat][j]->Write() ;
      hXZ[nscat][j]->Write() ;
      hYZ[nscat][j]->Write() ;
    }
    fout->cd() ;
  }
  fout->Close() ;
  delete fout ;

  cout << "Output written to '" << ofname << "'." << endl ;

  cout << "Done!" << endl ;

  return 0 ;
}