/* This file is part of MAUS: http://micewww.pp.rl.ac.uk:8080/projects/maus
 *
 * MAUS is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * MAUS is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MAUS.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include "src/common_cpp/Recon/EMR/TrackMomentum.hh"
#include "Utils/Squeak.hh"

namespace TrackMomentum {

// Masses
double _m;		// Mass of impinging particle [MeV/c^2]
double _z;		// Charge of the impinging particle [e]
double _me;		// Electron mass [MeV/c^2]

// Material properties
double _Z;		// Effective atomic number [e]
double _A;		// Relative atomic mass [g/mole]
double _rau;	 	// Density [g/cm^3]
double _I;		// Mean excitation potential [MeV]

// Empirical parameters of the delta function
double _C;
double _a;
double _k;
double _X0;
double _X1;

// Bethe-Bloch constant
double _K;		// [MeV*cm^2/g]

TF1* f(0);

double csda_momentum(double range,
		     std::string particle) {

  if ( range < 0 )
    return -1;

  if ( !initialize(particle) )
    return -1;

  if ( !f )
    f = new TF1("f", rangefunction, 0, 1000, 1);

  f->SetParameter(0, _m);
  double p = f->GetX(range, 0, 1000, pow(10, -6));
//   delete f;

  return p; // Total momentum [MeV/c]
}

double csda_momentum_error(double momentum,
			   double erange,
			   std::string particle) {

  if ( !initialize(particle) )
    return -1;

  if ( momentum < .05*_m ) // Limit of Bethe-Bloch validity
    return -1;

  if (!f)
    f = new TF1("f", rangefunction, 0, 1000, 1);

  f->SetParameter(0, _m);
  double dpdr = 1./(f->Derivative(momentum));
//   delete f;

  return dpdr*erange; // [MeV/c]
}

bool initialize(std::string particle) {

  // !!!TODO!!! get the properties from G4
  // Electron mass
  _me = .511;

  // Mass of impinging particle
  if ( particle == "muon" ) {
    _m = 105.66;
    _z = 0.;
  } else if ( particle == "pion" ) {
    _m = 139.57;
    _z = 0.;
  } else if ( particle == "electron" ) {
    _m = .511;
    _z = -1.;
  } else if ( particle == "positron" ) {
    _m = .511;
    _z = 1.;
  } else {
    MAUS::Squeak::mout(MAUS::Squeak::warning)
    << "WARNING in csda_momentum. The particle type is not supported,"
    << " the momentum will not be reconstructed."
    << std::endl;
    return false;
  }

  // Properties of polystryrene from the PDG (C8H8)
  double Cfrac = 48./56;
  double Hfrac = 8./56;
  _Z = pow(Cfrac*pow(6., 2.94)+Hfrac*pow(1., 2.94), 1./2.94);
  _A = _Z/0.53768;
  _rau = 1.060;
  _I = 68.7*1e-6;

  _C = 3.2999;
  _a = 0.16454;
  _k = 3.2224;
  _X0 = 0.1647;
  _X1 = 2.5031;

  // Bethe-Bloch mean ionization energy loss constant
  _K = 0.307;

  return true;
}

double tmax(double beta, double gamma, double M) {

  return 2*_me*beta*beta*gamma*gamma/(1+2*gamma*_me/M+pow(_me/M, 2));
}

double tauFel(double beta, double tau) {

  return 1.-pow(beta, 2)+(pow(tau, 2)/8.-(2.*tau+1.)*log(2.))/pow(tau+1., 2);
}

double tauFpos(double beta, double tau) {

  return 2.*log(2.) - (pow(beta, 2)/12.)*(23. + 14./(tau+2.)
		    + 10./pow(tau+2., 2) + 4./pow(tau+2., 3));
}

double deltaF(double alpha) {

  double X = log(alpha)/log(10); // Base 10 logarithm of alpha

  if ( X < _X0 ) {
    return 0;
  } else if ( X >= _X0 && X < _X1 ) {
    return 2*log(10)*X - _C + _a*pow(_X1-X, _k);
  } else {
    return 2*log(10)*X - _C;
  }
}

double bethe(double *x, double *par) {

  double beta = x[0]/sqrt(pow(x[0], 2)+1);
  double gamma = sqrt(pow(x[0], 2)+1);
  double tau = gamma-1;

  double dedx(0.);
  if ( _m > _me ) { 		// Heavy ion Bethe-Bloch
     dedx = 0.1*_K*_rau*_Z/(_A*beta*beta)
	    * (.5*log(2*_me*pow(x[0], 2)*tmax(beta, gamma, par[0])/pow(_I, 2))
	       - beta*beta - .5*deltaF(x[0]));
  } else if ( _z < 0 ) { 	// Electron Bethe-Bloch
     dedx = 0.1*_K*_rau*_Z/(_A*beta*beta)
	    * (log(_me*pow(x[0], 2)*tmax(beta, gamma, par[0])/(2*pow(_I, 2)))
	       + tauFel(beta, tau) - deltaF(x[0]));
  } else if ( _z > 0 ) {	// Positron Bethe-Bloch
     dedx = 0.1*_K*_rau*_Z/(_A*beta*beta)
	    * (log(_me*pow(x[0], 2)*tmax(beta, gamma, par[0])/(2*pow(_I, 2)))
	       + tauFpos(beta, tau) - deltaF(x[0]));
  }

  return dedx; // [MeV/mm]
}

double bremsstrahlung(double *x, double *par) {

  // Physics constants
  double alpha = 1./137.036;		// Fine-structure constant
  double NA = 6.02214179*1e23;		// Avogadro constant
  double kC = 8.987*1e9; 		// Coulomb constant [N*m^2/C^2]
  double e = 1.602*1e-19;	 	// Electron charge [C]
  double c = 299792458; 		// Speed of light [m/s]
  double Mkg = _m*1e6*e/(c*c);		// Particle mass in [kg]

  // Define function
  double gamma = sqrt(x[0]*x[0]+1);
  double tau = gamma-1;

  // -<dE/dx> = (4*n*Z*alpha^3*(hbar*c)^2/(m_e^2*c^4))*E*ln(183/Z^{1/3})
  // n = N_A*Z*rho/A, electron density in the target
  // alpha = e^2/(4*pi*eps_0*hbar*c) = kC*e^2/(hbar*c), Sommerfeld's constant
  double dedx = 0.1*4*alpha*(NA*_rau*_Z/_A)*_Z * 1e4*pow(kC*e*e/(Mkg*c*c), 2)
	      * tau*_m * log(183./pow(_Z, 1./3));

  return dedx; // [MeV/mm]
}

double spower(double *x, double *par) {

  return bethe(x, par) /*+ bremsstrahlung(x, par)*/; // [MeV/mm]
}

double invspower(double *x, double *par) {

  return 1./spower(x, par); // [mm/MeV]
}

double invspower_jac(double *x, double *par) {

  if ( x[0] < 0.01 )
      return 0.;

  double beta = x[0]/sqrt(pow(x[0], 2)+1);
  double dxde = invspower(x, par);

  return dxde*beta*par[0];
}

TF1 *f_bb(0);

double rangefunction(double *x, double *par) {

  double alpha = x[0]/par[0];

  if ( !f_bb )
    f_bb = new TF1("f", invspower_jac, 0, 100, 1);

  f_bb->SetParameter(0, par[0]);
  double R = f_bb->Integral(0, alpha, f->GetParameters(), 1e-3);
//   delete f_bb;

  return R; // mm
}
} // namespace TrackMomentum