//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of resCalc
//      Uses existing GFTracks and calculates
//      residua in XY.
//      Developed as part of the GEM-TPC project
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Felix Boehmer                     (original author)
//      Physik Department E18, TUM
//
//-----------------------------------------------------------

#include "TpcRefGFTrkResCalc.h"
#include "TClonesArray.h"
#include "TpcCluster.h"
#include "TpcResidual.h"
#include "TpcRefResidualCollection.h"
#include "TpcSPHit.h"
#include "TpcAlignmentManager.h"

#include "GFTools.h"
#include "GFTrack.h"
#include "GFException.h"
#include "GFAbsTrackRep.h"
#include "GFRecoHitFactory.h"
#include "GFRecoHitProducer.h"
#include "RKTrackRep.h"

#include <FairRootManager.h>

#include <TClass.h>


#include <string>
#include <iostream>


ClassImp(TpcRefGFTrkResCalc)

    TpcRefGFTrkResCalc::TpcRefGFTrkResCalc()
      : fGFTrackBranchName("TpcTrackPostFit"),
      fTpcSPHitBranchName("TpcSPHit"),
      fRefGFTrackBranchName("CdcTrackPostFit"),
      fTpcClusterBranchName("TpcCluster"),
      fNRep(-1),
      fDetID(-1),
      fRequireGoodFit(true),
      fUseClusters(false),
      fTrackMomentumCut(0.5),
      fTrackPosCut(30.),
      fTrackdPhiCut(10),
      fAlMan(NULL),
      fQAFileName("TpcRefGFTrkResCalcQA.root")
{
  dPhiHist=new TH1D("dPhi","dPhi",1000,-360,360);
  hists.push_back(dPhiHist);
  dPhiVsPhiHist=new TH2D("dPhiVsPhi","dPhiVsPhi",1000,-360,360,1000,-360,360);
  hists.push_back(dPhiVsPhiHist);
  dPhiPosHist=new TH1D("dPhiPos","dPhiPos",1000,-360,360);
  hists.push_back(dPhiPosHist);
  dPhiVsPhiPosHist=new TH2D("dPhiVsPhiPos","dPhiVsPhiPos",1000,-360,360,1000,-360,360);
  hists.push_back(dPhiVsPhiPosHist);

  rRefHist=new TH1D("rRefHist","rRefHist",1000,0,150);
  hists.push_back(rRefHist);
  xyRefHist=new TH2D("xyRefHist","xyRefHist",1000,-150,150,1000,-150,150);
  hists.push_back(xyRefHist);
  momRefHist=new TH1D("momRefHist","momRefHist",1000,0,10);
  hists.push_back(momRefHist);
  phiRefHist=new TH1D("phiRefHist","phiRefHist",1000,-360,360);
  hists.push_back(phiRefHist);
  dPhiRefTrkPosHist=new TH1D("dPhiRefTrackPos","dPhiRefTrackPos",1000,-360,360);
  hists.push_back(dPhiRefTrkPosHist);

  rRefCutHist=new TH1D("rRefCutHist","rRefCutHist",1000,0,150);
  hists.push_back(rRefCutHist);
  xyRefCutHist=new TH2D("xyRefCutHist","xyRefCutHist",1000,-150,150,1000,-150,150);
  hists.push_back(xyRefCutHist);
  momRefCutHist=new TH1D("momRefCutHist","momRefCutHist",1000,0,10);
  hists.push_back(momRefCutHist);
  phiRefCutHist=new TH1D("phiRefCutHist","phiRefCutHist",1000,360,360);
  hists.push_back(phiRefCutHist);
  dPhiRefCutTrkPosHist=new TH1D("dPhiRefCutTrackPos","dPhiRefCutTrackPos",1000,-360,360);
  hists.push_back(dPhiRefCutTrkPosHist);
  dPhiRefCutTrkPosHist2=new TH1D("dPhiRefCutTrackPos_rot","dPhiRefCutTrackPos_rot",1000,-360,360);
  hists.push_back(dPhiRefCutTrkPosHist2);


  rHist=new TH1D("rHist","rHist",1000,0,150);
  hists.push_back(rHist);
  momHist=new TH1D("momHist","momHist",1000,0,10);
  hists.push_back(momHist);
  phiHist=new TH1D("phiHist","phiHist",1000,-360,360);
  hists.push_back(phiHist);
  dPhiTpcTrkPosHist=new TH1D("dPhiTpcTrackPos","dPhiTpcTrackPos",1000,-360,360);
  hists.push_back(dPhiTpcTrkPosHist);

  fNExpectedBranches = 4;
}

TpcRefGFTrkResCalc::~TpcRefGFTrkResCalc() {

}

bool TpcRefGFTrkResCalc::init() {
  if(fBranchMap.size() != fNExpectedBranches) {
    std::cout<<"TpcRefGFTrkResCalc::init(): "<<std::endl
        <<"   ERROR: Wrong number of input branches ("
        <<fBranchMap.size()<<")!"<<std::endl;
    return true;
  }
  //check for proper data types:
  bool foundGFTrack=false;
  bool foundRefGFTrack=false;
  bool foundSPHitArray=false;
  bool foundClusterArray=false;
  std::string GFTrackRefType("GFTrack");
  std::string RefGFTrackRefType("GFTrack");
  std::string tpcSPHitType("TpcSPHit");
  std::string tpcClusterType("TpcCluster");

  //identify the needed branches in the map
  //  std::map<TString, TClonesArray*>::const_iterator it;
  //  for(it=fBranchMap.begin(); it!=fBranchMap.end(); it++) {
  //    std::string type((it->second)->GetClass()->GetName());
  //    if(type.find(GFTrackRefType)<type.size()) {
  //      fGFTrackArray=it->second;
  //      fGFTrackBranchName = it->first;
  //      foundGFTrack=true;
  //      std::cout<<"TpcRefGFTrkResCalc::init(): "<<std::endl
  //          <<"   INFO: got "<<type<<" as Object type for GFTracks with branch name "<<it->first
  //          <<std::endl;
  //    }
  //    if(type.find(tpcClusterRefType)<type.size()) {
  //      fTpcSPHitArray=it->second;
  //      foundTpcCluster=true;
  //      std::cout<<"TpcRefGFTrkResCalc::init(): "<<std::endl
  //          <<"   INFO: got "<<type<<" as Object type for TPC Clusters with branch name "<<it->first
  //          <<std::endl;
  //    }
  //  }
  if(fBranchMap.count(fRefGFTrackBranchName)){
    std::cout<<"reftrack branch found:"<<fRefGFTrackBranchName<<std::endl;
    std::string type(fBranchMap[fRefGFTrackBranchName]->GetClass()->GetName());
    if(type.find(RefGFTrackRefType)<type.size()){
      fRefGFTrackArray=fBranchMap[fRefGFTrackBranchName];
      foundRefGFTrack=true;
    }
  }
  if(fBranchMap.count(fGFTrackBranchName)){
    std::cout<<"track branch found:"<<fRefGFTrackBranchName<<std::endl;
    std::string type(fBranchMap[fGFTrackBranchName]->GetClass()->GetName());
    if(type.find(GFTrackRefType)<type.size()){
      fGFTrackArray=fBranchMap[fGFTrackBranchName];
      foundGFTrack=true;
    }
  }
  if(fBranchMap.count(fTpcSPHitBranchName)){
    std::cout<<"SPHit branch found:"<<fTpcSPHitBranchName<<std::endl;
    std::string type(fBranchMap[fTpcSPHitBranchName]->GetClass()->GetName());
    if(type.find(tpcSPHitType)<type.size()){
      fTpcSPHitArray=fBranchMap[fTpcSPHitBranchName];
      foundSPHitArray=true;
    }
  }
  if(!(foundGFTrack&&foundSPHitArray&&foundRefGFTrack)) {
    std::cout<<"TpcRefGFTrkResCalc::init(): "<<std::endl
        <<"   ERROR: Wrong Class Types in input branches!"<<std::endl;
    return true;
  }

  //  fDetID = FairRootManager::Instance()->GetBranchId(fTpcClusterBranchName);
  //  fFactory->addProducer(fDetID, new GFRecoHitProducer<TpcCluster, TpcSPHit>(fTpcClusterArray));
  //
  //all ok
  if(fUseClusters){
    fAlMan=TpcAlignmentManager::getInstance();
    if(fBranchMap.count(fTpcClusterBranchName)){
      std::cout<<"TpcCluster branch found:"<<fTpcClusterBranchName<<std::endl;
      std::string type(fBranchMap[fTpcClusterBranchName]->GetClass()->GetName());
      if(type.find(tpcClusterType)<type.size()){
        fTpcClusterArray=fBranchMap[fTpcClusterBranchName];
        foundClusterArray=true;
      }
    }
    if(!foundClusterArray){
      std::cout<<"TpcRefGFTrkResCalc::init(): "<<std::endl
          <<"   ERROR: TpcClusters enabled but not found!"<<std::endl;
      return true;
    }
  }
  return false;
}


//main functionality
int TpcRefGFTrkResCalc::calc() {

  //clear results
  for(unsigned int i=0; i<fResults.size(); i++)
    delete fResults[i];

  fResults.clear();


  unsigned int nRefTracks = fRefGFTrackArray->GetEntriesFast();
  // GF trackfit loop

  for(unsigned int irtr=0; irtr<nRefTracks; ++irtr) {
    //std::cout<<"reftrack #"<<irtr<<std::endl;
    GFTrack* iRefTrack = (GFTrack*)((*fRefGFTrackArray)[irtr]);
    GFAbsTrackRep* refRep;
    if(fNRep>=0)
      refRep = iRefTrack->getTrackRep(fNRep);
    else
      refRep = iRefTrack->getCardinalRep();
    int val = refRep->getStatusFlag();

    if(val!=0 && fRequireGoodFit)
      continue;


    //cut to remove curlers in CDC, require position to be at TPC-edge of CDC
    //TODO configurable
    TVector3 refMom=refRep->getMom();
    TVector3 refPos=refRep->getPos();

    dPhiRefTrkPosHist->Fill((refMom.Phi()-refPos.Phi())*180/TMath::Pi());
    rRefHist->Fill(refPos.Perp());
    xyRefHist->Fill(refPos.X(),refPos.Y());
    momRefHist->Fill(refMom.Mag());
    phiRefHist->Fill(refMom.Phi()*180/TMath::Pi());

    //    if(refPos.Perp()>fTrackPosCut){
    //      continue;
    //    }
    //cut to only use well define tracks, ie high momentum
    //TODO configurable

    if(refMom.Mag()<fTrackMomentumCut){
      continue;
    }
    dPhiRefCutTrkPosHist->Fill((refMom.Phi()-refPos.Phi())*180/TMath::Pi());
    rRefCutHist->Fill(refPos.Perp());
    xyRefCutHist->Fill(refPos.X(),refPos.Y());
    momRefCutHist->Fill(refMom.Mag());
    phiRefCutHist->Fill(refMom.Phi()*180/TMath::Pi());

    //create residualCollection for this track
    fResults.push_back(new TpcRefResidualCollection(fTpcSPHitBranchName.Data(),irtr));

    unsigned int nTracks = fGFTrackArray->GetEntriesFast();

    //refMom.RotateZ(TMath::Pi());
    double refPhi=refMom.Phi();
    if(fabs(refPhi-refPos.Phi())*180/TMath::Pi()>150){
      refMom.RotateZ(TMath::Pi());
      refPhi=refMom.Phi();
      //   std::cerr<<"Error:: Rotating track!"<<std::endl;
    }
    dPhiRefCutTrkPosHist2->Fill((refMom.Phi()-refPos.Phi())*180/TMath::Pi());
    //sort out tracks not fitting
    for(unsigned int itr=0;itr<nTracks;++itr){
      GFTrack* iTrack=(GFTrack*)((*fGFTrackArray)[itr]);
      GFAbsTrackRep* rep= iTrack->getCardinalRep();
      GFAbsTrackRep* repClone = rep->clone();
      TVector3 mom = rep->getMom();
      TVector3 pos = rep->getPos();
      if(mom.Mag()<0.2 ||rep->getStatusFlag()!=0){
        continue;
      }
      try{
        rep->extrapolateToCylinder(20,pos,mom);
      }catch(GFException& ex){
        std::cout<<"TpcRefGFTrackResCalc::calc(): Error during extrap to cylinder mom: \n"
                <<"     "<<ex.what()<<std::endl;
        continue;
      }
      rHist->Fill(pos.Perp());
      momHist->Fill(mom.Mag());
      phiHist->Fill(mom.Phi()*180/TMath::Pi());

      dPhiTpcTrkPosHist->Fill((mom.Phi()-pos.Phi())*180/TMath::Pi());


      dPhiVsPhiHist->Fill(mom.Phi()*180/TMath::Pi(),(mom.Phi()-refPhi)*180/TMath::Pi());
      dPhiHist->Fill((mom.Phi()-refMom.Phi())*180/TMath::Pi());

      dPhiVsPhiPosHist->Fill(pos.Phi()*180/TMath::Pi(),(pos.Phi()-refPos.Phi())*180/TMath::Pi());
      dPhiPosHist->Fill((pos.Phi()-refPos.Phi())*180/TMath::Pi());

      if(fabs(pos.Phi()-refPos.Phi())*180/TMath::Pi()>this->fTrackdPhiCut){
        continue;
      }

      GFTrackCand cand=iTrack->getCand();
      unsigned int nHits = cand.getNHits();
      for(unsigned int iHit=0;iHit<nHits;++iHit){
        int hitId, detId;

        cand.getHit(iHit,detId,hitId);

        TpcSPHit* recoHit = (TpcSPHit*) fTpcSPHitArray->At(hitId);
        if(recoHit==NULL){
          std::cout<<"nooooooooooooooooooooooooooooooooo "<<iHit<<" "<<hitId<<" "<<detId<<std::endl;
          std::cout<<fTpcSPHitArray->GetEntries()<<std::endl;
        }
        TVector3 hPos;
        TMatrixT<double> cov(3,3) ;

        if(fUseClusters){
          int iCl=recoHit->getClusterID();
          TpcCluster* cl= (TpcCluster*) fTpcClusterArray->At(recoHit->getClusterID());
          if(cl==0){
            std::cout<<"no cluster at id: "<<iCl<<" nCl: "<<fTpcClusterArray->GetEntriesFast()<<std::endl;
            std::cout<<"hitid: "<<hitId<<" ihit: "<<iHit<<" nhit: "<<fTpcSPHitArray->GetEntriesFast()<<std::endl;
            continue;
          }
          hPos=cl->pos();
          hPos=fAlMan->localToMaster("tpc",hPos);
          TVector3 sig = cl->sig();
          cov[0][0] = sig.X()*sig.X();
          cov[1][1] = sig.Y()*sig.Y();
          cov[2][2] = sig.Z()*sig.Z();
          cov=fAlMan->localToMaster("tpc",cov);
        }else{
          TVectorT<double> rc = ((GFAbsRecoHit*)recoHit)->getRawHitCoord();
          cov =((GFAbsRecoHit*)recoHit)->getRawHitCov();
          hPos = TVector3(rc[0], rc[1], rc[2]);
        }
        double tracklength=0;
        bool ok=true;
        TVector3 poca, norm;
        TVector3 pos1,mom1;
        TMatrixDSym cov6;
        try {

          refRep->extrapolateToPoint(hPos, poca, norm);
          TVectorT<double> dummy;
          GFDetPlane plane(poca,norm);
          //        plane.setO(poca);
          //        plane.setNormal(norm);
          //        RKTrackRep* trackrepCopy=(RKTrackRep*)rep->clone();
          tracklength=refRep->extrapolate(plane);

          refRep->getPosMomCov(pos1,mom1,cov6);

        }
        catch(GFException& ex) {
            std::cout<<"TpcRefGFTrackResCalc::calc(): Error during extrapolation: \n"
                    <<"     "<<ex.what()<<std::endl;
          ok=false;
        }
        if(ok) {  //add residual
          TpcResidual* theRes = new TpcResidual();
          theRes->setResXYZ(poca-hPos);
          theRes->setRefPos(poca);
          theRes->setHitPos(hPos);
          theRes->setTrackDir(mom1.Unit());
          theRes->setHitIndex(recoHit->getClusterID());
          theRes->setHitCov(cov);
          theRes->setTrackLength(tracklength);
          fResults.back()->addResidual(theRes);
        }
      }

    }//end GFTrack loo[
  }//end refGFTrack loop
}

void TpcRefGFTrkResCalc::writeHists(){
  TFile* file = new TFile(fQAFileName,"recreate");
  file->cd();
  for(unsigned int i =0;i<hists.size();++i){
    hists[i]->Write();
  }
  file->Close();
}