#include "PndForwardTrackFinderTask.h"

//ClassImp(PndForwardTrackFinderTask);

PndForwardTrackFinderTask::~PndForwardTrackFinderTask() {

}

void PndForwardTrackFinderTask::SetParContainers() {

}

InitStatus PndForwardTrackFinderTask::Init() {
  fIoman = FairRootManager::Instance();
  //load hits
  fHits = (TClonesArray*) fIoman->GetObject("FTSHit");
  //planes
  planes = fIoman->Register("FtsTrkPlanes", "PndTrack","FTS", kFALSE);
  planesHit = fIoman->Register("FtsTrkPlanesHits", "PndFtsHit","FTS", kFALSE);
  //tracklets
  fFirstTrackCandArray = fIoman->Register("FtsTrkFirstTrackCand", "PndTrackCand","FTS", kFALSE);
  //corrected tracklets
  correctedTracklets = fIoman->Register("FtsTrkCorrectedTracklets", "PndTrack","FTS", kFALSE);
  correctedHits = fIoman->Register("FtsTrkCorrectedHits", "PndFtsHit","FTS", kFALSE);
  //layer tracks
  fLayerTracks = fIoman->Register("FtsTrkLayerTracks", "PndTrack","FTS", kFALSE);
  fLayerTracksHits = fIoman->Register("FtsTrkLayerTracksHits", "PndFtsHit","FTS", kFALSE);
  //global tracks befor and after magnet
  fGlobalTracks = fIoman->Register("FtsTrkGlobalTracks", "PndTrack","FTS", kFALSE);
  fGlobalTracksHits = fIoman->Register("FtsTrkGlobalTracksHits", "PndFtsHit","FTS", kFALSE);
  //final Solution
  fFinalSolution = fIoman->Register("FtsTrkFinalSolution", "PndTrack","FTS", kTRUE);
  //track collection
  fTrackCollection = fIoman->Register("FtsTrkTrackCollection", "PndTrack","FTS", kTRUE);
  fTrackCollectionHits = fIoman->Register("FtsTrkTrackCollectionHits", "PndFtsHit","FTS", kTRUE);
  //fill the root-branches array
  rootBranches[0]=fLayerTracks;
  rootBranches[1]=fLayerTracksHits;
  rootBranches[2]=fGlobalTracks;
  rootBranches[3]=fGlobalTracksHits;
  //init the algorithm classes
  fPndFtsCellTrackletGenerator = new PndFtsCellTrackletGenerator();
  fPndFtsLineApproximator = new PndFtsLineApproximator();
  fPndTrackCombiner = new PndTrackCombiner(planes,planes);
  fModuleCombiner = new PndModuleCombiner();
  fWayFollower = new PndWayFollower(planes,planesHit);

  //for statistics:
  for(int i=0;i<20;i++)dists[i]=0;
  for(int i=0;i<20;i++)angles[i]=0;
  return kSUCCESS;
}

void PndForwardTrackFinderTask::Exec(Option_t* opt) {
  if (fVerbose > 0) std::cout
      << ">>>>>>>>>>>>>>====================PndForwardTrackFinderTask::Exec====================<<<<<<<<<<<<<<"
      << std::endl;
  if (fVerbose > 0) std::cout << "EventNumber:" << eventNumber++ << std::endl;
  //reset the TrackletGenerator
  fPndFtsCellTrackletGenerator->reset();
  //get the hits and add to the TrackletGenerator
  std::vector<PndFtsHit*> hits;
  for (int i = 0; i < fHits->GetEntries(); i++) {
    PndFtsHit* hit = (PndFtsHit*) (fHits->At(i));
    hit->SetEntryNr(FairLink(-1, fIoman->GetEntryNr(), fIoman->GetBranchId("FTSHit"), i));
    //hit->SetEntryNr(FairLink(-1, fIoman->GetEntryNr(), fIoman->GetBranchId("FTSHitErr"), i));
    hits.push_back(hit);
    fOriginalHits[hit->GetTubeID()]=hit;
  }
  fPndFtsCellTrackletGenerator->setHits(hits); //init cellular automaton

  fPndFtsCellTrackletGenerator->findTracks();//start cellular automaton

  //create first track cand (from output of cellular automaton)
  map<Int_t, vector<PndTrackCand>> result = fPndFtsCellTrackletGenerator->getTracklets();
  fFirstTrackCand.clear();
  for (map<Int_t, vector<PndTrackCand>>::iterator it = result.begin();
      it != result.end(); it++) {
    for(int i=0;i<it->second.size();i++)
      fFirstTrackCand.push_back(it->second[i]);
  }
  if (fVerbose > 0) cout << fFirstTrackCand.size() << " Tracks found " << endl;
  if(fFirstTrackCand.size() > 100){ //more than 100 tracks -> too many for reconstruction (high runtime)
    fFirstTrackCand.clear();
    result.clear();
    if (fVerbose > 0) cout << "Too many tracks for reconstruction" << endl;
  }

  //save data in root-branch
  for (int i = 0; i < fFirstTrackCand.size(); i++) {
    PndTrackCand* myCand = new ((*fFirstTrackCandArray)[i]) PndTrackCand(fFirstTrackCand[i]);
  }
  if (fVerbose > 0) cout << "-------  1111    -----" << endl;

  //create expanded track cands
  map<Int_t, vector<PndFtsExpandedTrackCand>> expandedTrackCands;
  for (map<Int_t, vector<PndTrackCand>>::iterator it = result.begin();
        it != result.end(); it++) {
    vector<PndTrackCand> myCands = it->second;
    for(int i=0;i<myCands.size();i++){
      PndTrackCand myCand = myCands[i];
      //start the approximation of lines
      PndFtsExpandedTrackCand c = fPndFtsLineApproximator->createExpandedTrackCand(myCand);
      if(c.getLineApproximations().size()!=0)
        expandedTrackCands[it->first].push_back(c);
    }
  }
  if (fVerbose > 0) cout << "-------  2222    -----" << endl;
  //////saveCorrectedTracklets(expandedTrackCands); //save in root branch

  //combine the expanded track cands
  //combinedX = combined lines in module X
  if (fVerbose > 0) cout << "-------  3333    -----" << endl;
  fPndTrackCombiner->init(expandedTrackCands);
  if (fVerbose > 0) cout << "----------------------combine module 1" << endl;
  vector<PndLineApproximation> combined1 = fPndTrackCombiner->combine(0);
  if (fVerbose > 0) cout << "-------------------------combine module 2" << endl;
  vector<PndLineApproximation> combined2 = fPndTrackCombiner->combine(1);
  if (fVerbose > 0) cout << "-----------------------------combine module 5" << endl;
  vector<PndLineApproximation> combined5 = fPndTrackCombiner->combine(4);
  if (fVerbose > 0) cout << "-------------------------------combine module 6" << endl;
  vector<PndLineApproximation> combined6 = fPndTrackCombiner->combine(5);

  //save the combined lines in root branch
  combinedTracksHitNum = 0;
  combinedTracksTrackNum = 0;
  //cout << "save layer tracks " << endl;
   /// FIXME: DO WE NEED THAT??
  //saveCombined(combined1,0,fIoman->GetBranchId("LayerTracksHits"),kFALSE);
  //saveCombined(combined2,0,fIoman->GetBranchId("LayerTracksHits"),kFALSE);
  //saveCombined(combined5,0,fIoman->GetBranchId("LayerTracksHits"),kFALSE);
  //saveCombined(combined6,0,fIoman->GetBranchId("LayerTracksHits"),kFALSE);

  //combination of the stations
  if (fVerbose > 0) cout<<"combination of the stations"<<endl;
  // mXY = combination of station X and Y
  vector<PndLineApproximation> m12 = fModuleCombiner->combineModules(combined1,combined2);
  vector<PndLineApproximation> m56 = fModuleCombiner->combineModules(combined5,combined6);

  //way follower
  if (fVerbose > 0) cout<<"way follower"<<endl;
  fWayFollower->init(m12,expandedTrackCands); //init the way follower
  //start the way follower (8 = doublelayer for the begin of the magnatic field)
  vector<PndTrackCollection> ptc = fWayFollower->followLines(8);

  //createStatictcs(ptc,m56);

  // combine modules
  if (fVerbose > 0) cout<<"combine modules"<<endl;
  fModuleCombiner->combineModules(ptc,m56); //combine 1-4 with 5+6
  fModuleCombiner->addUnusedHits(hits,ptc); //step 6: add unused hits
   /// FIXME: DO WE NEED THAT??
  //saveTrackCollection(ptc,kFALSE); //save in root branch

  //save solution in root branch
  combinedTracksHitNum = 0;
  combinedTracksTrackNum = 0;
  if (fVerbose > 0) cout << "save global " << endl;
  //saveCombined(m12,2,fIoman->GetBranchId("GlobalTracksHits"),kFALSE);
  if (fVerbose > 0) cout << "save final solution" << endl;
  saveTrackCollection(ptc,kTRUE);
  if (fVerbose > 0) cout << "we're done here!" << endl;
}

void PndForwardTrackFinderTask::FinishEvent() {
  fFirstTrackCandArray->Delete();
  fLayerTracks->Delete();
  correctedHits->Delete();
  correctedTracklets->Delete();
  fLayerTracksHits->Delete();
  fGlobalTracks->Delete();
  fGlobalTracksHits->Delete();
  planes->Delete();
  planesHit->Delete();
  fTrackCollection->Delete();
  fTrackCollectionHits->Delete();
  fFinalSolution->Delete();
}

void PndForwardTrackFinderTask::Finish() {
  //output statistics:
  cout << "Dists:" << endl;
  for(int i=0;i<20;i++) cout << dists[i] <<";";
  cout << endl <<"Angles:" << endl;
  for(int i=0;i<20;i++) cout << angles[i] <<";";
}

/**
 * saved the vector combined in root branch
 * branch: element id of the branch title in rootBranches array
 * branch id: id of the branch
 * useOrgHits: use the original hits or the corrected hits
 * */
void PndForwardTrackFinderTask::saveCombined(vector<PndLineApproximation> combined,Int_t branch,Int_t branchID,Bool_t useOrgHits){
  TVector3 v(1, 1, 1);
  for (int i = 0; i < combined.size(); i++) {
    PndLineApproximation l = combined[i];
    PndLine l2 = l.getLine();
    PndTrackCand c;
    FairTrackParP tp1(v, v, v, v, 1, v, v, v);
    FairTrackParP tp2(l2.getP1(), 3 * l2.getDir().Unit(), v, v, 1, v, v, v);
    for (int j = 0; j < l.getHits().size(); j++) {
      PndFtsHit* hit = l.getHits()[j];
      if(!useOrgHits){
        hit->SetEntryNr(FairLink(-1, fIoman->GetEntryNr(),branchID, combinedTracksHitNum));
        PndFtsHit* myHit = new ((*(rootBranches[branch+1]))[combinedTracksHitNum++]) PndFtsHit(*hit);
      } else {
        hit = fOriginalHits[hit->GetTubeID()];
      }
      c.AddHit(hit->GetEntryNr(), j);
    }
    PndTrack* myCand = new ((*(rootBranches[branch]))[combinedTracksTrackNum++]) PndTrack(tp1, tp2, c);
  }
}

/**
 * save the corrected tracklets res in root branch (linear)
 * */
void PndForwardTrackFinderTask::saveCorrectedTracklets(map<Int_t, vector<PndFtsExpandedTrackCand>> res){
  Int_t trackNum = 0;
  Int_t hitNum= 0;
  for (map<Int_t, vector<PndFtsExpandedTrackCand>>::iterator it = res.begin();it != res.end(); it++) {
    vector<PndFtsExpandedTrackCand> etc = it->second;
    for(int i=0;i<etc.size();i++){
      PndFtsExpandedTrackCand cand = etc[i];
      vector<PndLineApproximation> approxs = cand.getLineApproximations();
      for(int j=0;j<approxs.size();j++){
        PndLineApproximation approx = approxs[j];
        PndTrackCand trackCand;
        for (int k = 0; k < approx.getHits().size(); k++) {
          PndFtsHit* hit = approx.getHits()[k];
          hit->SetEntryNr(FairLink(-1, fIoman->GetEntryNr(),fIoman->GetBranchId("FtsTrkCorrectedHits"),hitNum));
          PndFtsHit* myHit = new ((*correctedHits)[hitNum++]) PndFtsHit(*hit);
          trackCand.AddHit(hit->GetEntryNr(), k);
        }
        TVector3 v(1, 1, 1);
        FairTrackParP tp1(approx.getLine().getP1(),
            approx.getLine().getDir().Unit(), v, v, 1, v, v, v);
        FairTrackParP tp2(v, v.Unit(), v, v, 1, v, v, v);
        //////FairTrackParP tp1(approx.getLine().getP1(),
            //////3 * approx.getLine().getDir().Unit(), v, v, 1, v, v, v);
        //////FairTrackParP tp2(v, 2 * v, v, v, 1, v, v, v);
        PndTrack* myCand = new ((*correctedTracklets)[trackNum++]) PndTrack(tp1, tp2, trackCand);
      }
    }
  }
}

/*
 * save the track collection in root branch
 * useOrgHits: use the original hits or the corrected hits
 * **/
void PndForwardTrackFinderTask::saveTrackCollection(vector<PndTrackCollection> coll,Bool_t withOrgHits){
  Int_t trackNum = 0;
  Int_t hitNum = 0;
  TVector3 v(1, 1, 1);
  for(int k=0;k<coll.size();k++){
    PndTrackCollection c = coll[k];
    if (c.getNLines() < 2) continue;
    if(withOrgHits){
      PndTrack* myCand = new ((*fFinalSolution)[trackNum++]) PndTrack(c.getPndTrack(fOriginalHits));
    } else {
      vector<PndLineApproximation> *approx = c.getLines();
      for(int i=0;i<approx->size();i++){
        PndLineApproximation a = (*approx)[i];
        PndTrackCand trackCand;
        for(int j=0;j<a.getHits().size();j++){
          PndFtsHit* hit = a.getHits()[j];
          hit->SetEntryNr(FairLink(-1, fIoman->GetEntryNr(),fIoman->GetBranchId("FtsTrkTrackCollectionHits"),hitNum));
          PndFtsHit* myHit = new ((*fTrackCollectionHits)[hitNum++]) PndFtsHit(*hit);
          trackCand.AddHit(hit->GetEntryNr(), j);
        }
        FairTrackParP tp1(a.getLine().getP1(),3 * a.getLine().getDir().Unit(), v, v, 1, v, v, v);
        FairTrackParP tp2(v, 2 * v, v, v, 1, v, v, v);
        PndTrack* myCand = new ((*fTrackCollection)[trackNum++]) PndTrack(tp1, tp2, trackCand);
      }
      FairTrackParP tp1(c.getCurrLine().getP1(),3 * c.getCurrLine().getDir().Unit(), v, v, 1, v, v, v);
      FairTrackParP tp2(v, 2 * v, v, v, 1, v, v, v);
      PndTrackCand  trackCand;
      PndTrack* myCand = new ((*fTrackCollection)[trackNum++]) PndTrack(tp1, tp2, trackCand);
    }
  }
}

//create a statistic (not needed for the algorithm)
void PndForwardTrackFinderTask::createStatictcs(vector<PndTrackCollection> c,vector<PndLineApproximation> a){
  PndFtsLineComparator comparator(0,0);
  for(int i=0;i<c.size();i++){
    PndTrackCollection coll = c[i];
    PndLineApproximation appr = getBest(coll,a);
    if(appr.getLine().getRating()==-1) continue;
    PndLine l1 = coll.getLastLine();
    PndLine l2 = appr.getLine();
    comparator.setZValue((l1.getP1()[2]+l2.getP1()[2])/2);
    Double_t dist = comparator.getDist2D(l1.getP1(),l2.getP1(),l1.getDir(),l2.getDir());
    Double_t angle = comparator.getAngle2D(l1.getDir(),l2.getDir());
    cout << "Dist=" << (Int_t)(dist/5) << endl;
    cout << "Angle=" << (Int_t)(angle/5) << endl;;
    if(dist>=100) dists[19]++;
    else dists[(Int_t)(dist/5)]++;
    if(angle>=100) angles[19]++;
    else angles[(Int_t)(angle/5)]++;
  }
}

//for a statistic (not needed for the algorithm)
PndLineApproximation PndForwardTrackFinderTask::getBest(PndTrackCollection &c,vector<PndLineApproximation> a){
  Int_t mc = getMcId(c.getHits());
  for(int i=0;i<a.size();i++){
    Int_t mc2 = getMcId(a[i].getHits());
    if(mc==mc2) return a[i];
  }
  PndLineApproximation appr;
  appr.getLine().setRating(-1);
  return appr;
}
//for a statistic (not needed for the algorithm)
Int_t PndForwardTrackFinderTask::getMcId(vector<PndFtsHit*> hits){
  Int_t mcTrackBranchId = fIoman->GetBranchId("MCTrack");
  vector<Int_t> trackIDs;
  for(int i=0;i<hits.size();i++){
    PndFtsHit* hit = fOriginalHits[hits[i]->GetTubeID()];
    FairLink link = hit->GetEntryNr();
    PndFtsPoint* p = (PndFtsPoint*) fIoman->GetCloneOfLinkData(link);
    FairMultiLinkedData mld = p->GetLinksWithType(mcTrackBranchId);
    FairLink linksToMc = mld.GetLink(0);
    int trackID = linksToMc.GetIndex();
    trackIDs.push_back(trackID);
  }
  Int_t maxID = 0;
  Int_t maxVal = -1;
  map<Int_t,Int_t> m;
  for(int i=0;i<trackIDs.size();i++){
    m[trackIDs[i]]++;
    if(m[trackIDs[i]] > maxVal){
      maxVal = m[trackIDs[i]];
      maxID = trackIDs[i];
    }
  }
  return maxID;
}

ClassImp(PndForwardTrackFinderTask);