#include "Spectrometer.h"
#include <vector>
#include <iostream>
#include <math.h>
#include "Hit.h"
using namespace std;

Spectrometer::Spectrometer():  //constructor default values
                               //RUN
                               numberOfEvents(5000),
                               beamMomentum(0.05),
                               thetaxz(0.0),
                               //GEOMETRY
                               spectrometerLength(30.),
                               integralBdL(0.5),
                               magLength(10.),
                               numberOfPlanes(6),
                               distBetweenPlanes(10.),
                               pixelSize(0.002),
                               resolution(0.0006),
                               ySize(2.),
                               zSize(1.),
                               tailAmplitude(0.1),
                               tailWidth(0.0018),
                               multScattAngle(0.00002),
                               //NOISE
                               noiseOccupancy(0.00001),
                               hitEfficiency(0.97),
                               Cut1(16.),
                               Cut2(16.),
                               Cut3(25.),
                               // STAT
                               numberOfReconstructedTracks(0),
                               nTotalHits(0),
                               nNoiseHitsOnTrack(0),
                               numberOfInefficientTracks(0),
                               numberOfRejectedTracks(0),
                               debug(true)

  {

    // Overwrite data members here to remember the default.
    // But then these data members must not be declared const.

    //spectrometerLength=60.;
    //beamMomentum=1.;
    //noiseOccupancy=0.0001;
    //numberOfPlanes=4;

    // Some data members are derived from others:
    distBetweenPlanes= spectrometerLength/(numberOfPlanes-1.);

    // The hit efficiency depends on the noise occupancy:
    hitEfficiency=1.-0.06/fmaxf(log10(noiseOccupancy)+7.0,0.06);

    //output file for histograms
    outfile = new TFile("spectrometer.root", "RECREATE");

    // Book histograms
    h1 = new TH1F("h1","y0 residuals",100,-.005,.005);
    h2 = new TH1F("h2","z0 residuals",100,-.005,.005);
    h3 = new TH1F("h3","ty residuals",100,-.01,.01);
    h4 = new TH1F("h4","tz residuals",100,-.01,.01);
    h5 = new TH1F("h5","1/p residuals",100,-0.2/beamMomentum,0.2/beamMomentum);
    h6 = new TH1F("h6","y0 pull",100,-10.,10.);
    h7 = new TH1F("h7","z0 pull",100,-10.,10.);
    h8 = new TH1F("h8"," ty pull ",100,-10.,10.);
    h9 = new TH1F("h9"," tz pull ",100,-10.,10.);
    h10 = new TH1F("h10"," 1/p pull ",100,-10.,10.);
    h11 = new TH1F("h11"," z chisquared at plane 3 ",80,0.,24.);
    h12 = new TH1F("h12"," total chisquared ",80,0.,24.);

}
TruthTrack Spectrometer::SimulateEvent() {
    //==========================================================================================
    // Simulate the event
    //==========================================================================================

    TruthTrack t(0.,thetaxz,0.,0.,1./beamMomentum);

    // Clear Hits
    hitdata.clear();

    // propagate track through the planes before the magnet
    PropagateStraight(0,t);
    float* par=t.GetPar();
    //par[0] =z0
    //par[1] =dz/dx
    //par[2] =y0
    //par[3] =dy/x

    //Trace through magnet (to a good approximation),
    float dtheta=0.003*integralBdL/beamMomentum;// bend angle in the magnet
    par[2]+=(par[3]*distBetweenPlanes/2.);      // move in y to the magnet center
    float ang=atan(par[3])+dtheta;
    par[2]+=(tan(ang)*distBetweenPlanes/2.);    // and then on to the next plane
    par[3]=tan(ang);                            // update angle
    par[0]+=(distBetweenPlanes*par[1]);         // in z all the way through magnet
    t.SetPar(par);

    //Propagate track through the rest of the planes after the magnet
    PropagateStraight(numberOfPlanes/2,t);

    return t;
}

void Spectrometer::PropagateStraight(int firstplane, TruthTrack& t) { 
    //
    // propagate from plane to plane in a field free region
    //-----------------------------------------------------------------------

    float* par=t.GetPar();
    int nh=(int)hitdata.size();

    for(int j=firstplane; j<firstplane+numberOfPlanes/2; j++) {    

      if(debug) cout << " propagating track. Now at plane " << j << endl;
      numHit[j]=0;
      float y,z,r;

      if(j==firstplane) {
        y=par[2];                            // track impact in first plane
        z=par[0];                         
      }else{
      //extrapolate track to next plane
        par[2]+=(distBetweenPlanes*par[3]);
        par[0]+=(distBetweenPlanes*par[1]);
        y=par[2];                            // track impact in second plane
        z=par[0];                         
      }
       //add multiple scattering
        r=rdm.Gaus()/beamMomentum;
        par[3]+=(r*multScattAngle);           // update angle due to MS
        r=rdm.Gaus()/beamMomentum;
        par[1]+=(r*multScattAngle);           // update angle due to MS
        if(debug) cout << " track y z at plane " << j << " "
                       << " " << y << " "
                       << z << " angles " << par[3] << " " << par[1] << endl;

        r=rdm.Uniform();
        bool ishit = false; //has this plane been hit for the first time? 

        // If this gives a hit (allowing for inefficiency and tracks missing the chip)
        if(r<hitEfficiency && fabs(y)<ySize && fabs(z)<zSize) {
          float r1=rdm.Gaus();
          y+=r1*resolution;                    // then smear impact by resolution
          r1=rdm.Uniform();                    // sometimes give extra smear
          if(r1<tailAmplitude) {
            float r2=rdm.Gaus();
            y+=r2*tailWidth;                   // extra smear (resolution tail)
          }
          float r3=rdm.Gaus();
          z+=r3*resolution;                    // smear impact by resolution
          r3=rdm.Uniform();                    // do we need extra smear?
          if(r3<tailAmplitude) {
            float r2=rdm.Gaus();
            z+=r2*tailWidth;                   // extra smear (resolution tail)
          }

          Hit h(z,y,j,false);
          hitdata.push_back(h);       // update vector of hits in data
          if(!ishit) firstHit[j]=nh;  // index of first hit in this plane
          ishit = true;
          t.AddHit(nh);               // update vector of hit indices on track
          if(debug) cout << " hit " << nh << " at plane " << j <<
	     	          " y " << y << " z " << z << " expect y z "
			 << par[2] << " " << par[0] << endl;
          nh++;                       //update total number of hits
          numHit[j]++;                //update number of hits in this plane
      }

      
      // add noise to 500x500 pixel area ( around real hit )
      int noise=rdm.Poisson(250000*noiseOccupancy);
      if(noise>0) {
        for(int k=0;k<noise;k++) {
          r=rdm.Uniform();                  // placement of noise hit
          float ynoise= y+(r-0.5)*(500.*pixelSize);
          r=rdm.Uniform();
          float znoise= z+(r-0.5)*(500.*pixelSize);
          Hit hnoise(znoise,ynoise,j,true);
          hitdata.push_back(hnoise);        // update vector of hits in data
          if(!ishit) firstHit[j]=nh;
          ishit=true;
          nh++;
          numHit[j]++;
        }
      }

    }                                       //end plane simulation loop
}

float Spectrometer::ReconstructEvent() {
    // =======================================================================
    //Reconstruct one track in 4 or 6 planes
    //Hits in the four first planes are mandatory.
    //One missing plane among the two last in the 6 plane setup is allowed. 
    // =======================================================================
  
  float chi2min=10000000.;     // Initialize output
  bestTrack.SetChi2(chi2min);

    if(debug) {
      for(int i=0;i<numberOfPlanes;i++) cout << "  Hits in plane " << i
                                             << " : " << numHit[i] << endl;
      cout << " total hits " << hitdata.size() << endl;
    }
    // There must be hits in the first four planes  
    if(numHit[0]==0 || numHit[1]==0 || numHit[2]==0 || numHit[3]==0) return chi2min;

    // loop over the hits in the first plane
    // =======================================================================
    for( int i0=firstHit[0];i0< firstHit[0]+numHit[0]; i0++) {
 
      // skip hit in first plane if it is outside the beam profile
      if(fabs(hitdata.at(i0).GetY())>10.*pixelSize) continue;


      // loop over the hits in the second plane
      // =====================================================================
      for( int i1=firstHit[1];i1< firstHit[1]+numHit[1]; i1++) {

        double s2=resolution*resolution;

	//consider the xz plane - a non-bending plane
        //track state at plane 1
        TMatrixD z(2,1);
        z[0][0]= hitdata.at(i1).GetZ();
        z[1][0]= (hitdata.at(i1).GetZ()-hitdata.at(i0).GetZ())/distBetweenPlanes ;
	//its covariance
        TMatrixD Cz(2,2);
        Cz[0][0]=s2;
        Cz[0][1]=s2/distBetweenPlanes;
        Cz[1][0]=Cz[0][1];
        Cz[1][1]=2*s2/distBetweenPlanes/distBetweenPlanes;
        //its hits
        vector<int> hits;
        hits.push_back(i0);
        hits.push_back(i1);
        if(debug) cout << "  Seed track hits: " << i0 << " " << i1 << endl;
	//store it all in two separate track copies
        RecoTrack t2(z,Cz,hits);
        RecoTrack t3(z,Cz,hits);

        // loop over hits in the third plane
        //===========================================================
        for( int i2=firstHit[2];i2< firstHit[2]+numHit[2]; i2++) {

	  float chi2 = KalmanFilter(2,i2,t2,t3);  // t3 is the updated track
          if(debug) cout << " Hit in third plane, i2,chi2  " << i2 << " " << chi2 << endl;
 	  if(chi2 > Cut1) continue;
 
          //loop over hits in the fourth plane
          //====================================================================================
          for( int i3=firstHit[3];i3< firstHit[3]+numHit[3]; i3++) {
            RecoTrack t4(t3.GetPar(),t3.GetCov(),t3.GetHits());
	    chi2 = KalmanFilter(3,i3,t3,t4);      //t4 is the updated track
            if(debug) cout << " Hit in fourth plane, i3,chi2  " << i3 << " " << chi2  << endl;
            h11->Fill(chi2);

 	    if(chi2 > Cut1 ) continue;

            if(numberOfPlanes<6) {
	        // now we have a track candidate for the four-plane configuration

	        // make a global chi2 fit and store only the best track
                // ===============================================================================
                TMatrixD x(5,1);
                TMatrixD C(5,5);
                x[4][0] = 1./beamMomentum; //use here a fixed momentum estimate
                hits=t4.GetHits();         // get hit list
                RecoTrack tg(x,C,hits);
                chi2= GlobalChi2(tg);      
                if(chi2<chi2min) {
		  bestTrack.SetPar(tg.GetPar()); 
		  bestTrack.SetCov(tg.GetCov()); 
		  bestTrack.SetHits(tg.GetHits()); 
		  bestTrack.SetChi2(tg.GetChi2()); 
		  bestTrack.SetNdf(tg.GetNdf()); 
                  chi2min=chi2;
		}
                continue;
	    }


            //loop over hits in the fifth plane
            //====================================================================================

            for( int i4=firstHit[4];i4< firstHit[4]+numHit[4]; i4++) {
              RecoTrack t5(t4.GetPar(),t4.GetCov(),t4.GetHits());
	      chi2 = KalmanFilter(4,i4,t4,t5); //t5 is the updated track (only if chi2 < Cut1)
              if(debug) cout << " Hit in fifth plane, i4,chi2  " << i4 << " " << chi2  << endl;
              bool ishit=false;
 	      if(chi2 < Cut1 ) ishit=true; // we allow for a "hole" in one of the two last planes

              //loop over hits in the sixth plane
              //====================================================================================
              for( int i5=firstHit[5];i5< firstHit[5]+numHit[5]; i5++) {
                RecoTrack t6(t5.GetPar(),t5.GetCov(),t5.GetHits());
  	        chi2 = KalmanFilter(5,i5,t5,t6);
                if(debug) cout << " Hit in sixth plane, i5,chi2  " << i5 << " " << chi2  << endl;
 	        if(chi2 > Cut1 && !ishit) continue;

	        // now we have a track candidate for the six-plane configuration

	        // make a global chi2 fit and store only the best track
                // ===============================================================================
                TMatrixD x(5,1);
                TMatrixD C(5,5);
                x[4][0] = 1./beamMomentum; //use here a fixed momentum estimate
                hits=t6.GetHits();         // get hit list
                int nh=(int)hits.size();
                vector<int> usedHits=bestTrack.GetHits();
                // do not doublicate solutions !
                if(nh==(int)usedHits.size() &&
                   (hits[nh-1]==usedHits[nh-1] ||
                    hits[nh-2]==usedHits[nh-2])
                ) continue;
		// choose the best solution
                RecoTrack tg(x,C,hits);
                chi2= GlobalChi2(tg);
                if(debug) cout << " Global fit chi2  " << chi2 << endl;
                if(chi2<chi2min) {
		// store it
		  bestTrack.SetPar(tg.GetPar()); 
		  bestTrack.SetCov(tg.GetCov()); 
		  bestTrack.SetHits(tg.GetHits()); 
		  bestTrack.SetChi2(tg.GetChi2()); 
		  bestTrack.SetNdf(tg.GetNdf()); 
                  chi2min=chi2;
		}
 	      }
	    }
	  }
	}
      }
    }
    return chi2min;
}
float Spectrometer::KalmanFilter(int p1, int ihit, RecoTrack t, RecoTrack& tout) {
  // Propagate a track candidate to detector plane p1
  // as a straight line in x-z (the non-bending plane).
  // Update the track parameters and their error matrix in x-z using hit
  // number ihit in p1. Return the updated track and the chisquared of ihit.

  //the first few lines are here

        if(debug) cout << " in Kalman Filter " << endl; 
        double s2=resolution*resolution;
        double pinv = 1./beamMomentum; //use here a fixed momentum estimate
  	double t0=multScattAngle*pinv; //average multiple scattering angle
        TMatrixD z=t.GetPar();         // current track parameters
        TMatrixD C=t.GetCov();         // current covariance matrix
        vector<int> hits=t.GetHits();  // get hit list
        int plane0= hitdata.at(hits.back()).GetPlane(); //current plane
        double d=distBetweenPlanes*(p1-plane0); // propagation distance
        if(debug) cout << "   propagating from plane " << plane0 <<
          " to " << p1 << " dist " <<  d << endl; 

        TMatrixD Fz(2,2);  //propagator Fz to next plane
        TMatrixD FTz(2,2); //transposed Fz matrix
        Fz[0][0]=1;     
        Fz[0][1]=d;
        Fz[1][0]=0;     
        Fz[1][1]=1;

        FTz=Fz.T();  //the transpose of F
        Fz.T();

        if(debug) cout << "   adding mult scatt " << endl; 
        //multiple scattering contribution to covariance of extrapolation
        TMatrixD Qz(2,2);
        Qz[0][0]=t0*t0*d*d;
        Qz[0][1]=t0*t0*d;
        Qz[1][0]=t0*t0*d;
        Qz[1][1]=t0*t0;

 	//covariance of extrapolation
        TMatrixD Cpz(2,2);
        Cpz = Fz*C*FTz+Qz;

        //predicted state at plane p1
        TMatrixD zpred(2,1);
        zpred = Fz*z;

        //covariance matrix of updated state 
        TMatrixD Cinv(2,2); //inverse covariance matrix of updated state
        TMatrixD Minv(2,2); //inverse covariance matrix of the new measurement
        Minv.Zero();
        Minv[0][0]=1./s2;
        Cinv = Cpz.Invert() + Minv; //add inverted covariances
        C=Cinv.Invert();            //update covariance

        //update track with new measurement
        TMatrixD znew(2,1);       //new measurement weighted by 1/s2
        znew[0][0]=hitdata.at(ihit).GetZ()/s2;
        znew[1][0]=0.;
        z=C*(Cpz*zpred + znew);  //update track params - remember Cpz is inverted
        Cpz.Invert();


        //The residual of the new measurement wrt the prediction 
        double r =(hitdata.at(ihit).GetZ()-zpred[0][0]);
        double R = s2 + Cpz[0][0];
        float chi2=r*r/R;

        if(chi2<Cut1) {            //update track
          hits.push_back(ihit);    //update hit list
          int ndf= hits.size()-2;       
          tout.SetPar(z); 
          tout.SetCov(C);
          tout.SetHits(hits);
          tout.SetChi2(chi2);
          tout.SetNdf(ndf);
	}
 
        return chi2;
}


//**********************************************************************
float Spectrometer::GlobalChi2(RecoTrack& t) {
  //global chi2-fit to x-y hits in a number of pixel planes

  TMatrixD x(5,1);
  x=t.GetPar();
  TMatrixD C(5,5);
  C=t.GetCov();

  float pinv=fabs(x[4][0]);      // we have to estimate the momentum first
                                 // in order to estimate the multiple scattering
  float  dpdt=0.003*integralBdL; // d(momentum)/d(tan(delta_theta))
  double t0=multScattAngle*pinv; // MS angle using an estimated 1/p
  double t2=t0*t0;
  double s2=resolution*resolution;

  vector<int> hits=t.GetHits();       // get hit list
  const int nplane=hits.size();       // number of hit pixel planes
  const int ndim=2*nplane;            // 2 times that (z and y measurements)

  //
  // The column vector of the measurements
  // - first the nplane z measurements then the nplane y measurements 
  TMatrixD m(ndim,1);
  double dist[nplane];
  if(debug) cout << " Global fit to " << nplane << " hits " << endl;
  for (int i=0 ;i<nplane;i++) {
    Hit h=hitdata.at(hits.at(i));
    m[i][0]=h.GetZ();
    m[i+nplane][0]=h.GetY();
    if(i<nplane-1) {
      Hit hnext=hitdata.at(hits.at(i+1));
      int ndis = hnext.GetPlane() - h.GetPlane(); //we have a complication  since we,
      dist[i]=distBetweenPlanes*ndis;    //in principle, allow holes (missing planes)
    } else {
      dist[i]=distBetweenPlanes;
    }
    if(debug) cout << "     Hit " << i << " z " <<  m[i][0] << " y " << m[i+nplane][0]
                   << " dx " << dist[i] << endl;
  }

  //
  // The Covariance matrix of the measurements, including MS induced correlations
  // This is probably not right in the cases of missing planes
  TMatrixD V(ndim,ndim);
  V.Zero();
  V[0][0]=s2;
  V[nplane][nplane]=s2;
  for(int i=1; i<nplane; i++) {
    for(int j=i; j<nplane; j++) {
      V[i][j] = V[i-1][j-1] + i*j*t2*dist[i-1]*dist[j-1];
       V[i+nplane][j+nplane] = V[i][j];

       if(j>i) {
         V[j][i] = V[i][j];
         V[j+nplane][i+nplane] = V[i+nplane][j+nplane];
       }
       if(debug) cout << " V meas " << i << " " << j << " " << V[i][j] << " " << V[i+nplane][j+nplane] << endl;
    }
  }

  //
  // The track state vector is defined as
  // x = (z0, z', y0, y', 1/p) with
  // track impact z0, y0 and slopes z'=dz/dx and y'=dy/dx all given at plane 0
  //
  // H projects the track state vector on the measurement base
 TMatrixD H(ndim,5);
 H.Zero();
  for(int i=0; i<nplane; i++) {
    Hit h=hitdata.at(hits.at(i));
    int plane=h.GetPlane();
    int j=i+nplane;
    H[i][0]=1;
    H[j][2]=1;
    H[i][1]=plane*distBetweenPlanes;
    H[j][3]=plane*distBetweenPlanes;
    if(plane>numberOfPlanes/2-1) H[j][4]=dpdt*distBetweenPlanes*(0.5+plane-numberOfPlanes/2);
    if(debug) cout << " Plane " << plane << " H:  " << H[i][0] << " " << H[i][1] <<
     	 " " << H[j][2] << " " << H[j][3] << " " <<  H[j][4] << endl;
  }

  TMatrixD HT=H.T(); //The transposed
  H.T();             //Do this because H got also transposed.
  //
  //
  C=HT*V.Invert()*H;
  C.Invert();
  // V is now inverse
  // C is now the covariance matrix of the fitted track state
  // x is the fitted track state vector
  x = C*HT*V*m;

  // R is the residual vector
  TMatrixD R=m-H*x;
  TMatrixD RT=R.T();
  R.T();          //Do this because R got also transposed
  TMatrixD Chi2=RT*V*R;
  
  if(Chi2[0][0]<Cut3) {
    t.SetPar(x);             //update track
    t.SetCov(C);
    t.SetChi2(Chi2[0][0]);
    t.SetNdf(ndim-5);
  }
  return Chi2[0][0];  
}
//*****************************************************************************************************
void Spectrometer::PrintEvent(int event){


  TMatrixD par= bestTrack.GetPar();
  TMatrixD cov= bestTrack.GetCov();
  float chi2= bestTrack.GetChi2();
  int ndf= bestTrack.GetNdf();
  vector<int> hits= bestTrack.GetHits();

  h12->Fill(chi2);

  //Print summary of this event
  if(chi2 < Cut3) {

    numberOfReconstructedTracks++;
    if(debug) cout << " " << endl;
    if(debug) cout << " New reconstructed track: " << endl;
    if(debug) cout << " z0 " << par[0][0] << " +- " << sqrt(cov[0][0]) << endl;
    if(debug) cout << " dz/dx " << par[1][0] << " +- " << sqrt(cov[1][1])  << endl;
    if(debug) cout << " y0 " << par[2][0]  << " +- " << sqrt(cov[2][2]) << endl;
    if(debug) cout << " dy/dx " << par[3][0]  << " +- " << sqrt(cov[2][2]) << endl;
    if(debug) cout << " q/p " << par[4][0]  << " +- " << sqrt(cov[2][2]) << endl;
    if(debug) cout << " Chi2 " << chi2  << "  for ndf " << ndf << endl;

    //Fill histograms
    h1->Fill( par[2][0] );
    h2->Fill( par[0][0] );
    h3->Fill( par[3][0] );
    h4->Fill( par[1][0]-thetaxz );
    h5->Fill( par[4][0]- 1/beamMomentum );
    h6->Fill( par[2][0]/sqrt(cov[2][2]) );
    h7->Fill( par[0][0]/sqrt(cov[0][0]) );
    h8->Fill( par[3][0]/sqrt(cov[3][3]) );
    h9->Fill( (par[1][0]-thetaxz)/sqrt(cov[1][1]) );
    h10->Fill( (par[4][0]- 1/beamMomentum)/sqrt(cov[4][4]) );
    int nh=(int)hits.size();
    if(nh<numberOfPlanes) numberOfInefficientTracks++;
    for(int i=0; i<nh; i++) {
      nTotalHits++;
      if(hitdata.at(hits.at(i)).IsNoise()) nNoiseHitsOnTrack++;
    }
    if(debug) cout << " number of hits " << nh << endl;

  } else {
    if(debug) cout << " fit failed  chisquared=" << chi2  << " ndf= " << ndf << endl;
    numberOfRejectedTracks++;
  }

  if(event>1) debug=false;
  if(event%1000==0) cout << " event " << event << endl;

  if(event==numberOfEvents-1) {
    cout << " " << endl;
    cout << " Generated Tracks " << numberOfEvents << endl;
    cout << " Reconstructed Tracks " << numberOfReconstructedTracks << endl;
    cout << " Tracks with missing hits " << numberOfInefficientTracks << endl;
    cout << " Tracks lost to quality cuts " << numberOfRejectedTracks << endl;
    cout << " Total hits on track " << nTotalHits << endl;
    cout << " Total noise hits on track " << nNoiseHitsOnTrack << endl;
    outfile->Write();
  }

}
//*********************************************************************************************