#define PIPLUS 0
#define PIMINUS 1
#define HADRONS 2
#define QUARKS 6

#include <sstream>
#include <string>
//#include "FourVector.hh"


#ifndef __CINT__
#error You must be trying to compile :(....
#include "TROOT.h"
#include "TTree.h"
#include "TFile.h"
#include "TObject.h"
#include "TH1.h"
#include "TH2.h"
#endif //__CINT__

char * hadsText[HADRONS]={"#pi^{+}","#pi^{-}"};
char * quarksText[QUARKS]={"d","u","s","#bar{d}","#bar{u}","#bar{s}"};

Int_t check_fiducial(Int_t , Float_t , Float_t );
Float_t recalculate_z(Int_t , Float_t, Float_t );

Int_t yield_x_z(char * , char * ,
		Int_t richUnfold=1, Int_t data=1,
		Int_t zfixedBins=0, Int_t zBins=6,
		Float_t w2Min=10, Float_t mxMin=0.0,
		Float_t ptMax=30.0, Int_t fiducialCut=2, 
		Int_t xfixedBins=1, Int_t xBins=7);

Int_t yield_x_z(char * InFileRoot, char * Outfile,
		Int_t richUnfold, Int_t data,
		Int_t zfixedBins, Int_t zBins, 
		Float_t w2Min, Float_t mxMin, 
		Float_t ptMax, Int_t fiducialCut, 
		Int_t xfixedBins, Int_t xBins)
{
  Float_t q2,x,w2, ethx, ethy, pt;
  Float_t pdf[6],qp[3];
  Float_t thx, thy;
  Float_t z,mx, xf;
  Long64_t events,tracks;
  Int_t tr_event,tr_run, qflav, itype;
  Int_t ev_event,ev_run;

  string fileName;
  ostringstream outs;
  cout << "==>>>  w2min: " << w2Min << " mxmin:" ;
  cout << mxMin << " ptmax:" << ptMax << endl;

  //This is the code in chain mode == multiple root trees...
  //Can read one root file or many very nice ...
 //   TChain * T = new TChain("trackTree");
//   TChain * E = new TChain("eventTree");
  
//   string pattern;
//   outs << InFileRoot << "*.root" << ends;
//   pattern=outs.str();
//   outs.seekp(0);
//   cout << pattern << endl;
//   E->Add(pattern.c_str());
//   T->Add(pattern.c_str());
  TFile * file = new TFile(InFileRoot);
  TTree * T = file->Get("trackTree");
  TTree * E = file->Get("eventTree");

  T->SetBranchAddress("ievent", &tr_event);
  T->SetBranchAddress("irun", &tr_run);
  T->SetBranchAddress("iType", &itype);
  T->SetBranchAddress("z", &z);
  T->SetBranchAddress("M_X",&mx);
  T->SetBranchAddress("pT",&pt);
  T->SetBranchAddress("xf",&xf);
  if (data)
    {  
      Float_t pweight[3];
      T->SetBranchAddress("pMatrixWeight",pweight);
      Int_t csb;
      E->SetBranchAddress("CSB",&csb);
    }

  E->SetBranchAddress("ievent",&ev_event);
  E->SetBranchAddress("irun",&ev_run);
  E->SetBranchAddress("qflav",&qflav);
  E->SetBranchAddress("Q2",&q2);
  E->SetBranchAddress("W2",&w2);
  E->SetBranchAddress("x",&x);
  E->SetBranchAddress("qp",qp);

  if ( !data )
    {
      Float_t eventWeight;  
      E->SetBranchAddress("weight",&eventWeight);
    }
  
  


  switch ( fiducialCut )
    {
    case 1:
      Float_t eventFidData[2];
      E->SetBranchAddress("thx",&eventFidData[0]);
      E->SetBranchAddress("thy",&eventFidData[1]);
      Float_t trackFidData[2];
      T->SetBranchAddress("thx",&trackFidData[0]);
      T->SetBranchAddress("thy",&trackFidData[1]);
      break;
    case 2:
      Float_t eventFidData[7];
      E->SetBranchAddress("calX",&eventFidData[0]);
      E->SetBranchAddress("calY",&eventFidData[1]);
      E->SetBranchAddress("frontClampX",&eventFidData[2]);
      E->SetBranchAddress("frontClampY",&eventFidData[3]);
      E->SetBranchAddress("rearClampX",&eventFidData[4]);
      E->SetBranchAddress("rearClampY",&eventFidData[5]);
      E->SetBranchAddress("rearClampYother",&eventFidData[6]);
      Float_t trackFidData[7];
      T->SetBranchAddress("calX",&trackFidData[0]);
      T->SetBranchAddress("calY",&trackFidData[1]);
      T->SetBranchAddress("frontClampX",&trackFidData[2]);
      T->SetBranchAddress("frontClampY",&trackFidData[3]);
      T->SetBranchAddress("rearClampX",&trackFidData[4]);
      T->SetBranchAddress("rearClampY",&trackFidData[5]);
      T->SetBranchAddress("rearClampYother",&trackFidData[6]);
      break;
    default:
      Float_t * eventFidData;
      Float_t * trackFidData;
      break;
    }

#ifdef __CINT__
  //  gSystem->AddIncludePath("-I$HOME/cvsworking/rdepm/sripts");
  gROOT->ProcessLine(".L ..//scripts/fiducial_check.C");
#endif

  //  ofstream fout(Outfile, ios::trunc);
  TFile * outFile=new TFile(Outfile,"RECREATE");
  outFile->SetCompressionLevel(9);
  TTree * yieldData = new TTree("yieldData","Pion yield from target");

  Float_t xpproot,q2pproot,zpproot,npproot,dnpproot,pdfpproot[6];
  yieldData->Branch("xpp",&xpproot,"xpp/F");
  yieldData->Branch("q2pp",&q2pproot,"q2pp/F");
  yieldData->Branch("zpp",&zpproot,"zpp/F");
  yieldData->Branch("npp",&npproot,"npp/F");
  yieldData->Branch("dnpp",&dnpproot,"dnpp/F");
  yieldData->Branch("pdfpp",pdfpproot,"pdfpp[6]/F");

  Float_t xpmroot,q2pmroot,zpmroot,npmroot,dnpmroot,pdfpmroot[6];
  yieldData->Branch("xpm",&xpmroot,"xpm/F");
  yieldData->Branch("q2pm",&q2pmroot,"q2pm/F");
  yieldData->Branch("zpm",&zpmroot,"zpm/F");
  yieldData->Branch("npm",&npmroot,"npm/F");
  yieldData->Branch("dnpm",&dnpmroot,"dnpm/F");
  yieldData->Branch("pdfpm",pdfpmroot,"pdfpm[6]/F");

  Int_t xbin, zbin;
  yieldData->Branch("xbin",&xbin,"xbin/I");
  yieldData->Branch("zbin",&zbin,"zbin/I");

  Float_t * xBinEdges =new Float_t [xBins+1];
  Float_t * zBinEdges = new Float_t[zBins+1];

#ifdef __CINT__
  gROOT->ProcessLine(".L ../scripts/binning.C");
#endif // __CINT__

  binning(xBins, xfixedBins, zBins, zfixedBins, xBinEdges, zBinEdges);

  //All the arrays to hold our data
  TMatrixF * xAvg = new TMatrixF(HADRONS,xBins*zBins);
  TMatrixF * q2Avg = new TMatrixF(HADRONS,xBins*zBins);
  TMatrixF * zAvg = new TMatrixF(HADRONS,xBins*zBins);
  TMatrixF * semiInclusiveYield= new TMatrixF(HADRONS,xBins*zBins);
  TMatrixF * semiInclusiveYield2= new TMatrixF(HADRONS,xBins*zBins);
  TMatrixF * semiInclusiveYieldError= new TMatrixF(HADRONS,xBins*zBins);

//   FourVector * xAvg= new FourVector(HADRONS,xBins,zBins);
//   FourVector * q2Avg= new FourVector(HADRONS,xBins,zBins);
//   FourVector * zAvg= new FourVector(HADRONS,xBins,zBins);
//   FourVector * semiInclusiveYield= new FourVector(HADRONS,xBins,zBins);
//   FourVector * semiInclusiveYield2= new FourVector(HADRONS,xBins,zBins);
//   FourVector * semiInclusiveYieldError= new FourVector(HADRONS,xBins,zBins);

  //Build an index with run and event so we 
  //can navigate from track to event tree...
  E->BuildIndex("irun","ievent");
  
  tracks=T->GetEntries();
  cout << "total tracks=" << tracks << endl;
  //Loop over all the tracks....
  //  for (Int_t i=0; i < tracks; i++)
  for (Int_t i=0; i < 1000; i++)
    {
      T->GetEntry(i);
      if (/////////////////////////////
	  // These need to be added to the smearing matrix as well still
	  //pt < ptMax &&
	  //////////////////////////////
	  mx > mxMin &&
	  z >= zBinEdges[0] && 
	  z <= zBinEdges[zBins] &&
	  xf > 0.1  &&
	  fiducial_check( trackFidData, fiducialCut ) )
	{
	  //Get the event corresponding to this track
	  //cout << "Fetching event " << tr_event << endl;
	  E->GetEntryWithIndex(tr_run,tr_event);
	  if ( w2 > w2Min &&
	       x >= xBinEdges[0] &&
	       x <= xBinEdges[xBins] &&
	       q2 > 1.0 &&
	       fiducial_check(eventFidData, fiducialCut) )
	    {
	      for(Int_t ix=0; ix < xBins; ix++)
		{
		  if ( x >= xBinEdges[ix] && 
		       x <= xBinEdges[ix+1] )
		    {
		      // Here we need to recalculate the z 
		      // value if this is a kaon or proton 
		      // if this is rich unfolding
		      if ( richUnfold )
			z=recalculate_z(itype,z,qp[0]);
		      
		      for(Int_t iz=0; iz < zBins; iz++)
			{
			  if ( z >= zBinEdges[iz] &&
			       z <= zBinEdges[iz+1] )
			    {
			      ////////////////////////////////////
			      // DATA path
			      ///////////////////////////////////
			      //count everybody if unfolding
			      if ( data )
				{
				  if ( richUnfold )
				    {
				      if ( 5 >= itype && 
					   itype >= 3)
					{
					  
					  semiInclusiveYield(PIPLUS, ix*zBins+iz)+=csb*pweight[abs(itype)-3];
					  semiInclusiveYield2(PIPLUS, ix*zBins+iz)+=pow( pweight[abs(itype)-3], 2 );
					  xAvg(PIPLUS, ix*zBins+iz)+=x*pweight[abs(itype)-3];
					  zAvg(PIPLUS, ix*zBins+iz)+=z*pweight[abs(itype)-3];
					  q2Avg(PIPLUS, ix*zBins+iz)+=q2*pweight[abs(itype)-3];

// 					  semiInclusiveYield->AddToValue( PIPLUS, ix, iz, 
// 									  csb*pweight[fabs(itype)-3]);
// 					  semiInclusiveYield2->AddToValue( PIPLUS, ix, iz, 
// 									   pow( pweight[fabs(itype)-3], 2 ) );
// 					  xAvg->AddToValue( PIPLUS, ix, iz, x*pweight[fabs(itype)-3] );
// 					  zAvg->AddToValue( PIPLUS, ix, iz, z*pweight[fabs(itype)-3] );
// 					  q2Avg->AddToValue( PIPLUS, ix, iz, q2*pweight[fabs(itype)-3] );
					  
					}
				      else if( -3 >= itype && 
					       itype >= -5)
					{

					  semiInclusiveYield(PIMINUS, ix*zBins+iz)+=csb*pweight[abs(itype)-3];
					  semiInclusiveYield2(PIMINUS, ix*zBins+iz)+=pow( pweight[abs(itype)-3], 2 );
					  xAvg(PIMINUS, ix*zBins+iz)+=x*pweight[abs(itype)-3];
					  zAvg(PIMINUS, ix*zBins+iz)+=z*pweight[abs(itype)-3];
					  q2Avg(PIMINUS, ix*zBins+iz)+=q2*pweight[abs(itype)-3];

// 					  semiInclusiveYield->AddToValue( PIMINUS, ix, iz, 
// 									  csb*pweight[fabs(itype)-3]);
// 					  semiInclusiveYield2->AddToValue( PIMINUS, ix, iz, 
// 									   pow( pweight[fabs(itype)-3], 2 ) );
// 					  xAvg->AddToValue( PIMINUS, ix, iz, x*pweight[fabs(itype)-3] );
// 					  zAvg->AddToValue( PIMINUS, ix, iz, z*pweight[fabs(itype)-3] );
// 					  q2Avg->AddToValue( PIMINUS, ix, iz, q2*pweight[fabs(itype)-3] );

					}
				    }
				  else //Only count pions and no weight
				    {
				      switch (itype)
					{
					case 3:
					  
					  semiInclusiveYield(PIPLUS, ix*zBins+iz)+=csb;
					  semiInclusiveYield2(PIPLUS, ix*zBins+iz)+=1;
					  xAvg(PIPLUS, ix*zBins+iz)+=x;
					  zAvg(PIPLUS, ix*zBins+iz)+=z;
					  q2Avg(PIPLUS, ix*zBins+iz)+=q2;					
  
// 					  semiInclusiveYield->AddToValue( PIPLUS, ix, iz, csb);
// 					  semiInclusiveYield2->AddToValue( PIPLUS, ix, iz, 1);
// 					  xAvg->AddToValue( PIPLUS, ix, iz, x );
// 					  zAvg->AddToValue( PIPLUS, ix, iz, z );
// 					  q2Avg->AddToValue( PIPLUS, ix, iz, q2 );
					  break;
					case -3:

					  semiInclusiveYield(PIMINUS, ix*zBins+iz)+=csb;
					  semiInclusiveYield2(PIMINUS, ix*zBins+iz)+=1;
					  xAvg(PIMINUS, ix*zBins+iz)+=x;
					  zAvg(PIMINUS, ix*zBins+iz)+=z;
					  q2Avg(PIMINUS, ix*zBins+iz)+=q2;
					  
// 					  semiInclusiveYield->AddToValue( PIMINUS, ix, iz, csb);
// 					  semiInclusiveYield2->AddToValue( PIMINUS, ix, iz, 1);
// 					  xAvg->AddToValue( PIMINUS, ix, iz, x );
// 					  zAvg->AddToValue( PIMINUS, ix, iz, z );
// 					  q2Avg->AddToValue( PIMINUS, ix, iz, q2 );
					  break;
					default:
					  break;
					}
				    }
				}
			      ////////////////////////////////////
			      // MC path
			      ///////////////////////////////////
			      else 
				{
				  switch (itype)
				    {
				    case 3:

				      semiInclusiveYield(PIPLUS, ix*zBins+iz)+=eventWeight;
				      semiInclusiveYield2(PIPLUS, ix*zBins+iz)+=pow( eventWeight, 2 );
				      xAvg(PIPLUS, ix*zBins+iz)+=x*eventWeight;
				      zAvg(PIPLUS, ix*zBins+iz)+=z*eventWeight;
				      q2Avg(PIPLUS, ix*zBins+iz)+=q2*eventWeight;

// 				      semiInclusiveYield->AddToValue( PIPLUS, ix, iz, 
// 								      eventWeight);
// 				      semiInclusiveYield2->AddToValue( PIPLUS, ix, iz, 
// 								       pow( eventWeight, 2 ) );
// 				      xAvg->AddToValue( PIPLUS, ix, iz, x*eventWeight );
// 				      zAvg->AddToValue( PIPLUS, ix, iz, z*eventWeight );
// 				      q2Avg->AddToValue( PIPLUS, ix, iz, q2*eventWeight );
				      break;
				    case -3:

				      semiInclusiveYield(PIMINUS, ix*zBins+iz)+=eventWeight;
				      semiInclusiveYield2(PIMINUS, ix*zBins+iz)+=pow( eventWeight, 2 );
				      xAvg(PIMINUS, ix*zBins+iz)+=x*eventWeight;
				      zAvg(PIMINUS, ix*zBins+iz)+=z*eventWeight;
				      q2Avg(PIMINUS, ix*zBins+iz)+=q2*eventWeight;
				      
//  				      semiInclusiveYield->AddToValue( PIMINUS, ix, iz, 
// 								      eventWeight);
// 				      semiInclusiveYield2->AddToValue( PIMINUS, ix, iz, 
// 								       pow( eventWeight, 2 ) );
// 				      xAvg->AddToValue( PIMINUS, ix, iz, x*eventWeight );
// 				      zAvg->AddToValue( PIMINUS, ix, iz, z*eventWeight );
// 				      q2Avg->AddToValue( PIMINUS, ix, iz, q2*eventWeight );
				      break;
				    default:
				      break;
				    }
				}
			      
			      break;
			    } // zbin test
			} // loop over zbins
		      break;
		    } //xbin test
		} //loop over xbins
	    }//inclusive cuts
	}//semi-inclusive cuts
    }//track loop
  
  //Calculate averages
  for(Int_t ix=0; ix < xBins; ix++)
    {
      for(Int_t iz=0; iz < zBins; iz++)
	{
	  for (Int_t k=0;  k < HADRONS; k++)
	    {
	      
	      if( semiInclusiveYield(k,ix*zBins+iz) )
		{
 		  xAvg(k, ix*zBins+iz)=xAvg(k, ix*zBins+iz)/semiInclusiveYield(k,ix*zBins+iz);
 		  zAvg(k, ix*zBins+iz)=zAvg(k, ix*zBins+iz)/semiInclusiveYield(k,ix*zBins+iz);
 		  q2Avg(k, ix*zBins+iz)=q2Avg(k, ix*zBins+iz)/semiInclusiveYield(k,ix*zBins+iz);
		}
	      else
		{
 		  xAvg(k, ix*zBins+iz)=( xBinEdges[ix+1] - xBinEdges[ix] ) / 2 + xBinEdges[ix];
		  zAvg(k, ix*zBins+iz)=( zBinEdges[iz+1] - zBinEdges[iz] ) / 2 + zBinEdges[iz];
 		  q2Avg(k, ix*zBins+iz)=0;
		}
	      
	      semiInclusiveYieldError(k,ix*zBins+iz)=sqrt( semiInclusiveYield2(k,ix*zBins+iz) );
	      
	      if ( k == PIPLUS )
		{
		  xpproot=xAvg(k, ix*zBins+iz);
		  q2pproot=q2Avg(k, ix*zBins+iz);
		  zpproot=zAvg(k, ix*zBins+iz);
		  npproot=semiInclusiveYield(k, ix*zBins+iz);
		  dnpproot=semiInclusiveYieldError(k, ix*zBins+iz);
		}
	      else
		{
		  xpmroot=xAvg(k, ix*zBins+iz);
		  q2pmroot=q2Avg(k, ix*zBins+iz);
		  zpmroot=zAvg(k, ix*zBins+iz);
		  npmroot=semiInclusiveYield(k, ix*zBins+iz);
		  dnpmroot=semiInclusiveYieldError(k, ix*zBins+iz);
		}
	      
// 	      if( semiInclusiveYield(k,ix,iz) )
// 		{

// 		  xAvg->SetValue(k,ix,iz,xAvg(k,ix,iz)/semiInclusiveYield(k,ix,iz));
// 		  zAvg->SetValue(k,ix,iz,zAvg(k,ix,iz)/semiInclusiveYield(k,ix,iz));
// 		  q2Avg->SetValue(k,ix,iz,q2Avg(k,ix,iz)/semiInclusiveYield(k,ix,iz));

// 		}
// 	      else
// 		{
// 		  xAvg->SetValue(k,ix,iz, ( xBinEdges[ix+1] - xBinEdges[ix] ) / 2 + xBinEdges[ix]);
// 		  zAvg->SetValue(k,iz,iz, ( zBinEdges[iz+1] - zBinEdges[iz] ) / 2 + zBinEdges[iz]);
// 		  q2Avg->SetValue(k,ix,iz,0);
// 		}
	      
// 	      semiInclusiveYieldError->SetValue(k,ix,iz, sqrt( semiInclusiveYield2(k,ix,iz) ) );
	      
// 	      if ( k == PIPLUS )
// 		{
// 		  xpproot=xAvg(k, ix, iz);
// 		  q2pproot=q2Avg(k, ix, iz);
// 		  zpproot=zAvg(k, ix, iz);
// 		  npproot=semiInclusiveYield(k, ix, iz);
// 		  dnpproot=semiInclusiveYieldError(k, ix,  iz);
// 		}
// 	      else
// 		{
// 		  xpmroot=xAvg(k, ix, iz);
// 		  q2pmroot=q2Avg(k, ix, iz);
// 		  zpmroot=zAvg(k, ix, iz);
// 		  npmroot=semiInclusiveYield->GetValue(k,ix,iz);
// 		  dnpmroot=semiInclusiveYieldError(k, ix, iz);
// 		}

	      //Write to the file...
	      cout << xAvg(k, ix*zBins+iz) << " " << zAvg(k, ix*zBins+iz) << " " << q2Avg(k, ix*zBins+iz) << " ";
	      cout << semiInclusiveYield(k,ix*zBins+iz) << " ";
	      cout << semiInclusiveYieldError(k, ix*zBins+iz) << " ";
	    } // hadron Loop
	  //Fill up the data
	  xbin=ix;
	  zbin=iz;
	  yieldData->Fill();			      
	}// z loop
    }// x loop
  cout << endl;
  
  delete [] xBinEdges;
  delete [] zBinEdges;
  delete semiInclusiveYield;
  delete semiInclusiveYield2;
  delete semiInclusiveYieldError;
  delete xAvg;
  delete zAvg;
  delete q2Avg;

  outFile->Write();
  outFile->Close();

  cout << "The End" << endl;

  return 0;

}


Float_t recalculate_z(Int_t type, Float_t z_orig, Float_t nu)
{
  Float_t z_new;
  Float_t dM;
  Float_t M[3] = { 0.13957018, //pion
		   0.49367000, //kaon
		   0.93827200 }; //proton

  dM=pow(M[abs(type)-3],2)-pow(M[0],2);

  z_new=sqrt(pow(z_orig,2) + dM / pow(nu,2) );

  return z_new;
}