#include "Mille.h"
#include "TRandom.h"

#include <iostream>
using std::cout;
using std::endl;

#include <fstream>
using std::ofstream;

#include <iomanip>
using std::setprecision;
using std::setw;

#include <vector>
using std::vector;

//     define detector geometry
const int   NPLAN = 100;   // Number of detecor planes
const float DETX  = 10.0;  // x-pos of first plane
const float DISX  = 10.0;  // distance between planes
const float THCK  = 2.0;   // thickness of one plane
const float HEIT  = 100.0; // height of detector plane
const float DISPL = 0.1;   // max displacement 1 mm 
const float DRIFT = 0.02;  // max Vdrift deviation 2 % 

void WriteSteeringFile()
{

  ofstream steerFile;
  steerFile.open ("mp2str.txt");

  steerFile << "*            Default test steering file"<<endl;
  steerFile << "cfiles ! binary file are fortran"<<endl;
  steerFile << "mp2con.txt   ! constraints text file "<<endl;
  steerFile << "mp2tst.bin   ! binary data file"<<endl;
  steerFile << "*outlierrejection 100.0 ! reject if Chi^2/Ndf >"<<endl;
  steerFile << "*outliersuppression 3   ! 3 local_fit iterations"<<endl;
  steerFile << " "<<endl;
  steerFile << "*bandwidth 6 ! width of band matrix"<<endl;
  steerFile << "method diagonalization 5 0.001 ! diagonalization"<<endl;
  steerFile << "method fullGMRES       5 0.01 ! minimal residual"<<endl;
  steerFile << "method sparseGMRES     5 0.01 ! minimal residual"<<endl;
  steerFile << "method cholesky        5 0.001 ! Cholesky"<<endl;
  steerFile << "method bandcholesky    5 0.1   ! Cholesky"<<endl;
  steerFile << "method inversion       3 0.001 ! Gauss matrix inversion"<<endl;
  steerFile << "* last method is applied"<<endl;
  steerFile << " "<<endl;

  steerFile.close();

}

int CalculateUniqueLabel(int i)
{
  // i+1 because arrays in c start from 0, in fortran they start from 1 
  return 10+(i+1)*2;
}

void WriteConstrainFile(vector<float> &DEL)
{

  ofstream constrainFile;
  constrainFile.open ("mp2con.txt");
  constrainFile << " Constraint  0.0" << endl;

  for ( int i=0; i< NPLAN; i++) {
    int LABELT = CalculateUniqueLabel(i);
    constrainFile <<"    "<<setw(4) << LABELT<<"   1.00000"<<endl;
  }
 
  constrainFile << " Constraint 0.0" << endl;

  float SW=0.0;//                         ! tilt constraint 
  float SV=0.0;
  float S1=0.0;
  float S2=0.0;
  float DBAR = 0.5*(NPLAN-1)*DISX;
  float XBAR = DETX + 0.5*(NPLAN-1)*DISX;// ! +0.5*THCK

  for(int i=0; i< NPLAN; i++) {
    int LABELT = CalculateUniqueLabel(i);
    float X = DETX + i*DISX;//          !+0.5*THCK
    float WW=(X-XBAR)/DBAR;//
    constrainFile <<"    "<<setw(4) << LABELT<<setprecision(5) 
		  << setw(10)<<WW<<endl;

    S1 = S1+DEL[i];
    S2=S2+WW*DEL[i];
    SW=SW+WW;
    SV=SV+WW*WW; 
  }
  constrainFile.close();
 
  float DET=NPLAN*SV-SW*SW;
  float EPS=(SV*S1-SW*S2)/DET;  
  float ETA=(NPLAN*S2-SW*S1)/DET;
  for(int i=0; i< NPLAN; i++) {
    float X= DETX + i*DISX; 
    float WW=(X-XBAR)/DBAR;
    DEL[i]=DEL[i]-EPS-ETA*WW;//        ! correct displacement for constraints
  }

  float SUM1=0.0;
  float SUM2=0.0;
  for(int i=0; i< NPLAN; i++) {
    SUM1=SUM1+DEL[i]; 
    float X=DETX+i*DISX;//          !+0.5*THCK
    float WW=(X-XBAR)/DBAR;
    SUM2=SUM2+DEL[i]*WW;
  }

  cout<<" Check for constraints "<<SUM1<<", "<<SUM2<<endl;
}

int GenerateEvents(int IP, const vector <float> &EFF, 
		    const vector<float> &DEL, const vector<float> &SGM, 
		    const vector<float> &DVD, 
		    vector<float> &XHITS, vector<int> &IHITS,
		    vector<float> &YHITS, vector <float> &YDRFT,
		    vector<float> &SIGMA)
{
  // track definition with vertex and slope
  // Uniform distribution of the vertex between 0.45 and 0.55 of the toal
  // detector height
  // the slope is between 0 and 0.05  
  float YNULL= 0.5*HEIT + 0.1*HEIT*(gRandom->Uniform(1.)-1.);
  float SLOPE=(gRandom->Uniform(1.)-1.)*HEIT/((NPLAN-1)*DISX);

  if (IP != 0) {
    cout<<"Vertex: "<<YNULL<<"   Slope: "<<SLOPE<<endl;
  }

  int NHITS=-1;
  for ( int i=0; i< NPLAN; i++) {
    // take into account detector inefficencies
    float ran = gRandom->Uniform(1.);  
    if ( ran > EFF[i]) continue;

    // calculate true x and y positions for the track
    float X=DETX+i*DISX;                 // detector x position
    float YLIN  = YNULL + SLOPE*X;       // true y value at position  x

    float YBIAS = YLIN-DEL[i];           // y value in shifted detector 
    int NWIRE = (int)(1.0+YBIAS/4.0);    // wire number

    if (NWIRE < 0. || NWIRE > 25.) continue;

    NHITS++;                             // track hits the plane
    XHITS[NHITS]=X;
    IHITS[NHITS]=i;

    // calculate detector respons tacking ito account the detector
    // resolution and drift velocity differences 
    float GR=gRandom->Gaus(0.,1.);
    float YMEAS=SGM[i]*GR;            // hit smearing
    float YDVDS=0.0;                  // drift velocity correction
    YHITS[NHITS]=YBIAS+YMEAS+YDVDS;   // measured position
    float YWIRE=NWIRE*4.0-2.0;        // y position of wire
    YDRFT[NHITS]=YBIAS-YWIRE;         // drift length
    YDVDS=YDRFT[NHITS]*DVD[i];        // position error from drift velocity
    YHITS[NHITS]=YBIAS+YMEAS-YDVDS;   // reconstructed hit position
    SIGMA[NHITS]=SGM[i];

    if (IP != 0) {
      cout << setw(4) << NHITS << setw(4) << i<< setw(5) << X  
	   << setw(10) << YLIN << setw(10) << YBIAS << setw(13) << YMEAS
	   <<setw(8)<<SGM[i]<<setw(10)<<YHITS[NHITS]<<setw(10)<<GR<<setw(12)
	   <<DEL[i]<<endl;
    }
  }

 return NHITS;  

}

int main() 
{
  vector<float> EFF(NPLAN);    // efficency of detector planes
  vector<float> SGM(NPLAN);    // measurement error (sigma)
  vector<float> DVD(NPLAN);    // rel. drift velocity deviation

  vector<float> DEL(NPLAN);    // introduced detector shift

  // define detector efficencies and measuremnet errors
  for ( int i=0; i< NPLAN; i++) {
    EFF[i]=0.90;
    SGM[i]=0.0150;
  }
    // modify one plane (7)
  EFF[7]=0.1;     // low efficiency
  SGM[7]=0.0400;  // and bad resolution
  
  // misalign detector planes   
  for ( int i=0; i< NPLAN; i++) {
    DEL[i]=DISPL*gRandom->Gaus(0.,1.);
    DVD[i]=DRIFT*gRandom->Gaus(0.,1.);
  }

  // No shift for two detector planes
  DEL[10]=0.0;
  DEL[90]=0.0;

  WriteSteeringFile();

  WriteConstrainFile(DEL);

  //Open MillePede Output file
  //  Mille Mille("mp2tst.bin", false, true); // write human readable ascii file
  Mille  Mille("mp2tst.bin", true, false); // write binary file needed by pede

  //  events loop ------------------------------------------------------

  int NCOUNT=10000;
  int NTHITS=0;
  int NHITS;
  int NRECDS=0;
  float DERLC[2];
  float DERGL[2];
  int LABEL[2];

  vector<float> XHITS(100);
  vector<int> IHITS(100);
  vector<float> YHITS(100);
  vector<float> YDRFT(100);
  vector<float> SIGMA(100);

  for(int ICOUNT=0; ICOUNT< NCOUNT; ICOUNT++) {
    int IP=0;
    if (ICOUNT == 8759) IP=1;
    NHITS = GenerateEvents(IP, EFF, DEL, SGM, DVD,
		   XHITS, IHITS, YHITS, YDRFT, SIGMA);
    
    if (NHITS > 90) cout<<"hits are "<<NHITS<<endl;
    if (IP != 0) cout << "******************"<<endl;
    for(int i=0; i<NHITS; i++) {
      DERLC[0]=1.0;
      DERLC[1]=XHITS[i];
      DERGL[0]=1.0;
      DERGL[1]=YDRFT[i];
      LABEL[0]=CalculateUniqueLabel(IHITS[i]);
      LABEL[1]=500 + (IHITS[i]+1);
      if (IHITS[i] == 99) cout<<"hits are 99"<<endl;
      if (IP != 0) {
	cout << setw(4) << DERLC[0] << setw(4)  << DERLC[1] 
	     << setw(4) << DERGL[0] << setw(10) << DERGL[1] 
	     << setw(5) << LABEL[0] << setw(5)  << LABEL[1] 
	     <<setw(10) << YHITS[i] << setw(10) << SIGMA[i]
	     << endl;
      }
      Mille.mille(2,DERLC,2,DERGL,LABEL,YHITS[i],SIGMA[i]);
      NTHITS++;   // count hits
    }
    if (IP != 0) cout << "******************"<<endl;
    Mille.end();
    NRECDS=NRECDS+1;   // count records
  }
  cout<<""<<endl;
  cout<<"Shifts and drift velocity deviations:"<<endl;
  ofstream shiftFile;
  shiftFile.open ("shifts.txt");
  for(int i=0; i<NPLAN; i++) {
    cout<<setw(4)<<i<<setw(12)<<DEL[i]<<setw(12)<<DVD[i]<<endl;
    int label = CalculateUniqueLabel(i);
    shiftFile << setw(5) << label << setw(12)<<DEL[i] << endl;
    label = 500 + (i+1);
    shiftFile << setw(5) << label << setw(12)<<DVD[i] << endl;
  }
  shiftFile.close ();

  cout<<NCOUNT<<" tracks generated with "<<NTHITS<<" hits."<<endl;
  cout<<NRECDS<<" records written."<<endl;
 
  cout<<"Total Hits: "<<NTHITS<<endl;
  cout<<"Total events: "<<NRECDS<<endl;
}