//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of resCalc
//      Uses existing CDC circle fits (XY) and calculates
//      residua in XY.
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Felix Boehmer                     (original author)
//      Francesco Cusanno                 (adapted for the alignment manager)
//      Physik Department E18, TUM
//
//-----------------------------------------------------------

#include "TClonesArray.h"
#include "TClass.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TFile.h"

#include "TpcCdcFit2DResCalc_Alignment.h"
#include "CdcTrack.h"
#include "TpcCluster.h"
#include "TpcResidual.h"
#include "TpcRefResidualCollection.h"
#include "TpcAlignmentManager.h"
#include "TpcSPHit.h"
#include "MatchingTuple.h"

#include "GFTrack.h"
#include "GFAbsTrackRep.h"
#include "GFDetPlane.h"
#include "GFException.h"
#include "RKTrackRep.h"

#include <string>
#include <iostream>
#include <sstream>
#include <fstream>
#include <cmath>

TpcCdcFit2DResCalc_Alignment::TpcCdcFit2DResCalc_Alignment()
    : fUseClusters(false), fDistCut(-1), fMatching(true),
      fOrigTransfName("tpc"), fCdcTrackBranchName("CdcTrack"),
      fTpcTrackBranchName("TpcTrackPostFit"),
      fTpcClusterBranchName("TpcCluster"), fTpcSPHitBranchName("TpcSPHit"),
      fFopiTupleBranchName("FopiTrackTuples"), fQA("TpcCdcFit2DResCalc_Alignment_QA.root") {
  fNExpectedBranches = 3;
}

TpcCdcFit2DResCalc_Alignment::~TpcCdcFit2DResCalc_Alignment() {    fQA.writeHists(); }

bool TpcCdcFit2DResCalc_Alignment::init() {
  if(fMatching){
      fNExpectedBranches+=1;
  }
  if (fBranchMap.size() != fNExpectedBranches) {
    std::cout << "TpcCdcFit2DResCalc_Alignment::init(): " << std::endl
              << "   ERROR: Wrong number of input branches ("
              << fBranchMap.size() << ")!" << std::endl;
    return true;
  }
  // check for proper data types:
  std::string cdcTrackRefType("CdcTrack");
  std::string tpcClusterType("TpcCluster");
  std::string tpcSPHitType("TpcSPHit");
  std::string tpcTrackType("GFTrack");
  bool foundCdcTrack = false;
  bool foundTpcCluster = false;
  bool foundSPHit = false;
  bool foundTpcTrack = false;
  // identify the needed branches in the map

  fAlMan = TpcAlignmentManager::getInstance();
  if (fBranchMap.count(fCdcTrackBranchName)) {
    std::cout << "cdctrack branch found:" << fCdcTrackBranchName << std::endl;
    std::string type(fBranchMap[fCdcTrackBranchName]->GetClass()->GetName());
    if (type.find(cdcTrackRefType) < type.size()) {
      fCdcTrackArray = fBranchMap[fCdcTrackBranchName];
      foundCdcTrack = true;
    }
  }
  if (fBranchMap.count(fTpcTrackBranchName)) {
    std::cout << "track branch found:" << fTpcTrackBranchName << std::endl;
    std::string type(fBranchMap[fTpcTrackBranchName]->GetClass()->GetName());
    if (type.find(tpcTrackType) < type.size()) {
      fTpcTrackArray = fBranchMap[fTpcTrackBranchName];
      foundTpcTrack = 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];
      foundSPHit = true;
    }
  }
  if (!(foundTpcTrack && foundSPHit && foundCdcTrack)) {
    std::cout << "CdcTrackTpcHitResCalc::init(): " << std::endl
              << "   ERROR: Wrong Class Types in input branches!" << std::endl;
    return true;
  }

  if (fUseClusters) {
    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];
        foundTpcCluster = true;
      }
    }
    if (!foundTpcCluster) {
      std::cout << "CdcTrackTpcHitResCalc::init(): " << std::endl
                << "   ERROR: TpcClusters enabled but not found!" << std::endl;
      return true;
    }
  }
  if (fMatching) {
    bool foundTupleArray = false;
    if (fBranchMap.count(fFopiTupleBranchName)) {
      std::cout << "FopiTuple branch found:" << fFopiTupleBranchName
                << std::endl;
      std::string type(fBranchMap[fFopiTupleBranchName]->GetClass()->GetName());
      if (type.find("MatchingTuple") < type.size()) {
        fFopiTupleArray = fBranchMap[fFopiTupleBranchName];
        foundTupleArray = true;
      }
    }

    if (!(foundTupleArray)) {
      std::cerr
          << "CdcTrackTpcHitResCalc::init(): " << std::endl
          << "   ERROR: FopiMatchingTuples and tracks enabled but not found!"
          << std::endl;
      return true;
    }
  }
  std::ofstream txtoutfile("TpcCdcFit2DResCalc_Alignment_settings.txt");
  txtoutfile << "Settings for TpcCdcFit2DResCalc_Alignment" << std::endl;
  txtoutfile << "Original applied alignment: " << fOrigTransfName << std::endl;
  txtoutfile << "Use SPHits(0) or Cluster(1): " << fUseClusters << std::endl;
  txtoutfile << "Fiducial cut, max distance from circle: " << fDistCut
             << std::endl;
  txtoutfile << "Momentum cut, minimum momentum for ref tracks: " << fPtMin
             << ", " << fPtMax << std::endl;
  // All ok
  return false;
}

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

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

  // get instance of alignment manager
  TpcAlignmentManager *alMan = TpcAlignmentManager::getInstance();

  unsigned int nTracks = fCdcTrackArray->GetEntriesFast();
  // //CDC trackfit loop
  std::map<int, std::vector<int> > matchedTpcTracks;
  if (fMatching) {
    int nTuples = fFopiTupleArray->GetEntriesFast();
    for (unsigned iTuple = 0; iTuple < nTuples; ++iTuple) {
      MatchingTuple *theTuple = (MatchingTuple *)fFopiTupleArray->At(iTuple);
      if (theTuple->isMatched()) {
        if (theTuple->hasBranch(fCdcTrackBranchName) and
            theTuple->hasBranch(fTpcTrackBranchName)) {
          matchedTpcTracks[theTuple->getIdFromBranch(fCdcTrackBranchName)]
              .push_back(theTuple->getIdFromBranch(fTpcTrackBranchName));
        }
      }
    }
  }
  for (unsigned int itr = 0; itr < nTracks; itr++) {
    CdcTrack *itrack = (CdcTrack *)((*fCdcTrackArray)[itr]);
    // validity - I don't know how it's defined, but I'll use it for now...

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

    double r = fabs(itrack->GetRadius());
    int charge = itrack->GetRadius() / fabs(itrack->GetRadius());
    TVector3 M(itrack->GetMx(), itrack->GetMy(), 0.);
    unsigned int ncl = 0;
    if (fUseClusters) {
      ncl = fTpcClusterArray->GetEntriesFast();
    } else {
      ncl = fTpcSPHitArray->GetEntriesFast();
    }
    GFTrack *tpcTrack = NULL;
    double pt = r * 0.3 * 0.616/100;
    if (fMatching) {
      if (!matchedTpcTracks.count(itr) or matchedTpcTracks[itr].size() != 1) {
        continue;
      }
      int tpcId = matchedTpcTracks[itr][0];
      CdcTrack *cdcTrack = (CdcTrack *)fCdcTrackArray->At(itr);
      tpcTrack = (GFTrack *)fTpcTrackArray->At(tpcId);
      TVector3 cdcMom, cdcPos;
      getTrackCylinderIntersect(cdcTrack, 15, cdcPos, cdcMom);
      if(cdcPos.Perp()<10){
          continue;
      }
      cdcMom.SetZ(cdcMom.Perp()/TMath::Tan(cdcTrack->GetTheta()));
      TVector3 tpcMom, tpcPos;
      try {
        RKTrackRep *theRep = (RKTrackRep *)tpcTrack->getCardinalRep();
        theRep->extrapolateToCylinder(15, tpcPos, tpcMom);
      } catch (GFException &exp) {
        continue;
      }

      fQA.fillBasicHists("cdc", cdcMom, cdcPos, charge);
      fQA.fillBasicHists("cdc", cdcMom, cdcPos, 0);
      fQA.fillBasicHists("tpc", tpcMom, tpcPos, charge);
      fQA.fillBasicHists("tpc", tpcMom, tpcPos, 0);
      //std::cout<<"pt: "<<pt<<std::endl;
      //std::cout<<"r: "<<r<<std::endl;
      if (pt > fPtMax or pt < fPtMin) {
        continue;
      }
      fQA.fillBasicHists("cdc_cutPt", cdcMom, cdcPos, charge);
      fQA.fillBasicHists("cdc_cutPt", cdcMom, cdcPos, 0);
      fQA.fillBasicHists("tpc_cutPt", tpcMom, tpcPos, charge);
      fQA.fillBasicHists("tpc_cutPt", tpcMom, tpcPos, 0);

      
      fQA.fillMatchingHists(tpcMom, tpcPos, cdcMom, cdcPos, false, 0);
      fQA.fillMatchingHists(tpcMom, tpcPos, cdcMom, cdcPos, false, charge);

      double dTheta=(tpcMom.Theta()-cdcTrack->GetTheta())*TMath::RadToDeg();
      if(fabs(dTheta)>20){
	continue;
      }
      fQA.fillMatchingHists(tpcMom, tpcPos, cdcMom, cdcPos, true, 0);
      fQA.fillMatchingHists(tpcMom, tpcPos, cdcMom, cdcPos, true, charge);

      ncl=tpcTrack->getCand().getNHits();
    } else {
      if (pt > fPtMax or pt < fPtMin) {
        continue;
      }
    }
    // just do QA
    // continue;
    // calculate residuals in XY --------------------------

    for (unsigned int icl = 0; icl < ncl; icl++) {
      // solve for intersection of vector to point and circle fit ...
      TVector3 clPos;
      if (!fUseClusters) {
        int clusterId = icl;
        if (fMatching) {
            if(tpcTrack==NULL){
                continue;
            }
          int hitId, detId;
          tpcTrack->getCand().getHit(icl, detId, hitId);
          clusterId = hitId;
        }
        TpcSPHit *hit = (TpcSPHit *)(fTpcSPHitArray->At(clusterId));
        TVectorD hitPos = hit->getRawHitCoord();
        TVector3 hitVec(hitPos[0], hitPos[1], hitPos[2]);
        if (fOrigTransfName == "tpc") {
          clPos = hitVec;
        } else {
          clPos = alMan->masterToLocal(fOrigTransfName, hitVec);
          if(clPos.Perp()>13 or 
                  clPos.Perp()<7){
              continue;
          }
          clPos = alMan->localToMaster("tpc", clPos);
        }
      } else {
        TpcCluster *cl = (TpcCluster *)((*fTpcClusterArray)[icl]);
        clPos = alMan->localToMaster("tpc", cl->pos());
      }
      //
      
      TVector3 copy = clPos;
      clPos.SetZ(0.);
      
      TVector3 diff = clPos - M; //from m pointing towards clPos
      double dist = diff.Mag() - r; 
      diff.SetMag(dist);
      
      TVector3 refPos=clPos - diff;
      
      TVector3 refMom=diff;
      refMom.RotateZ(charge/abs(charge)*TMath::PiOver2());
      // fiducial volume cut
      if (fDistCut > 0 && fabs(dist) > fDistCut) {
        continue;
      }
      // save residual
      
      TpcResidual *resi = new TpcResidual();
      resi->setResXYZ(diff);
      resi->setHitPos(copy);
      resi->setHitIndex(icl);
      resi->setRefPos(refPos);
      resi->setTrackDir(refMom.Unit());
      resi->setCharge(charge);
      
      fResults.back()->addResidual(resi);
    }
    // end residual calculation  -------------------------
  } // end CDC trackfit loop
}

void TpcCdcFit2DResCalc_Alignment::writeHists() {
    fQA.writeHists();
}

void TpcCdcFit2DResCalc_Alignment::getTrackCylinderIntersect(CdcTrack *tr,
                                                             unsigned int r_cyl,
                                                             TVector3 &pos,
                                                             TVector3 &mom) {

  double r_tr = fabs(tr->GetRadius());
  int charge = tr->GetRadius() / fabs(tr->GetRadius());
  TVector3 trackCenter(tr->GetMx(), tr->GetMy(), 0.);
  double d = trackCenter.Mag();
  // track completely outside cylinder
  // track enclosing cylinder
  // track enclosed by cylinder
  if ((r_tr + r_cyl < d) or (r_cyl + d < r_tr) or (r_tr + d < r_cyl)) {
    return;
  }
  // x measured along conecting line between 0,0 and track center
  // y perpendicular
  double x =
      (trackCenter.Mag() * trackCenter.Mag() - r_tr * r_tr + r_cyl * r_cyl) /
      (2 * trackCenter.Mag());
  double y = sqrt(r_cyl * r_cyl - x * x);
  
  TVector3 xVec = trackCenter;
  xVec.SetMag(x);
  TVector3 yVec = trackCenter;
  yVec.RotateZ(TMath::PiOver2());
  yVec.SetMag(y);

  pos = xVec + charge / abs(charge) * yVec;
  mom =  trackCenter-pos;
  mom.RotateZ(charge/abs(charge)*TMath::PiOver2());
  mom.SetMag(r_tr * 0.3 * 0.616/100);

}