//ROOT includes
#include <TFile.h>
#include <TLorentzVector.h>
#include <TStopwatch.h>

//std-lib
#include <iostream>

//collaborating classes
#include "MatchingTuple.h"
#include "TpcdEdx.h"
#include "FopiPidHub.h"
#include "FopiPidInfo.h"

//FOPI2ROOT
#include "CdcTrack.h"
#include "RpcTrack.h"
#include "BarrelTrack.h"
//Genfit
#include <GFRaveConverters.h>
#include <GFRaveTrackParameters.h>
#include <GFRaveVertex.h>
#include <GFRaveVertexFactory.h>
#include <GFException.h>
#include <GFTrack.h>
#include <GFAbsTrackRep.h>

#include <GFTrack.h>
#include <GFAbsTrackRep.h>
#include <RKTrackRep.h>
//header
#include "V0Chain.h"

V0Chain::V0Chain(TChain *chain_, TString outFile_):
    trackBranchName("FopiTrackPostFit"),
    matchingTupleBranchName("TpcCdcMatchingPair"),
    trackTupleBranchName("FopiTrackTuples"),
    dEdxBranchName("dEdx"),
    rpcBranchName("RpcTrack"),
    cdcBranchName("CdcTrack"),
    barrelBranchName("BarrelTrack"),
    outFileName(outFile_),
    fVerbose(1),
    vanillaPID(true)
        
{
    chain=chain_;
    nEvents=0;
    trackArray=new TClonesArray("GFTrack");
    matchingTupleArray=new TClonesArray("MatchingTuple");
    trackTupleArray=new TClonesArray("MatchingTuple");
    dEdxArray=new TClonesArray("TpcdEdx");
    rpcArray=new TClonesArray("RpcTrack");
    cdcArray=new TClonesArray("CdcTrack");
    barrelArray=new TClonesArray("BarrelTrack");
    arrays.push_back(trackArray);
    arrays.push_back(matchingTupleArray);
    arrays.push_back(trackTupleArray);
    arrays.push_back(dEdxArray);
    arrays.push_back(rpcArray);
    //arrays.push_back(barrelArray);
    arrays.push_back(cdcArray);
}


void V0Chain::Init(){
    std::cout<<"V0Chain::Init()"<<std::endl;
    chain->SetBranchStatus("*",0);
    chain->SetBranchStatus(trackBranchName+".*",1);
    chain->SetBranchStatus(trackTupleBranchName+".*",1);
    //chain->SetBranchStatus(matchingTupleBranchName+".*",1);
    chain->SetBranchStatus(dEdxBranchName+".*",1);
    chain->SetBranchStatus(rpcBranchName+".*",1);
    //chain->SetBranchStatus(barrelBranchName+".*",1);
    chain->SetBranchStatus(cdcBranchName+".*",1);
    std::cout<<"chain->SetBranchAddress(cdcBranchName,&cdcArray);"<<std::endl;
    chain->SetBranchAddress(cdcBranchName,&cdcArray);
    std::cout<<"chain->SetBranchAddress(trackBranchName,&trackArray);"<<std::endl;
    chain->SetBranchAddress(trackBranchName,&trackArray);
    std::cout<<"chain->SetBranchAddress(matchingTupleBranchName,&matchingTupleArray);"<<std::endl;
    chain->SetBranchAddress(matchingTupleBranchName,&matchingTupleArray);
    std::cout<<"chain->SetBranchAddress(trackTupleBranchName,&trackTupleArray);"<<std::endl;
    chain->SetBranchAddress(trackTupleBranchName,&trackTupleArray);
    std::cout<<"chain->SetBranchAddress(dEdxBranchName,&dEdxArray);"<<std::endl;
    chain->SetBranchAddress(dEdxBranchName,&dEdxArray);
    std::cout<<"chain->SetBranchAddress(rpcBranchName,&rpcArray);"<<std::endl;
    chain->SetBranchAddress(rpcBranchName,&rpcArray);
    //chain->SetBranchAddress(barrelBranchName,&barrelArray);
    std::cout<<"opening "<<outFileName<<" as output file"<<std::endl;
    outFile=TFile::Open(outFileName,"recreate");
    outFile->cd();
    k0Tuple=new TNtuple( "k0tuple","k0Tuple","iev:vx:vy:vz:pocaX:pocaY:pocaZ:m:px:py:pz");
    lambdaTuple=new TNtuple( "lambdaTuple","lambdaTuple","iev:vx:vy:vz:pocaX:pocaY:pocaZ:m:px:py:pz");
    
    hub = new FopiPidHub();
    hub->setDetProbFile("tpc","/scratch/sdorheim/FOPI/newDev/tpc/FOPI/par/tpcPidTest2.root");
    hub->setDetProbFile("bar","/scratch/sdorheim/FOPI/newDev/tpc/FOPI/par/barPidTest2.root");
    hub->setDetProbFile("rpc","/scratch/sdorheim/FOPI/newDev/tpc/FOPI/par/rpcPidTest2.root");
    //Use FopiPidFractionSets from file to use MC fractions
    hub->setUseMCFractions("/scratch/sdorheim/FOPI/newDev/tpc/FOPI/par/MCFractionsTest.root");
    //Use calibration for a given detector, defining a reference region
    hub->setCalibrationFile("tpc", "/scratch/sdorheim/FOPI/newDev/tpc/FOPI/par/testCal2.root",
                           "tpcCalPion", "tpcCalProton",
                           3280, 3292);
    std::cout<<"hub: init() "<<hub->init()<<std::endl;
    
}


void V0Chain::Exec(){
    std::cout<<"V0Chain::Exec()"<<std::endl;
    int nEntries = chain->GetEntries();
    if(nEvents==0){
        nEvents=nEntries;
    }else if(nEvents>nEntries){
        nEvents=nEntries;
    }
    std::cout<<"Chain with "<<nEntries<<" events loaded, looping over "<<nEvents<<" events"<<std::endl;
    double CDCCONVERSION = 1./1.5e3;
    double TPCCONVERSION = 1.e-3;
    TDatabasePDG* fPDG=TDatabasePDG::Instance();
    double pionMass=fPDG->GetParticle(211)->Mass();
    double protonMass=fPDG->GetParticle(2212)->Mass();
    //Event loop
    int nPiPlus=0;
    int nPiMinus=0;
    int nProton=0;
    int nK0=0;
    int nLambda=0;
    TStopwatch clock;
    
    clock.Start();
    for(unsigned int i_ev=0;i_ev<nEvents;++i_ev){
        if(fVerbose){
            if(i_ev%1000==0){
                std::cout<<"."<<std::flush;
            }
            if(i_ev%10000==0){
                std::cout<<"------------------- "<<i_ev<<" ---------------------------------"<<std::endl;
                clock.Print();
                clock.Continue();
            }
        }
        
        chain->GetEvent(i_ev);
        std::vector<GFTrack*> pi_plus;
        std::vector<GFTrack*> pi_minus;
        std::vector<GFTrack*> protons;
        

        
        
        unsigned int nTracks=trackArray->GetEntries();
        for(unsigned i_tr=0;i_tr<nTracks;++i_tr){
            GFTrack* theTrack = (GFTrack*)trackArray->UncheckedAt(i_tr);
            if(theTrack->getCardinalRep()->getStatusFlag()!=0){
                continue;
            }
            double combMom =theTrack->getMom().Mag();
            double cdcMom=0;
            double charge=theTrack->getCharge();
            if(combMom<0.2 or combMom>0.8 ){
                   continue;
            }
            TpcdEdx* dEdx = (TpcdEdx*)dEdxArray->At(i_tr);
            MatchingTuple* mp = (MatchingTuple*) trackTupleArray->At(i_tr);
            CdcTrack* theCdcTrack= NULL;
            if(mp->isMatched()){ 
                //Track has CDC match
                if(mp->hasBranch(cdcBranchName)){
                     theCdcTrack= (CdcTrack*) cdcArray->At(mp->getIdFromBranch(cdcBranchName));
                     cdcMom=theCdcTrack->GetMom();
                }
            }
            if(vanillaPID){
                if(theCdcTrack!=NULL){
                    double mass = theCdcTrack->GetMass();
                    if(theCdcTrack->GetCharge()>0){
                        if(mass>0.8&&mass<1.2){
                            protons.push_back(theTrack);
                        }else if(mass>0.1&&mass<0.3){
                            pi_plus.push_back(theTrack);
                        }
                    }else{
                        if(mass>0.1&&mass<0.3){
                            pi_minus.push_back(theTrack);
                        }
                    }
                }
            }else{
                //--------------------------------------------------------------------------------
                //Building PidInfo
                //--------------------------------------------------------------------------------
                FopiPidInfo pidInfo(3280);
                pidInfo.setTpcdEdx(dEdx->truncMean()*TPCCONVERSION);
                if(theCdcTrack!=NULL){
                    pidInfo.setCdcdEdx(theCdcTrack->GetEloss()*CDCCONVERSION);
                    if(theCdcTrack->GetRpcHitID()!=-1){
                        pidInfo.setRpcVel(theCdcTrack->GetVel());
                    }else if (theCdcTrack->GetBarHitID()){
                        pidInfo.setBarVel(theCdcTrack->GetVel());
                    }
                }
                double pionFrac=0;
                double protonFrac=0;
                if(cdcMom>0){
                    pionFrac=hub->getFracLikelihood("pion",cdcMom,pidInfo);
                    protonFrac=hub->getFracLikelihood("proton",cdcMom,pidInfo);
                }
                //--------------------------------------------------------------------------------
                // Building particle lists
                //--------------------------------------------------------------------------------
                if(pionFrac>0.5){
                    if(charge>0){
                        pi_plus.push_back(theTrack);
                        nPiPlus++;
                    }else if (charge<0){
                        pi_minus.push_back(theTrack);
                        nPiMinus++;
                    }
                    if(protonFrac>0.5){
                        std::cout<<"proton&pion>.5, p_frac: "<<protonFrac<<", pionFrac: "<<pionFrac<<std::endl;
                    }
                }else if(protonFrac>0.5){
                    if(charge>0){
                        protons.push_back(theTrack);
                        nProton++;
                    }else{
                        std::cout<<"negative proton, p_frac: "<<protonFrac<<", pionFrac: "<<pionFrac<<std::endl;
                    }
                }
            }
           // std::cout<<"iTr: "<<i_tr<<" p: "<<protonFrac<<" pi: "<<pionFrac<<std::endl;
        }//end loop over tracks
        //--------------------------------------------------------------------------------
        // Building track pairs for K0 and Lambda candidates
        //--------------------------------------------------------------------------------
        //K0
        std::vector<std::vector<GFTrack*> > k0TrackPairs;
        for(unsigned int ipp=0;ipp<pi_plus.size();++ipp){
            for(unsigned int ipm=0;ipm<pi_minus.size();++ipm){
                std::vector<GFTrack*> trackPair;
                trackPair.push_back(pi_plus[ipp]);
                trackPair.push_back(pi_minus[ipm]);
                k0TrackPairs.push_back(trackPair);
            }
        }
        //Lambda
        std::vector<std::vector<GFTrack*> > lambdaTrackPairs;
        for(unsigned int ip=0;ip<protons.size();++ip){
            for(unsigned int ipm=0;ipm<pi_minus.size();++ipm){
                std::vector<GFTrack*> trackPair;
                trackPair.push_back(protons[ip]);
                trackPair.push_back(pi_minus[ipm]);
                lambdaTrackPairs.push_back(trackPair);
            }
        }
        
        if(fVerbose>10){
            std::cout<<i_ev
                    <<", pp: "<<pi_plus.size()
                   <<", pm: "<<pi_minus.size()
                  <<", p: "<<protons.size()
                 <<", k0: "<<k0TrackPairs.size()
                <<", lambda: "<<lambdaTrackPairs.size()<<std::endl;
        }
        if(k0TrackPairs.size()>0 or lambdaTrackPairs.size()>0){

            
            //--------------------------------------------------------------------------------
            // Rave fitting 
            //--------------------------------------------------------------------------------
            
            TMatrixDSym beamCov;
            beamCov.ResizeTo(3,3);
            beamCov.UnitMatrix();
            beamCov *= 2;
            GFRaveVertexFactory* vertexFactorySmooth = new GFRaveVertexFactory(0, false);
            vertexFactorySmooth->setBeamspot(TVector3(0,0,0),beamCov);
            vertexFactorySmooth->setMethod("kalman-smoothing:1");
            std::vector<GFRaveVertex*> k0VerticesSmooth;
            std::vector<GFRaveVertex*> lambdaVerticesSmooth;
            std::vector<GFRaveVertex*> *fVertexBuffer = new std::vector < GFRaveVertex* >;
            
            //k0
             if(fVerbose>10){
                 std::cout<<"vertex k0"<<std::endl;
             }
            for(unsigned int ik0=0;ik0<k0TrackPairs.size();++ik0){
                fVertexBuffer->clear();
                try{
                    vertexFactorySmooth->findVertices(fVertexBuffer, k0TrackPairs[ik0], true);
                }catch(GFException exp){
                    std::cout<<exp.what()<<std::endl;
                }

                if(fVertexBuffer->size()==0){
                    k0VerticesSmooth.push_back(NULL);
                }else{
                    GFRaveVertex* theVtx=fVertexBuffer->back();
                    k0VerticesSmooth.push_back(theVtx);
                }
            }
            //lambda
             if(fVerbose>10){
                 std::cout<<"vertex lambda"<<std::endl;
             }
            for(unsigned int iLambda=0;iLambda<lambdaTrackPairs.size();++iLambda){
                fVertexBuffer->clear();
                try{   
                    vertexFactorySmooth->findVertices(fVertexBuffer, lambdaTrackPairs[iLambda], true);
                }catch(GFException exp){
                    std::cout<<exp.what()<<std::endl;
                }
                if(fVertexBuffer->size()==0){
                    lambdaVerticesSmooth.push_back(NULL);
                }else{
                    GFRaveVertex* theVtx=fVertexBuffer->back();
                    lambdaVerticesSmooth.push_back(theVtx);
                }
            }
            delete vertexFactorySmooth;
            delete fVertexBuffer;
            
            //--------------------------------------------------------------------------------
            // calculating dist in poca
            //--------------------------------------------------------------------------------
                std::vector<TVector3> k0Poca;
                std::vector<TVector3> lambdaPoca;
                //k0
                if(fVerbose>10){
                    std::cout<<"POCA k0"<<std::endl;
                }
                for(unsigned int ik0=0;ik0<k0TrackPairs.size();++ik0){
                    if(k0TrackPairs[ik0].size()==2 && 
                            k0VerticesSmooth.at(ik0)!=NULL){
                        try{
                            TVector3 vtxPos=k0VerticesSmooth.at(ik0)->getPos();
                            TVector3 pos1(0,0,0);
                            TVector3 pos2(0,0,0);
                            TVector3 dummy(0,0,0);
                            GFAbsTrackRep* rep1=k0TrackPairs[ik0][0]->getCardinalRep();
                            GFAbsTrackRep* rep2=k0TrackPairs[ik0][1]->getCardinalRep();
                            rep1->extrapolateToPoint(vtxPos,pos1,dummy);
                            rep2->extrapolateToPoint(vtxPos,pos2,dummy);
                            k0Poca.push_back(pos1-pos2);
                        }catch(GFException exp){
                            k0Poca.push_back(TVector3(0,0,0));
                        }
                    }else{
                            k0Poca.push_back(TVector3(0,0,0));
                    }
                }
                if(fVerbose>10){
                    std::cout<<"POCA lambda"<<std::endl;
                }
                //lambda
                for(unsigned int iLambda=0;iLambda<lambdaTrackPairs.size();++iLambda){
                    if(lambdaTrackPairs[iLambda].size()==2 && 
                            lambdaVerticesSmooth.at(iLambda)!=NULL){
                        try{
                            TVector3 vtxPos=lambdaVerticesSmooth.at(iLambda)->getPos();
                            TVector3 pos1(0,0,0);
                            TVector3 pos2(0,0,0);
                            TVector3 dummy(0,0,0);
                            GFAbsTrackRep* rep1=lambdaTrackPairs[iLambda][0]->getCardinalRep();
                            GFAbsTrackRep* rep2=lambdaTrackPairs[iLambda][1]->getCardinalRep();
                            rep1->extrapolateToPoint(vtxPos,pos1,dummy);
                            rep2->extrapolateToPoint(vtxPos,pos2,dummy);
                            lambdaPoca.push_back(pos1-pos2);
                        }catch(GFException exp){
                            lambdaPoca.push_back(TVector3(0,0,0));
                        }
                    }else{
                        lambdaPoca.push_back(TVector3(0,0,0));
                    }
                }
            
            //--------------------------------------------------------------------------------
            // Adding 4-vectors
            //--------------------------------------------------------------------------------
            std::vector<TLorentzVector> k0_4vec;
            std::vector<TLorentzVector> lambda_4vec;
            //k0
             if(fVerbose>10){
                 std::cout<<"4vec k0"<<std::endl;
             }
            for(unsigned int ik0=0;ik0<k0TrackPairs.size();++ik0){
                GFRaveVertex* theVtx = k0VerticesSmooth[ik0];
                if(theVtx!=NULL){
                    GFRaveTrackParameters* tr1=theVtx->getParameters(0);
                    GFRaveTrackParameters* tr2=theVtx->getParameters(1);
                    TLorentzVector pi_plus_4;
                    TLorentzVector pi_minus_4;
                    if(tr1->getCharge()>0){
                        pi_plus_4.SetVectM(tr1->getMom(),pionMass);
                        pi_minus_4.SetVectM(tr2->getMom(),pionMass);
                    }else{
                        pi_plus_4.SetVectM(tr2->getMom(),pionMass);
                        pi_minus_4.SetVectM(tr1->getMom(),pionMass);
                    }
                    TLorentzVector k0=pi_plus_4+pi_minus_4;
                    k0_4vec.push_back(k0);
                }else{
                    k0_4vec.push_back(TLorentzVector());
                }
            }
            //lambda
             if(fVerbose>10){
                 std::cout<<"4vec lambda"<<std::endl;
             }
            for(unsigned int ilambda=0;ilambda<lambdaTrackPairs.size();++ilambda){
                GFRaveVertex* theVtx = lambdaVerticesSmooth[ilambda];
                if(theVtx!=NULL){
                    GFRaveTrackParameters* tr1=theVtx->getParameters(0);
                    GFRaveTrackParameters* tr2=theVtx->getParameters(1);
                    TLorentzVector proton_4;
                    TLorentzVector pi_minus_4;
                    if(tr1->getCharge()>0){
                        proton_4.SetVectM(tr1->getMom(),protonMass);
                        pi_minus_4.SetVectM(tr2->getMom(),pionMass);
                    }else{
                        proton_4.SetVectM(tr2->getMom(),protonMass);
                        pi_minus_4.SetVectM(tr1->getMom(),pionMass);
                    }
                    TLorentzVector lambda=proton_4+pi_minus_4;
                    lambda_4vec.push_back(lambda);
                }else{
                    lambda_4vec.push_back(TLorentzVector());
                }
            }
            //--------------------------------------------------------------------------------
            // Filling ntuples
            //--------------------------------------------------------------------------------
            //        k0Tuple=new TNtuple( "k0tuple","k0Tuple","iev:vx:vy:vz:poca:m:px:py:pz");
            //        lambdaTuple=new TNtuple( "lambdaTuple","lambdaTuple","iev:vx:vy:vz:poca:m:px:py:pz");
             if(fVerbose>10){
                 std::cout<<"ntuple k0"<<std::endl;
             }
            for(unsigned int ik0=0;ik0<k0TrackPairs.size();++ik0){
                GFRaveVertex* theVtx=k0VerticesSmooth[ik0];
                TVector3 vtxPos;
                TLorentzVector k0;
                if(theVtx!=NULL){
                    vtxPos=theVtx->getPos();
                    k0=k0_4vec.at(ik0);
                    delete theVtx;
                }
                TVector3  poca=k0Poca[ik0];
                //k0.Print();
                k0Tuple->Fill(i_ev,
                              vtxPos.X(),vtxPos.Y(),vtxPos.Z(),
                              poca.X(),poca.Y(),poca.Z(),
                              k0.M(),
                              k0.Vect().X(),k0.Vect().Y(),k0.Vect().Z());
                nK0++;
            }
             if(fVerbose>10){
                 std::cout<<"ntuple lambda "<<std::endl;
             }
            for(unsigned int ilambda=0;ilambda<lambdaTrackPairs.size();++ilambda){
                GFRaveVertex* theVtx=lambdaVerticesSmooth[ilambda];
                TVector3 vtxPos;
                TLorentzVector lambda;
                if(theVtx!=NULL){
                    vtxPos=theVtx->getPos();
                    lambda=lambda_4vec.at(ilambda);
                    delete theVtx;
                }
                TVector3 poca=lambdaPoca[ilambda];
                //lambda.Print();
                lambdaTuple->Fill(i_ev,
                                  vtxPos.X(),vtxPos.Y(),vtxPos.Z(),
                                  poca.X(),poca.Y(),poca.Z(),
                                  lambda.M(),
                                  lambda.Vect().X(),lambda.Vect().Y(),lambda.Vect().Z());
                nLambda++;
            }
        }// End Vertexing fun
    }//end loop over events
    outFile->cd();
    lambdaTuple->Write();
    k0Tuple->Write();
    outFile->Close();
    std::cout<<std::endl;
    std::cout<<"Pi_plus: "<<nPiPlus<<std::endl;
    std::cout<<"Pi_minus: "<<nPiMinus<<std::endl;
    std::cout<<"protons: "<<nProton<<std::endl;
    std::cout<<"Lambda: "<<nLambda<<std::endl;
    std::cout<<"k0: "<<nK0<<std::endl;
    
//    outTuple->Write();
//    outFile->Close();
}
ClassImp(V0Chain)