// -------------------------------------------------------------------------
// -----                  CbmLitTrackPropagator source file               -----
// -----                  Created 16/07/07  by A. Lebedev               -----
// -------------------------------------------------------------------------
 


#include "CbmLitTrackPropagator.h"

#include "CbmLitEnvironment.h"



//constructors
CbmLitTrackPropagator::CbmLitTrackPropagator(CbmTrackExtrapolator *extrapolator):
   CbmTrackPropagator("CbmLitTrackPropagator"),
   fExtrapolator(extrapolator)
{
   CbmLitEnvironment* env = CbmLitEnvironment::Instance();
   fvAllMaterials = env->GetMaterials();
   
   // set default properties for the algorithm
   Properties().AddProperty("fFms", 1.1);
   Properties().AddProperty("fMass", 0.105);
   Properties().AddProperty("fEnergyLoss", 0.003);
   Properties().AddProperty("fApplyEnergyLoss", true);
}

//Destructor
CbmLitTrackPropagator::~CbmLitTrackPropagator() 
{
   //
}

// Initialization
void CbmLitTrackPropagator::Initialize()
{
   //
   Properties().GetProperty("fFms", fFms);
   Properties().GetProperty("fMass", fMass);
   Properties().GetProperty("fEnergyLoss", fEnergyLoss);
   Properties().GetProperty("fApplyEnergyLoss", fApplyEnergyLoss);
}

// Finalization
void CbmLitTrackPropagator::Finalize()
{
   //
}

void CbmLitTrackPropagator::Propagate( 
                             const CbmTrackParam *pParamIn,
                             CbmTrackParam *pParamOut,
                             Double_t zOut)
{
   *pParamOut = *pParamIn;
   Propagate(pParamOut, zOut);
}

void CbmLitTrackPropagator::Propagate(
                             CbmTrackParam *pParam,
                             Double_t zOut)
                                      
{
   Initialize();

   Double_t zIn = pParam->GetZ();
   
   if (zOut > zIn) fDownstream = kTRUE;
   else fDownstream = kFALSE;
   
   //Add noise from materials
   std::vector<CbmLitMaterial*>::iterator it_min, it_max;
   CbmLitMaterial* m = new CbmLitMaterial();
   
   m->SetZ(zIn);
   it_min = lower_bound(fvAllMaterials.begin(), fvAllMaterials.end(),
                              m, CbmLitMaterial::CmpFindZ);
   m->SetZ(zOut);
   it_max = lower_bound(fvAllMaterials.begin(), fvAllMaterials.end(), 
                              m, CbmLitMaterial::CmpFindZ);
   
   for(std::vector<CbmLitMaterial*>::iterator i = it_min; i != it_max; ) {

      Double_t Z = (*i)->GetZ();
      Double_t Thickness = (*i)->GetThickness();
      
      fExtrapolator->Extrapolate(pParam, Z);   

      //if ((*i)->IsTraversed(pParam->GetX(), pParam->GetY())) {
        Double_t RadLength =  Thickness / (*i)->GetRadLength();
        if (Thickness < 5.) AddThinScatter(pParam, RadLength);
          else AddThickScatter(pParam, RadLength, Thickness);
          
        if (fApplyEnergyLoss) AddEnergyLoss(pParam, *i);
      //}
      
      if (fDownstream) i++; else i--;
   }
   
   fExtrapolator->Extrapolate(pParam, zOut); 
   
   delete m;
}


void CbmLitTrackPropagator::AddThickScatter(
                             CbmTrackParam* pParam,
                             Double_t RadLength,
                             Double_t Thickness)
{
 
//TODO momentum too small

   Double_t tx = pParam->GetTx();
   Double_t ty = pParam->GetTy();
   Double_t qp = pParam->GetQp();
    
   Double_t t3 = 0.;
   if (RadLength > 1e-20) {
      Double_t norm = 1 + tx * tx + ty * ty;
      
      Double_t RadThick = RadLength * TMath::Sqrt(norm);
        
      Double_t t1 =  1 + 0.038 * log(RadThick);
      Double_t t2 = fFms * 0.0136 * qp * t1;
   
      t3 = ( 1 + fMass * fMass * qp * qp ) * RadThick * t2 * t2 * norm;
   }
  
   Double_t Q33 = ( 1 + tx * tx ) * t3;
   Double_t Q44 = ( 1 + ty * ty ) * t3;
   Double_t Q34 = tx * ty * t3;
   
   Double_t T23 = (Thickness * Thickness) / 3.0;
   Double_t T2 = Thickness / 2.0;

   Double_t D = (fDownstream) ? 1. : -1. ;

      Double_t C[15];
      pParam->CovMatrix(C);
      
      C[0] += Q33 * T23; 
      C[1] += Q34 * T23; 
      C[2] += Q33 * D * T2; 
      C[3] += Q34 * D * T2;
                           
      C[5] += Q44 * T23;           
      C[6] += Q34 * D * T2; 
      C[7] += Q44 * D * T2;
                                                      
      C[9] += Q33;
      C[10] += Q34;
   
      C[12] += Q44;
      
      pParam->SetCovMatrix(C);
   
      //
      //C[24] += 0.0001;
}


 void CbmLitTrackPropagator::AddThinScatter(
                              CbmTrackParam* pParam,
                              Double_t RadLength)
{ 
 
//TODO momentum too small

   Double_t tx = pParam->GetTx();
   Double_t ty = pParam->GetTy();
   Double_t qp = pParam->GetQp();
    
   Double_t t3 = 0.;
   if (RadLength > 1e-20) {
      Double_t norm = 1 + tx * tx + ty * ty;
      
      Double_t RadThick = RadLength * TMath::Sqrt(norm);
        
      Double_t t1 =  1 + 0.038 * log(RadThick);
      Double_t t2 = fFms * 0.0136 * qp * t1;
   
      t3 = ( 1 + fMass * fMass * qp * qp ) * RadThick * t2 * t2 * norm;
   }
  
  // Double_t Q33 = ( 1 + tx * tx ) * t3;
  // Double_t Q44 = ( 1 + ty * ty ) * t3;
  // Double_t Q34 = tx * ty * t3;
   

      Double_t C[15];
      pParam->CovMatrix(C);
      C[9] += ( 1 + tx * tx ) * t3;//Q33;
      C[12] += ( 1 + ty * ty ) * t3;//Q44;
      C[10] += tx * ty * t3;//Q34;
      pParam->SetCovMatrix(C);
}
 

void CbmLitTrackPropagator::AddEnergyLoss( 
                             CbmTrackParam* pParam,
                             CbmLitMaterial* Material)
{

   // dE/dx energy loss (Bethe-Block)
   Double_t tx = pParam->GetTx();
   Double_t ty = pParam->GetTy();
   Double_t qp = pParam->GetQp();  
   
   Double_t norm = TMath::Sqrt(1 + tx * tx + ty * ty);
   
   Double_t bbLoss = Material->GetThickness() * norm * 
                     Material->GetDensity() * fEnergyLoss * 
                     Material->GetZeff() / Material->GetAeff();
   
   //TODO bbEloss too big
   
   if (fDownstream) {
      bbLoss *= -1.0;
   }
   
   if (qp > 0.) {
      qp = 1. / ((1./qp) + bbLoss);
   }
   else{
      if (qp != 0. && ((1./qp) - bbLoss) != 0.) 
         qp = 1. / ((1./qp) - bbLoss);
   }
   
   pParam->SetQp(qp);
   
}

 
void CbmLitTrackPropagator::AddElectronEnergyLoss(
                             CbmTrackParam* pParam,
                             CbmLitMaterial* Material)
{
   //energy loss for electrons
 
   Double_t RadThick;
   
   Double_t tx = pParam->GetTx();
   Double_t ty = pParam->GetTy();
   Double_t qp = pParam->GetQp(); 
   
   Double_t norm = TMath::Sqrt(1 + tx * tx + ty * ty);
   
   if (fDownstream){
      RadThick = - (Material->GetThickness() / Material->GetRadLength()) * norm;
   }
   else{
      RadThick = (Material->GetThickness() / Material->GetRadLength()) * norm;
   }
   
   // TODO protect against t too big
   
   Double_t cov = pParam->GetCovariance(4, 4);
   cov += qp * qp * 
           (TMath::Exp(-RadThick * TMath::Log(3.0)/TMath::Log(2.0)) - 
            TMath::Exp(-2.0 * RadThick));
   pParam->SetCovariance(4, 4, cov);
                    
   pParam->SetQp(pParam->GetQp() * TMath::Exp(-RadThick));
}

ClassImp(CbmLitTrackPropagator)