//  Geant4 simulation for MuSR
//  AUTHOR: Toni SHIROKA, Paul Scherrer Institut, PSI
//  DATE  : 2008-05
//

#include "lem4PrimaryGeneratorAction.hh"
#include "lem4DetectorConstruction.hh"
#include "lem4PrimaryGeneratorMessenger.hh"
#include "G4Event.hh"
#include "G4ParticleGun.hh"
#include "G4ParticleTable.hh"
#include "G4ParticleDefinition.hh"
#include "Randomize.hh"
#include "G4ios.hh"
#include "G4ParticleGun.hh"
#include "G4UnitsTable.hh"
#include "globals.hh"
#include "G4Gamma.hh"
#include "G4ThreeVector.hh"
#include "G4RunManager.hh"
#include "time.h"
#include <iomanip>
#include "lem4RootOutput.hh"   //cks for storing some info in the Root output file
#include "lem4ErrorMessage.hh"

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo....
 
lem4PrimaryGeneratorAction::lem4PrimaryGeneratorAction(
                                            lem4DetectorConstruction* lem4DC)
  :lem4Detector(lem4DC), x0(0), y0(0), z0(-124*cm), xSigma(0), ySigma(0), zSigma(0), 
   rMaxAllowed(1e10*mm), zMinAllowed(-1e10*mm), zMaxAllowed(1e10*mm),
   p0(0), pSigma(0), pMinAllowed(0), pMaxAllowed(1e10*mm),
   xangle0(0), yangle0(0), xangleSigma(0), yangleSigma(0), pitch(0),
   UnpolarisedMuonBeam(false), xPolarisIni(1.), yPolarisIni(0.), zPolarisIni(0.),
   muonDecayTimeMin(-1), muonDecayTimeMax(-1), muonMeanLife(2197.03*ns),
   takeMuonsFromTurtleFile(false)
						      //, firstCall(true)
{
  G4int n_particle = 1;
  //cksdel  particleGun  = new lem4ParticleGun(n_particle);
  particleGun  = new G4ParticleGun(n_particle);
  
  //create a messenger for this class
  gunMessenger = new lem4PrimaryGeneratorMessenger(this);

  // default particle kinematic
  G4ParticleTable* particleTable = G4ParticleTable::GetParticleTable();
  G4ParticleDefinition* particle= particleTable->FindParticle("mu+");
  particleGun->SetParticleDefinition(particle);
  mu_mass = particle->GetPDGMass();
} 

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo....

lem4PrimaryGeneratorAction::~lem4PrimaryGeneratorAction()
{
  delete particleGun;
  delete gunMessenger;
  if (takeMuonsFromTurtleFile) {fclose(fTurtleFile);}
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo....

void lem4PrimaryGeneratorAction::GeneratePrimaries(G4Event* anEvent)
{
  // This function is called at the begining of event.

  //  // First of all check the setting of some variables.  Do it only once (for the first event).
  //  if (firstCall) {
  //    firstCall=false;
  //    char message[200];
  //    cout<<"********************  called just once  **********************"<<G4endl;
  //    // Check that the restriction of the generated position is reasonable:
  //    if ( sqrt(rMaxAllowed-x0)*(rMaxAllowed-x0)+(rMaxAllowed-y0)*(rMaxAllowed-y0) > 
  //         5*sqrt(xSigma*xSigma+ySigma*ySigma)  ) {
  //      sprintf(message,"lem4PrimaryGeneratorAction::GeneratePrimaries: Too strict restriction on the generated radius.");
  //      lem4ErrorMessage::GetInstance()->lem4Error(FATAL,"message",true);
  //    }
  //    // Check that the restriction on the z-coordinate of the generated postion is reasonable:
  //    if ( ((-zMaxAllowed+z0) > 5*fabs(zSigma)) || ((zMinAllowed-z0) > 5*fabs(zSigma)) ) {
  //      sprintf(message,"lem4PrimaryGeneratorAction::GeneratePrimaries: Too strict restriction on the generated z-coordinate.");
  //      lem4ErrorMessage::GetInstance()->lem4Error(FATAL,"message",true);
  //    }
  //  }
  
  G4double x, y, z;
  G4double p;// ke_offset;
  G4double xangle, yangle;

  //ke_offset = 3.73*keV; // Kin. energy offset due to a 3.73 kV acceleration at the Carbon foil
  
  if (takeMuonsFromTurtleFile) {
    char  line[501];  
    G4int checkNrOfCounts=0;
    do {
      float xTmp, yTmp, xAngleTmp, yAngleTmp, pTmp;
      if (feof(fTurtleFile)) {rewind(fTurtleFile);G4cout<<"End of Turtle file."<<G4endl;}
      fgets(line,500,fTurtleFile);
      //    sscanf(&line[0],"%g %g %g %g %g %g %g",&x,&xangle,&y,&yangle,&p,&w,&pol);
      sscanf(&line[0],"%g %g %g %g %g",&xTmp,&xAngleTmp,&yTmp,&yAngleTmp,&pTmp);
      xangle = xAngleTmp*mrad;
      yangle = yAngleTmp*mrad;
      x      = xTmp*mm + z0*xangle ;       // usually z0 is negative
      y      = yTmp*mm + z0*yangle ;
      p      = pTmp*GeV;
      // add some beam offset and beam tilt, if requested:
      xangle = xangle + xangle0;
      yangle = yangle + yangle0;
      x      = x + x0;
      y      = y + y0;
      //      G4cout<<": ";
      checkNrOfCounts++;
      if (checkNrOfCounts>1000) {
	G4cout<<"lem4PrimaryGeneratorAction::GeneratePrimaries:  Too strict requirements on the r position!"<<G4endl;
      }
    } while( (x*x+y*y)>(rMaxAllowed*rMaxAllowed) );
    z=z0;
    //    G4cout<<"x,y,z=("<<x/mm<<","<<y/mm<<","<<z/mm<<"),  angles="<<xangle/mrad<<","<<yangle/mrad<<"  p="<<p/MeV<<G4endl;
  }

  else {          // Generate the starting position of the muon by random
    //  rMaxAllowed  ... maximal radius, within which the muon can be generated
    //  x0, y0, z0   ... central point around which the muons are generated
    //  xSigma, ySigma, zSigma  ... sigma of the (gaussian) distributions of the beam
    //  x, y, z      ... actual initial position of the generated muon
    
    G4int checkNrOfCounts=0;
    do {
      if (xSigma>0)      {x = G4RandGauss::shoot(x0,xSigma);}           //  Gaussian distribution
      else if (xSigma<0) {x = x0 + xSigma*(G4UniformRand()*2.-1.);}     //  Uniform step distribution
      else               { x = x0;}                                     //  Point-like

      if (ySigma>0)      {y = G4RandGauss::shoot(y0,ySigma);}
      else if (ySigma<0) {y = y0 + ySigma*(G4UniformRand()*2.-1.);}
      else               {y = y0;}

      if (zSigma>0)      {z = G4RandGauss::shoot(z0,zSigma);}
      else if (zSigma<0) {z = z0 + zSigma*(G4UniformRand()*2.-1.);}
      else               {z = z0;}

      checkNrOfCounts++;
      if (checkNrOfCounts>1000) {
	G4cout<<"lem4PrimaryGeneratorAction::GeneratePrimaries:  Too strict requirements on the r or z position!"<<G4endl;
      }
    } while( ((x*x+y*y)>(rMaxAllowed*rMaxAllowed))||(z>zMaxAllowed)||(z<zMinAllowed) );
      // The generated muon has to stay within some well defined region, e.g. within the beampipe.
  
      
    // If the user defines energy instead of momentum
       // if (E0==0 && p0>0) {p = p0;}
       // if (E0>0 && p0==0) {p0 = std::sqrt(E0*E0 + 2*mu_mass*E0);}
       // else if (E0>0 && p0>0) {
       //   G4cout<<"Define either kinetic energy or momentum, but not both!"<<G4endl;
       //   G4cout << "S T O P    F O R C E D" << G4endl;
       //   exit(1);
       // }
    
              
    // Now generate the momentum
    checkNrOfCounts=0;
    do {
      if (pSigma>0) {p = G4RandGauss::shoot(p0,pSigma);}
      else {p=p0;}
      checkNrOfCounts++;
      if (checkNrOfCounts>1000) {
         G4cout<<"lem4PrimaryGeneratorAction::GeneratePrimaries:  Too strict requirements on the momentum!"<<G4endl;
      }
    } while ( (p>pMaxAllowed)||(p<pMinAllowed) );
    //G4cout<<"Kamil: momentum p="<<p/MeV<<" MeV"<<G4endl;
    
    // Add some initial angle (px and py component of the momentum)
    if (xangleSigma>0) { xangle = G4RandGauss::shoot(xangle0,xangleSigma); }
    else { xangle = xangle0; }
    //  Add the beam tilt, which depends on the distance from the beam centre.
    if (xSigma>0) {xangle += - pitch * (x-x0)/(1*mm);} //xSigma; } // Changed to absolute units of pitch
    
    if (yangleSigma>0) { yangle = G4RandGauss::shoot(yangle0,yangleSigma); }
    else { yangle = yangle0; }
    //  Add the beam tilt, which depends on the distance from the beam centre.
    if (ySigma>0) {yangle += - pitch * (y-y0)/(1*mm);} //ySigma; } // Changed to absolute units of pitch

  }  // end of the part specific for the muons generated by random rather then from TURTLE
  

  // Calculate the final momentum
  G4double px, py, pz;
  px = p*sin(xangle);
  py = p*sin(yangle);
  pz = std::sqrt(p*p - px*px - py*py);

  // Assign spin 
  G4double xpolaris=0, ypolaris=0, zpolaris=0;
  if (UnpolarisedMuonBeam) {
    G4cout<<"lem4PrimaryGeneratorAction.cc:  Unpolarised muon beam not yet implemented!"<<G4endl;
    G4cout<<"      =========>  S T O P"<<G4endl;
    exit(0);
    //  xpolaris = ;
    //  ypolaris = ;
    //  zpolaris = ;
  }
  else {
    xpolaris = xPolarisIni;
    ypolaris = yPolarisIni;
    zpolaris = zPolarisIni;
  }
  
  //  if (InitialMuonPolarization==3) {                       // longitudinal
  //    zpolaris=-1;
  //  }
  //  else if (InitialMuonPolarization==2) {                  // transverse
  //    xpolaris=1;
  //  }
  //  else if (InitialMuonPolarization==1) {                  // longitudinal
  //    xpolaris = -px/p;
  //    ypolaris = -py/p;
  //    zpolaris = -pz/p;
  //  }
  //  else {                                                  // transverse
  //    xpolaris = -pz/p;
  //  ypolaris = -py/p;
  //    zpolaris = -px/p;
  //  }

  particleGun->SetParticlePosition(G4ThreeVector(x,y,z));
  //particleGun->SetParticleMomentum(G4ThreeVector(px,py,pz));
  G4double particleEnergy = std::sqrt(p*p+mu_mass*mu_mass)-mu_mass;
  //particleGun->SetParticleEnergy(particleEnergy+ke_offset);
  particleGun->SetParticleEnergy(particleEnergy);
  particleGun->SetParticleMomentumDirection(G4ThreeVector(px,py,pz));
  particleGun->SetParticlePolarization(G4ThreeVector(xpolaris,ypolaris,zpolaris));
  particleGun->GeneratePrimaryVertex(anEvent);

  //  G4cout<<"lem4PrimaryGeneratorAction: Parameters:"<<G4endl;
  //  G4cout<<"     x0,y0,z0="<<x0/mm<<","<<y0/mm<<","<<z0/mm<<"   Sigma="<<xSigma/mm<<","<<ySigma/mm<<","<<zSigma/mm<<G4endl;
  //  G4cout<<"     rMaxAllowed="<<rMaxAllowed/mm<<"   zMaxAllowed="<<zMaxAllowed/mm<<"   zMinAllowed="<<zMinAllowed/mm<<G4endl;
  //  G4cout<<"     p0="<<p0/MeV<<"   pSigma="<<pSigma/MeV
  //	<<"   pMinAllowed="<<pMinAllowed/MeV<<"  pMaxAllowed=<<"<<pMaxAllowed/MeV<<G4endl;
  //  G4cout<<"     angle0="<<xangle0/deg<<","<<yangle0/deg<<",nic"
  //	<<"   Sigma="<<xangleSigma/deg<<","<<yangleSigma/deg<<",nic"<<G4endl;
  //  G4cout<<"     pitch="<<pitch/deg<<G4endl;
  //
  //  G4cout<<"lem4PrimaryGeneratorAction: Generated muon:"<<G4endl;
  //  G4cout<<"     x,y,z="<<x/mm<<","<<y/mm<<","<<z/mm<<"      angle="<<xangle/deg<<","<< yangle/deg<<",nic"<<G4endl;
  //  G4cout<<"     p="<<px/MeV<<","<<py/MeV<<","<<pz/MeV<<"    E="<< (particleGun->GetParticleEnergy())/MeV<<G4endl;
  //  G4cout<<"     polarisation="<<xpolaris<<","<<ypolaris<<","<<zpolaris<<G4endl;

  


  // if requested by "/gun/decaytimelimits", set the decay time of the muon such that it is within
  // the required time window.  Otherwise the decay time is set internally by Geant.
  if (muonDecayTimeMax>0.) {
    //    G4cout<<"muonDecayTimeMin="<<muonDecayTimeMin/ns<<" ns ,  muonDecayTimeMax="<<muonDecayTimeMax/ns
    //	  <<" ns ,   muonMeanLife="<<muonMeanLife/ns<<" ns."<<G4endl;
    // find the primary muon
    //    G4double decaytime;
    //    do {                                        // generate the decaytime within the specified time window
    //      decaytime = -muonMeanLife*log(1-G4UniformRand());
    //    } while ((decaytime<muonDecayTimeMin)||(decaytime>muonDecayTimeMax));
    //
    //  This algorithm should be much faster then the previous one:
    G4PrimaryParticle* generatedMuon = anEvent->GetPrimaryVertex(0)->GetPrimary(0);
    //    G4double decayLowerLimit = 1-exp(-muonDecayTimeMin/muonMeanLife);
    //    G4double decayUpperLimit = 1-exp(-muonDecayTimeMax/muonMeanLife);
    //    G4double randomVal       = G4UniformRand()*(decayUpperLimit-decayLowerLimit) + decayLowerLimit;
    //    G4double decaytime       = -muonMeanLife*log(1-randomVal);
    //
    //  The following code is numerically more stable compared to the commented lines above:
    G4double expMin    = exp(-muonDecayTimeMin/muonMeanLife);
    G4double expMax    = exp(-muonDecayTimeMax/muonMeanLife);
    G4double decaytime = -muonMeanLife * log(G4UniformRand()*(expMax-expMin)+expMin);
    //
    //    G4cout<<"decaytime="<<decaytime/ns<<"ns."<< G4endl;
    generatedMuon->SetProperTime(decaytime);
  }

  // Save variables into ROOT output file:
  lem4RootOutput* myRootOutput = lem4RootOutput::GetRootInstance();
  myRootOutput->SetInitialMuonParameters(x,y,z,px,py,pz,xpolaris,ypolaris,zpolaris);
  
}

///////////////////////////////////////////////////////////////////////
//
//


void lem4PrimaryGeneratorAction::SetInitialMuonPolariz(G4ThreeVector vIniPol)
{
  G4double magnitude=vIniPol.mag();
  if(magnitude<0.00000001) {
    G4cout<< "Unpolarised initial muons"<<G4endl;
    UnpolarisedMuonBeam=true;
  }
  else {
    xPolarisIni=vIniPol(0)/magnitude;
    yPolarisIni=vIniPol(1)/magnitude;
    zPolarisIni=vIniPol(2)/magnitude;
    G4cout<< "Initial Muon Polarisation set to ("<<xPolarisIni<<","<<yPolarisIni<<","<<zPolarisIni<<")"<<G4endl;
  }
}

void lem4PrimaryGeneratorAction::SetMuonDecayTimeLimits(G4ThreeVector decayTimeLimits) {
  muonDecayTimeMin = decayTimeLimits[0];
  muonDecayTimeMax = decayTimeLimits[1];
  muonMeanLife     = decayTimeLimits[2];
  // store the muon decay time parameters to the Root output
  lem4RootOutput* myRootOutput = lem4RootOutput::GetRootInstance();
  myRootOutput->StoreGeantParameter(2,muonDecayTimeMin/microsecond);
  myRootOutput->StoreGeantParameter(3,muonDecayTimeMax/microsecond); 
  myRootOutput->StoreGeantParameter(4,muonMeanLife/microsecond); 
}

void lem4PrimaryGeneratorAction::SetTurtleInput(G4String turtleFileName) {
  takeMuonsFromTurtleFile = true;
  fTurtleFile = fopen(turtleFileName.c_str(),"r");
  if (fTurtleFile==NULL) {
    G4cout << "E R R O R :    Failed to open TURTLE input file \"" << turtleFileName 
	   <<"\"."<< G4endl;
    G4cout << "S T O P    F O R C E D" << G4endl;
    exit(1);
  }
  else {G4cout << "Turtle input file  \"" << turtleFileName <<"\" opened."<< G4endl;}
}