//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//     Analysis task for lambdas
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sverre Doerheim
//
//-----------------------------------------------------------

// Panda Headers ----------------------

// This Class' Header ------------------
#include "FopiLambdaAnaTask.h"
#include "FopiLambdaCand.h"

// C/C++ Headers ----------------------
#include <iostream>
#include <vector>
#include <map>
#include <assert.h>

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "TClonesArray.h"
#include "GFTrack.h"
#include "GFDetPlane.h"
#include "GFFieldManager.h"
#include "GFKalman.h"

#include "GFRecoHitFactory.h"
#include "GFException.h"

#include "GFDetPlane.h"

#include "RKTrackRep.h"
//#include "FairGeanePro.h"

#include "PndDetectorList.h"

#include "PndFieldAdaptor.h"
#include "FairField.h"
#include "FairRunAna.h"
#include "GFTrackProximity.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TH2I.h"
#include "TLorentzVector.h"
#include "TDatabasePDG.h"



// Class Member definitions -----------
using namespace std;


FopiLambdaAnaTask::FopiLambdaAnaTask()
  : FairTask("FopiLambda"), fPersistence(kFALSE), farmCutMaxPt(0)
{
  fTrackBranchName = "TrackPostFitComplete";
  fVerbose = 0;

} 


FopiLambdaAnaTask::~FopiLambdaAnaTask()
{

}

InitStatus
FopiLambdaAnaTask::Init(){
  std::cerr<<"FopiLambdaAnaTask::Init()"<<std::endl;
  //Get ROOT Manager
  FairRootManager* ioman= FairRootManager::Instance();
  
  if(ioman==0) {
    Error("FopiLambdaAnaTask::Init","RootManager not instantiated!");
    return kERROR;
  }
  
  // Get input collection
  fTrackArray=(TClonesArray*) ioman->GetObject(fTrackBranchName);

  if(fTrackArray==0) {
    Error("FopiLambdaAnaTask::Init","track-array not found!");
    return kERROR;
  }
  
  vtxResX_1=new TH1D("vtxRexX_1","vtxRexX_1",2000,-5,5);
  vtxResY_1=new TH1D("vtxRexY_1","vtxRexY_1",2000,-5,5);
  vtxResZ_1=new TH1D("vtxRexZ_1","vtxRexZ_1",2000,-5,5);

  vtxResX_2=new TH1D("vtxRexX_2","vtxRexX_2",2000,-5,5);
  vtxResY_2=new TH1D("vtxRexY_2","vtxRexY_2",2000,-5,5);
  vtxResZ_2=new TH1D("vtxRexZ_2","vtxRexZ_2",2000,-5,5);
  
  vtxDist=new TH1D("vtxDist","vtxDist",2000,0,20);

  vtxXY=new TH2D("vtxXY","vtxXY",2000,-20,20,2000,-20,20);
  vtxZX=new TH2D("vtxZX","vtxZX",2000,-100,100,2000,-20,20);
  vtxZY=new TH2D("vtxZY","vtxZY",2000,-100,100,2000,-20,20);

  for(int i=0;i<2;++i){
    lambdaMass.push_back(new TH1D(TString::Format("lamdba_mass_%i",i),TString::Format("lambda_mass_%i",i),5000,0,3));
    nLambdaCands.push_back(new TH1D(TString::Format("n_lambda_cands_%i",i),TString::Format("n_lambda_cands_%i",i),10,0,10));
    //Combinatorics
    armenteros.push_back(new TH2D(TString::Format("armenteros_all_combi_%i",i),TString::Format("armenteros_all_combi_%i",i),200,-1,1,200,0,1.15));
  }
  pdg=new TH1D("pdg","pdg",10000,-5000,4999);
  charge=new TH1D("charge","charge",10000,-5000,4999);
  cout<<"FopiLambdaAnaTask::Init 1"<<endl;
  FairField* field=FairRunAna::Instance()->GetField();
  cout<<"FopiLambdaAnaTask::Init 2"<<endl;
  GFFieldManager::getInstance()->init(new PndFieldAdaptor(field));
  cout<<"FopiLambdaAnaTask::Init 3"<<endl;
  farmCutMaxPt=0.11;
  farmCutMaxAlpha=0.7+0.25;
  farmCutMinAlpha=0.7-0.25;
  vtxCut=1;
  nEv=0;
  cout<<"FopiLambdaAnaTask::Init 4"<<endl;
  fLambdaOutArray = new TClonesArray("FopiLambdaCand");
  cout<<"FopiLambdaAnaTask::Init 5"<<endl;
  ioman->Register("FopiLambdaCand","Tpc",fLambdaOutArray,true);
  fPDG=TDatabasePDG::Instance();
  return kSUCCESS;
}


void
FopiLambdaAnaTask::Exec(Option_t* opt){
  fLambdaOutArray->Delete();
  //if(nEv%1000==0)
    std::cerr<<"FopiLambdaAnaTask::Exec "<<nEv<<"########################################"<<std::endl;
  nEv++;
  unsigned int ntracks = fTrackArray->GetEntriesFast();
  vector<GFTrack*> protons;
  vector<GFTrack*> pi_min;

  std::map<int,std::vector<TLorentzVector > > l_protons;
  std::map<int,std::vector<TLorentzVector > >  l_pi_min;

  //--------------------------------------------------------------------------------
  //Looping over tracks and filling particle lists
  //--------------------------------------------------------------------------------
  //  cout<<"nTracks: "<<ntracks<<endl;
  for(unsigned int itr=0;  itr<ntracks; itr++) {
    //    std::cout<<"  ... processing TPC track no. "<<itr<<std::endl;
    GFTrack* track = (GFTrack*) (*fTrackArray)[itr];
    //GFTrackCand cand= track->getCand();
    GFAbsTrackRep* rep = track->getCardinalRep();

    pdg->Fill(rep->getPDG());
    charge->Fill(rep->getCharge());
    TLorentzVector p;
    GFKalman fitter;
    // fitter.setNumIterations(3);
    // if(rep->getPDG()==2212||rep->getPDG()==-211){
    //   try{
    // 	fitter.processTrack(track);
    //   }catch(GFException e){

    //   }
    // }
    switch(rep->getPDG()){
    case 2212:
      protons.push_back(track);
      p.SetVectM(track->getMom(),fPDG->GetParticle(2212)->Mass());
      l_protons[0].push_back(p);
      break;
    case -211:
      pi_min.push_back(track);
      p.SetVectM(track->getMom(),fPDG->GetParticle(211)->Mass());
      l_pi_min[0].push_back(p);
      break;
    default:
      break;
    }

  }
  // cout<<"#p:"<<protons.size()<<endl;
  // cout<<"#ap:"<<antiProtons.size()<<endl;
  // cout<<"#pi-:"<<pi_min.size()<<endl;
  // cout<<"#pi+:"<<pi_plus.size()<<endl;
  //--------------------------------------------------------------------------------
  //combining lambda candidates
  //--------------------------------------------------------------------------------

                
  //loping over cases
  //0 Mom from raw tracks


  int nlcands=0;
  double l_mass=1.115684;
  TLorentzVector bestLCand;
  double bestMassDiff=99999;

  TVector3 lVtxResBest;

  for(int p=0;p<l_protons[0].size();++p){
    for(int pi=0;pi<l_pi_min[0].size();++pi){
      //      cout<<"p:"<<p<<" pi:"<<pi<<endl;
      bool failed=false;
      TVector3 pocaTr,pocaP,pocaPi;
      failed=Poca(protons.at(p), pi_min.at(pi), 
		  l_protons[0].at(p),l_pi_min[0].at(pi), 
		  pocaTr,pocaP,pocaPi);
      TVector3 dist = pocaP-pocaPi;
      // cout<<"dist: ";dist.Print();
      // cout<<"pocaPi: ";pocaPi.Print();
      // cout<<"pocaP: ";pocaP.Print();
      
      // vtxDist->Fill(dist.Mag());
      // cout<<"failed"<<failed<<endl;
      for(int i=0;i<2;++i){
	if(!failed) 	  {
	  //	  cout<<"i__:"<<i<<endl;
	  TLorentzVector lambda=l_pi_min[0].at(pi)+l_protons[0].at(p);
	  double pt, alpha;
	  if(i==0){
            //            cout<<"i==0"<<endl;
	    //RKTrackRep* lRep=new RKTrackRep(pocaTr,lambda.Vect(),3122);
            //            cout<<"rk init"<<endl;
            //            pocaTr.Print();
            //            lambda.Vect().Print();
	    TVector3 a(0,0,-40);
	    TVector3 b(0,0,100);
	    TVector3 pocaZL(0,0,0),pocaZP,pocaZPi,normVec,poca_onwire;
	    try{
              //              cout<<"rk extrap"<<endl;
              //	      lRep->extrapolateToLine(a,b,pocaZL,normVec,poca_onwire);
	      protons.at(p)->getCardinalRep()->extrapolateToLine(a,b,pocaZP,normVec,poca_onwire);
	      pi_min.at(pi)->getCardinalRep()->extrapolateToLine(a,b,pocaZPi,normVec,poca_onwire);
	    } catch(GFException &exp){
	      cout<<exp.what()<<endl;
	      failed = true;
	    } 
	  
	    if(!failed){
              TMatrixT<double> err;
              err.ResizeTo(3,3);
              err.Zero();
              err[0][0]=dist.Mag();

	      FopiLambdaCand* cand=new((*fLambdaOutArray)[fLambdaOutArray->GetEntriesFast()]) FopiLambdaCand(lambda,
                                                                                                             pocaTr);
	      FillArmenterosPlot(lambda,l_protons[i].at(p),l_pi_min[i].at(pi),pt,alpha);
	      lambdaMass.at(i)->Fill(lambda.M());
	      armenteros.at(i)->Fill(alpha,pt);
	      
	    }
	  }
	  nlcands++;
	  if(i==1) {
	    //	    cout<<"teeeeeest"<<dist.Mag()<<":"<<vtxCut<<endl;
	    if(dist.Mag()<vtxCut){
              //              cout<<"i==1"<<endl;
	      FillArmenterosPlot(lambda,l_protons[0].at(p),l_pi_min[0].at(pi),pt,alpha);
	      lambdaMass.at(i)->Fill(lambda.M());
	      armenteros.at(i)->Fill(alpha,pt);
	    }
	  }
	}
      }
    }//end looping over pions
  }//end looping over protons


  nLambdaCands.at(0)->Fill(nlcands);
    


  return;
}
void FopiLambdaAnaTask::FillArmenterosPlot(TLorentzVector lambda,
                                           TLorentzVector pos,
                                           TLorentzVector neg,
                                           double& PT,
                                           double& alpha){
  TVector3 z=lambda.Vect().Unit();
  TVector3 dirPos=pos.Vect().Unit();
  TVector3 dirNeg=neg.Vect().Unit();
  
  TVector3 y=dirPos.Cross(dirNeg).Unit();
  TVector3 x=y.Cross(z).Unit();
  TRotation *rot=new TRotation();

  rot->RotateAxes(x,y,z);
  rot->Invert();
  TLorentzVector pos_rot(pos);
  TLorentzVector neg_rot(neg);
  TLorentzVector lambda_rot(lambda);
  pos_rot.Transform(*rot);
  neg_rot.Transform(*rot);
  lambda_rot.Transform(*rot);
  alpha=(pos_rot.Pz()-neg_rot.Pz())/(pos_rot.Pz()+(neg_rot.Pz()));
  PT=pos_rot.Pt();
  delete rot;
}

void FopiLambdaAnaTask::WriteHistograms(const TString& fname) const {
  TFile* rOut = new TFile(fname, "recreate");
  cout<<"pdg-write"<<endl;
  pdg->Write();
  cout<<"charge-write"<<endl;
  charge->Write();
  //cout<<"vtx2"<<endl;

  cout<<"vtxX2-write"<<endl;;
  vtxResX_2->Write();
  cout<<"vtxY2-write"<<endl;
  vtxResY_2->Write();
  cout<<"vtxZ2-write"<<endl;
  vtxResZ_2->Write();

  vtxXY->Write();
  vtxZX->Write();
  vtxZY->Write();
  vtxDist->Write();

  vtxResX_1->Write();
  cout<<"vtxX1-write"<<endl;
  vtxResY_1->Write();
  cout<<"vtxY1-write"<<endl;
  vtxResZ_1->Write();
  cout<<"vtxZ1-write"<<endl;

  for(int i=0;i<2;++i){
    cout<<i<<endl;
    lambdaMass.at(i)->Write();
    cout<<"test1"<<endl;
    nLambdaCands.at(i)->Write();
    cout<<"test2"<<endl;
    armenteros.at(i)->Write();
    cout<<"test10"<<endl;
  }
 


  rOut->cd();
  rOut->Close();
}



bool FopiLambdaAnaTask::Poca(GFTrack* proton,
                             GFTrack* pion,
			     TLorentzVector &l_proton,
                             TLorentzVector &l_pion,
                             TVector3& pocaTr,
                             TVector3& pocaP,
                             TVector3& pocaPi){
  cout<<"FopiLambdaAnaTask::Poca"<<endl;
  RKTrackRep*  proton_rep =new RKTrackRep(*((RKTrackRep* )proton->getCardinalRep()));
  RKTrackRep*  pion_rep   =new RKTrackRep(*((RKTrackRep* )pion->getCardinalRep()));
  TVector3 poca1,norm1;
  TVector3 poca2,norm2;
  GFDetPlane fPlane1,fPlane2;
  bool failed =false;
  try{
    TVector3 startPoint1=proton_rep->GFAbsTrackRep::getPos();
    TVector3 startPoint2=pion_rep->GFAbsTrackRep::getPos();
    
    TVector3 dist1=startPoint2-startPoint1;

    vtxResX_1->Fill(dist1.X());
    vtxResY_1->Fill(dist1.Y());
    vtxResZ_1->Fill(dist1.Z());
    startPoint1.Print();
    startPoint2.Print();
    pocaTr=trackProximity((GFAbsTrackRep*)proton_rep,(GFAbsTrackRep*)pion_rep);
    
    proton_rep->extrapolateToPoint(pocaTr, poca1, norm1);
    fPlane1.setO(pocaTr);
    fPlane1.setNormal(norm1);
    pion_rep->extrapolateToPoint(pocaTr, poca2, norm2);
    fPlane2.setO(pocaTr);
    fPlane2.setNormal(norm2);

    TVector3 mom1=proton_rep->getMom(fPlane1);
    TVector3 mom2=pion_rep->getMom(fPlane2);
    l_proton.SetVectM(mom1,fPDG->GetParticle(2212)->Mass());
    l_pion.SetVectM(mom2,fPDG->GetParticle(211)->Mass());
      
    startPoint1=proton_rep->GFAbsTrackRep::getPos();
    startPoint2=pion_rep->GFAbsTrackRep::getPos();
    
    vtxXY->Fill(pocaTr.X(),pocaTr.Y());
    vtxZX->Fill(pocaTr.Z(),pocaTr.X());
    vtxZY->Fill(pocaTr.Z(),pocaTr.Y());

    dist1=startPoint2-startPoint1;
    pocaP=poca1;
    pocaPi=poca2;
    vtxResX_2->Fill(dist1.X());
    vtxResY_2->Fill(dist1.Y());
    vtxResZ_2->Fill(dist1.Z());
    

    delete proton_rep;
    delete  pion_rep  ;
  }
  catch(GFException &exp){
    cout<<exp.what()<<endl;
    failed = true;
  }
  catch(...){
    cout<<"Noooooooooooooooooooooooooooooooooooooooooooooooooo"<<endl;
  }
  return failed;
}