// -------------------------------------------------------------------------
// -----                   PndLmdGeaneTask source file                 -----
// -----                 Created 18/07/08  by T.Stockmanns             -----
// -----           modified for Lmd by M. Michel & A.Karavdina         -----
// -------------------------------------------------------------------------
// libc includes
#include <iostream>

// Root includes
#include "TROOT.h"
#include "TClonesArray.h"
#include "TVector3.h"
#include <TMatrixDSym.h>


// framework includes
#include "FairRootManager.h"
#include "PndLmdGeaneTask.h"
#include "FairRun.h"
#include "FairRuntimeDb.h"
#include "FairHit.h"
#include "PndMCTrack.h"
#include "FairBaseParSet.h"
#include "TGeant3.h"
#include "FairTrackParH.h"
#include "FairTrackParP.h"
#include "TDatabasePDG.h"

// PndSds includes
#include "PndSdsMCPoint.h"
#include "PndSdsHit.h"
//PndLmd includes
#include "PndLinTrack.h"

#include <vector>
#include <map>


// -----   Default constructor   -------------------------------------------
PndLmdGeaneTask::PndLmdGeaneTask() : FairTask("Geane Task for PANDA Lmd"), fEventNr(0), fUseMVDPoint(false)
{
}
// -------------------------------------------------------------------------


PndLmdGeaneTask::PndLmdGeaneTask(Double_t pBeam,TVector3 IP): FairTask("Geane Task for PANDA Lmd"), fEventNr(0), fUseMVDPoint(false)
{
  fPbeam = pBeam;
  // if(fPDGid!=-2212){ //calculate momentum if particle is not antiproton
  //     }
  cout<<"Beam Momentum for particle with PDGid#"<<fPDGid<<" this run is "<<fPbeam<<endl;
  vtx = IP;
  cout<<"Interaction Point:"<<endl;
  vtx.Print();
}


// -----   Destructor   ----------------------------------------------------
PndLmdGeaneTask::~PndLmdGeaneTask()
{
}

// -----   Public method Init   --------------------------------------------
InitStatus PndLmdGeaneTask::Init()
{

  // Get RootManager
  FairRootManager* ioman = FairRootManager::Instance();
  if ( !ioman){
    std::cout << "-E- PndLmdGeaneTask::Init: "<< "RootManager not instantiated!" << std::endl;
    return kFATAL;
  }

  // Get input arrays
  fMCTracks = (TClonesArray*) ioman->GetObject("MCTrack");
  if (!fMCTracks){
    std::cout << "-W- PndLmdGeaneTask::Init: "<< "No MCTrack" << " array!" << std::endl;
    return kERROR;
  }

  fMCHits = (TClonesArray*) ioman->GetObject("LMDPoint");
  if ( !fMCHits)	{
    std::cout << "-W- PndLmdGeaneTask::Init: "<< "No LMDPoint"<<" array!" << std::endl;
    return kERROR;
  }

  fTracks = (TClonesArray*) ioman->GetObject("LMDTrack");
  if (!fTracks){
    std::cout << "-W- PndLmdGeaneTask::Init: "<< "No Track" << " array!" << std::endl;
    return kERROR;
  }

  fHits = (TClonesArray*) ioman->GetObject("LMDHitsStrip");
  if (!fHits){
    std::cout << "-W- PndLmdGeaneTask::Init: "<< "No Hits" << " array!" << std::endl;
    return kERROR;
  }

  fTrackParGeane = new TClonesArray("FairTrackParH");
  ioman->Register("GeaneTrackPar","Geane", fTrackParGeane, kTRUE);

  fTrackParIni = new TClonesArray("FairTrackParH");
  ioman->Register("GeaneTrackIni","Geane", fTrackParIni, kTRUE);

  fTrackParFinal = new TClonesArray("FairTrackParH");
  ioman->Register("GeaneTrackFinal","Geane", fTrackParFinal, kTRUE);

  // fTrackParGeane = new TClonesArray("FairTrackParP");
  // ioman->Register("GeaneTrackPar","Geane", fTrackParGeane, kTRUE);

  // fTrackParIni = new TClonesArray("FairTrackParP");
  // ioman->Register("GeaneTrackIni","Geane", fTrackParIni, kTRUE);

  // fTrackParFinal = new TClonesArray("FairTrackParP");
  // ioman->Register("GeaneTrackFinal","Geane", fTrackParFinal, kTRUE);

  fDetName = new TClonesArray("TObjString");
  ioman->Register("DetName", "Geane", fDetName, kTRUE);

  fPro = new FairGeanePro();
  fGeoH = PndGeoHandling::Instance();
  FairRun* fRun = FairRun::Instance();
  FairRuntimeDb* rtdb = fRun->GetRuntimeDb();
  // FairBaseParSet* par=(FairBaseParSet*)
  //   (rtdb->findContainer("FairBaseParSet"));
  // fPbeam = par->GetBeamMom();
  //  cout<<"Beam Momentum for this run is "<<fPbeam<<endl;
  return kSUCCESS;
}
// -------------------------------------------------------------------------
void PndLmdGeaneTask::SetParContainers()
{
  // Get Base Container
 /// FairRun* ana = FairRun::Instance();
 //s FairRuntimeDb* rtdb=ana->GetRuntimeDb();

}


// -----   Public method Exec   --------------------------------------------
void PndLmdGeaneTask::Exec(Option_t* opt)
{
  // cout<<"PndLmdGeaneTask::Exec starts!"<<endl;
  std::map<int, std::vector<int> > mcHitMap;//Track ->  MCHits
  fTrackParGeane->Delete();
  fTrackParIni->Delete();
  fTrackParFinal->Delete();
  fDetName->Delete();

  mcHitMap = AssignHitsToTracks();
  fEventNr++;

  //Charge & mass of particle
  Int_t PDGCode = -2212; //antiproton
  //  Int_t PDGCode = fPDGid;
  TDatabasePDG *fdbPDG = TDatabasePDG::Instance();
  TParticlePDG *fParticle = fdbPDG->GetParticle(PDGCode);
  Double_t  fCharge = fParticle->Charge();

  //go through all tracks
  int glI = fTracks->GetEntriesFast();
  // cout<<"glI = "<<glI<<endl;
  Int_t counterGeaneTrk = 0;
  for (Int_t i = 0; i<glI;i++){ 
    //  cout<<"PndLmdGeaneTask::Exec for track#"<<i<<endl;
      TVector3 StartPos, StartPosErr, StartMom, StartMomErr, StartO, StartU, StartV;
      int p = 0;

      ///Get parameters of real track
      PndLinTrack* recTrack = (PndLinTrack*)(fTracks->At(i));
      //   cout<<"recTrack="<<recTrack<<endl;
      TString 	DecName = recTrack->GetDetName();
      if(DecName!="Lumi") continue;
      // cout<<"DecName: "<<DecName<<endl;
      //Vector of particle momentum and starting point
      TVector3 DirVec =  recTrack->GetDirectionVec();
      StartMom = TVector3(DirVec.X()*fPbeam,DirVec.Y()*fPbeam,DirVec.Z()*fPbeam); 
      StartPos = recTrack->GetStartVec();
      // TVector3 StartPosOut;
      // StartPosOut.SetZ(StartPos.Z()-0.008);
      // StartPosOut.SetX(StartPos.X() - 0.008*DirVec.X());
      // StartPosOut.SetY(StartPos.Y() - 0.008*DirVec.Y());
      StartPosErr = recTrack->GetStartErrVec();
      //   StartPosErr.SetZ(1000.); //TEST
      // StartPosErr.SetZ(TMath::Hypot(StartPosErr.Z(),StartPosErr.X()*sin(2.326)));
      StartMomErr = (recTrack->GetDirectionErrVec())*fPbeam;
      
      // // cout<<"------------------------------------------"<<endl;      
      // cout<<"StartPos:"<<endl;
      // StartPos.Print();
      // cout<<"StartPosErr:"<<endl;
      // StartPosErr.Print();

      // cout<<"StartMom: "<<StartMom.Mag()<<endl;
      // StartMom.Print();
      // cout<<"StartMomErr: "<<StartMomErr.Mag()<<endl;
      // StartMomErr.Print();     

      TClonesArray& clref1 = *fTrackParIni;
      Int_t size1 = clref1.GetEntriesFast();

      ///Propagate to the PCA to a space point---------------------------------
      FairTrackParH *fStart = 
       	 new (clref1[size1]) FairTrackParH(StartPos, StartMom, StartPosErr, StartMomErr, fCharge);
      TClonesArray& clref = *fTrackParGeane;
      Int_t size = clref.GetEntriesFast();
      FairTrackParH *fRes = new(clref[size]) FairTrackParH();
      fPro->SetPoint(vtx);
      fPro->PropagateToPCA(1,-1);
      //  fPro->BackTrackToVertex();
      Bool_t isProp = fPro->Propagate(fStart, fRes, PDGCode);
      ///----------------------------------------------------------------------

      // ///Forwars propagate to the 1st plane to a space point---------------------------------
      // PndMCTrack* mctrk = (PndMCTrack*)(fMCTracks->At(i));
      // TVector3 MomMC = mctrk->GetMomentum();
      // TVector3 PosMC = mctrk->GetStartVertex();
      // FairTrackParH *fStart = 
      // 	new (clref1[size1]) FairTrackParH(PosMC, MomMC, StartPosErr, StartMomErr, fCharge);
      // TClonesArray& clref = *fTrackParGeane;
      // Int_t size = clref.GetEntriesFast();
      // FairTrackParH *fRes = new(clref[size]) FairTrackParH();
      // fPro->SetPoint(StartPos);
      // //fPro->PropagateToPCA(1,-1);
      // fPro->PropagateToPCA(1,1);
      // Bool_t isProp = fPro->Propagate(fStart, fRes, PDGCode);
      // ///----------------------------------------------------------------------


      // ///Propagate to virtual plane at PCA ------------------------------------
      // TVector3 oc = (0,0,0);
      // TVector3 dj(0,1,0);
      // TVector3 dk(-1,0,0);
      // FairTrackParP *fStart = 
      // 	new (clref1[size1]) FairTrackParP(StartPos, StartMom, StartPosErr, StartMomErr, fCharge, oc, dj, dk);
      // TClonesArray& clref = *fTrackParGeane;
      // Int_t size = clref.GetEntriesFast();
      // FairTrackParP *fRes = new(clref[size]) FairTrackParP();
      // fPro->SetPoint(vtx);
      // fPro->BackTrackToVirtualPlaneAtPCA(1);
      // Bool_t isProp =	fPro->Propagate(fStart, fRes, PDGCode);
      // cout<<"================= %%%% ===================="<<endl;
      // ///----------------------------------------------------------------------

    
    
     

      //------- TEST of calculation errors -------
      TVector3 gPos(fRes->GetX(),fRes->GetY(),fRes->GetZ());
      TVector3 gMom(fRes->GetPx(),fRes->GetPy(),fRes->GetPz());
      TVector3 gErrPos(fRes->GetDX(),fRes->GetDY(),fRes->GetDZ());
      TVector3 gErrMom(fRes->GetDPx(),fRes->GetDPy(),fRes->GetDPz());
      // cout<<" "<<endl;
      // cout<<"================= %%%% ===================="<<endl;
      
      // cout<<"gPos:"<<endl;
      // gPos.Print();
      // // cout<<"difference between final and initial point = "<<(-StartPos.Mag()+gPos.Mag())<<endl;
      // cout<<"gErrPos:"<<endl;
      // gErrPos.Print();
      // //  cout<<"difference between final and initial point error = "<<(-StartPosErr.Mag()+gErrPos.Mag())<<endl;

      // cout<<"gMom: "<<gMom.Mag()<<endl;
      // gMom.Print();   
      // // cout<<"difference between initial and final momentum = "<<(StartMom.Mag()-gMom.Mag())<<endl;
      // cout<<"gErrMom: "<<gErrMom.Mag()<<endl;
      // // cout<<"difference between initial and final momentum error= "<<(StartMomErr.Mag()-gErrMom.Mag())<<endl;
      // gErrMom.Print();
      // cout<<"================= %%%% ===================="<<endl;
      //------------------------------------------
      
      if(isProp == kTRUE){
	new((*fTrackParFinal)[counterGeaneTrk]) FairTrackParH(*(fRes)); //save Track
	//	new((*fTrackParFinal)[counterGeaneTrk]) FairTrackParP(*(fRes)); //save Track
	counterGeaneTrk++;
	//cout<<"***** isProp TRUE *****"<<endl;
      }
      else{
	cout<<"!!! Back-propagation with GEANE didn't return result !!!"<<endl;
	cout<<"StartPos:"<<endl;
	StartPos.Print();
	cout<<"StartPosErr:"<<endl;
	StartPosErr.Print();
	
	cout<<"StartMom: "<<StartMom.Mag()<<endl;
	StartMom.Print();
	cout<<"StartMomErr: "<<StartMomErr.Mag()<<endl;
	StartMomErr.Print();  
	new((*fTrackParFinal)[counterGeaneTrk]) FairTrackParH(); //save NULL
	//	new((*fTrackParFinal)[counterGeaneTrk]) FairTrackParP(); //save NULL
	counterGeaneTrk++;
      }
  }
  fMCTracks->Delete();
  fMCHits->Delete();
  //  cout<<"PndLmdGeaneTask::Exec END!"<<endl;
}

void PndLmdGeaneTask::Finish()
{
}

std::map<int, std::vector<int> > PndLmdGeaneTask::AssignHitsToTracks()
{
	std::map<int, std::vector<int> > result;
	for (int i = 0; i < fMCHits->GetEntriesFast(); i++){											//get all MC Hits
		PndSdsMCPoint* myPoint = (PndSdsMCPoint*)(fMCHits->At(i));									//sort MCHits with Tracks
		//PndMCTrack* myTrack = (PndMCTrack*)(fMCTracks->At(myPoint->GetTrackID()));
		result[myPoint->GetTrackID()].push_back(i);

	}
	return result;
}

ClassImp(PndLmdGeaneTask);