//////////////////////////////ALIGNMENT OF TEST CHAMBER IN RESPECT OF TEST BENCH TRACKS/////////////////////////
//Minized residuals between TpcClusters and GFTracks
//Residuals are calculated via TpcRefGFTrkResCalc
//Translations/Rotations of the chamber is done via the TpcAlignment Manager
//Name of the Test chamber ("TC01____") is hardcoded here
//
//Plots in world coordinates, blue is corrected, red is original distributions
//
//New alignement saved in ./newal.txt
//
//Needs to be compiled!
//Typical command in root : .x macro/tpc/TestBench/alignTestBenchAndTC.C++("run-2076.t0PSA.raw_decoded.root","macro/tpc/TestBench/COMPASS_alignment.tpcAligned.txt",5000)
//Author: Tiger

#include "TCanvas.h"
#include "TFile.h"
#include "TH2F.h"
#include "TH3F.h"
#include "TClonesArray.h"
#include "TGeoManager.h"
#include "TTree.h"
#include "TObjArray.h"
#include "TPolyMarker.h"
#include "TPolyMarker3D.h"
#include "TPolyLine.h"
#include "TVector3.h"
#include "TBox.h"
#include "TSystem.h"
#include "TROOT.h"
#include "TMinuit.h"

#include "TView3D.h"
#include "TLegend.h"
#include "TApplication.h"
#include "TMath.h"
#include "TStyle.h"
#include "TString.h"
#include "TColor.h"
#include "THStack.h"
#include "TVector3.h"
#include "tpc/par/TpcAlignmentManager.h"
//#include "tpc/TestBench/TBStripCluster.h"
#include "tpc/base/TpcCluster.h"
#include "tpc/TestBench/TBGEMCluster.h"
#include "tpc/TestBench/TBSICluster.h"

#include "tpc/base/TpcResidual.h"
#include "tpc/tools/residuals/TpcRefResidualCollection.h"
#include "tpc/base/TpcResidualCollection.h"
#include "tpc/tools/residuals/TpcRefGFTrkResCalc.h"
#include "GFAbsTrackRep.h"
#include "GFAbsRecoHit.h"
#include "GFConstField.h"
#include "GFDetPlane.h"
#include "GFException.h"
#include "GFFieldManager.h"
#include "GFKalman.h"
#include "GFTools.h"
#include "GFTrack.h"
#include "GFTrackCand.h"
#include "GenfitTools/trackrep/RKTrackRep/RKTrackRep.h"

#include <vector>
#include <cmath>
#include <iostream>
#include <algorithm>
#include <map>
#include <string>
#include <exception>

#define DEBUG 1
const int NDET = 8;
const int NDETXY = 4;
const double MAXRES = 0.4; //IMPORTANT cut on maximal distance bet clusters and tracks taking into account for al.
const char* TPCDET = "TC01____";
//class TBSICluster;
void fcn(Int_t &npar, Double_t *gin, Double_t &f, Double_t *par, Int_t iflag);
//GLOBALS FOR MINIMAZATION
std::vector<TpcResidual*> vres;
double MEANZ =0.;
TpcAlignmentManager * alMan;

void alignTestBenchAndTC(TString filename, std::string alfile = "macro/tpc/TestBench/COMPASS_alignment.txt", unsigned int EvMax=0)
{
 TGeoManager* geom = new TGeoManager("Geometry", "Geane geometry");
 TGeoManager::Import("/nfs/hicran/project/panda/SIM/vandenbm/fopiroot/tpc/TestBench/FOPIGeo.root");
 
 GFFieldManager::getInstance()->init(new GFConstField(0.,0.,0.));

 while(gROOT->GetListOfCanvases()->GetSize())
    delete gROOT->GetListOfCanvases()->At(0);
  gStyle->SetOptStat(10);
  gStyle->SetPalette(1);
  int a = 0;
  TFile* recofile = new TFile(filename);
  if(recofile->IsZombie()) {
    std::cerr<<"Reco file not existing! Aborting."<<std::endl;
    return;  }
  TTree* recotree = (TTree*)recofile->Get("cbmsim");
  if(!recotree) {
    std::cerr<<"Reco Tree not existing! Aborting."<<std::endl;
    return;  }
  std::map<TString, TClonesArray*> Tbranches;
  std::cout << "ARRAYS..." << std::endl;
  //I CAN HAZ ARRAYS :3
  TClonesArray* trTBArr=new TClonesArray("GFTrack");
  recotree->SetBranchAddress("TBTrackPostFit", &trTBArr);
  Tbranches["TBTrackPostFit"] = trTBArr;
  if(trTBArr==0) {std::cout<< "Init TB GFTrack-Array not found!"<< std::endl;return;}

  TClonesArray* tpcArr=new TClonesArray("TpcCluster");
  recotree->SetBranchAddress("TpcCluster", &tpcArr);
  Tbranches["TpcCluster"] = tpcArr;
  if(tpcArr==0) {std::cout<< "Init TpcCluster Array not found!"<< std::endl;return;}
  
  TClonesArray* tpcResArr=new TClonesArray("TpcRefResidualCollection");
  recotree->SetBranchAddress("TBTPCres", &tpcResArr);
  if(tpcResArr==0) {std::cout<< "Init TpcResidualCollection Array not found!"<< std::endl;return;}
  std::cout << "ALMAN..." << std::endl;
  alMan = TpcAlignmentManager::getInstance(alfile.c_str());
  if(alMan==0) {std::cout<< "Init TpcAlignmentManager:"<< alfile.c_str()<<" not found!"<< std::endl;return;}
  double phi,theta,psi, tmp;  
  std::cout << "EULER..." << std::endl;
  alMan->getEulerAngles(TPCDET, phi, theta, psi);
  std::cout << "TRANS..." << std::endl;
  TVector3 tpcpos = alMan->getTransformationVector(TPCDET);
  std::cout << "FILE : " << alfile.c_str() << std::endl;
  std::cout << "TPC xyz("<<tpcpos.X()<< ", "<<tpcpos.Y()<< ", "<<tpcpos.Z()<< "). Rot. angles("
	    << phi<<", " << theta << ", " << psi<< ")"<<std::endl;
  char oldtpcxyz[500];
  sprintf(oldtpcxyz,"OLD TPC xyz(%f, %f, %f). Rot. angles(%f, %f, %f)", tpcpos.X(),tpcpos.Y(),tpcpos.Z(),phi,theta,psi);
  std::cout << "Turing test passed." << std::endl;
  
  //TH1...
  TH1F * hresx = new TH1F("hresx","hresx",800,-1.*MAXRES,MAXRES);hresx->SetLineColor(kRed);
  TH1F * hresy = new TH1F("hresy","hresy",800,-1.*MAXRES,MAXRES);hresy->SetLineColor(kRed);
  TH1F * hresz = new TH1F("hresz","hresz",800,-0.02,0.02);hresz->SetLineColor(kRed);
  
  TH1F * hnesx = new TH1F("hnesx","hnesx",800,-1.*MAXRES,MAXRES);hnesx->SetLineColor(kBlue);
  TH1F * hnesy = new TH1F("hnesy","hnesy",800,-1.*MAXRES,MAXRES);hnesy->SetLineColor(kBlue);
  TH1F * hnesz = new TH1F("hnesz","hnesz",800,-0.02,0.02);hnesz->SetLineColor(kBlue);
  std::cout << "CALCULATOR : " << std::endl;
  //Residuals calculator
  TpcRefGFTrkResCalc* resgft = new TpcRefGFTrkResCalc();
  int test =   resgft->addBranchName("TpcCluster");
  if(test) {std::cout << "resgft->addBranchName(\"TpcCluster\"); is inconsistent" << std::endl;}
  test = resgft->addBranchName("TBTrackPostFit");
  if(test) {std::cout << "resgft->addBranchName(\"TBTrackPostFit\"); is inconsistent" << std::endl;}
  TVector3 uu(1.,0.,0.);  TVector3 vv(0.,1.,0.);  TVector3 oo(0.,0.,105.);//plane in front of the TC
  resgft->setProjectionPlane(oo,uu,vv);
  resgft->setTrackRepId(0);
  std::map<TString, TClonesArray*>* branches = resgft->getBranchMap();
  std::map<TString, TClonesArray*>::iterator it;
  for(it=branches->begin(); it!=branches->end(); it++) {
    TClonesArray* br = Tbranches[it->first];
    if(br==NULL) {
      std::string s("Branch ");
      s.append(it->first.Data());
      s.append(" requested by residual calculator, but not found in tree!");
      Fatal("TpcRefTrackResidualTask::Init()",s.c_str());
    }
    it->second = br;
  } 
  std::cout << "Self awareness ... " << std::endl;
  resgft->init();
  unsigned int MAX =recotree->GetEntries();
  std::cout << "Number of events: " << MAX  <<std::endl;
  if(!EvMax)EvMax=MAX;
  vres.clear();
  for(unsigned int k=0; k<MAX && k<EvMax; k++) 
    {
      //Set TpcClusters to current alignement
      recotree->GetEvent(k);
      int ntpccl = tpcArr->GetEntries();
      for(int c=0; c<ntpccl; c++) 
	{
	  TpcCluster* cl = (TpcCluster*)tpcArr->At(c);
	  TVector3 newxyz = alMan->localToMaster(TPCDET,cl->posuvw());
	  cl->SetPos(newxyz);
	}
      
      resgft->calc();
      std::vector<TpcRefResidualCollection*>* blka =  resgft->getResiduals() ;
      int nRes = blka->size();
      for(int ires=0; ires<nRes; ires++) 
	{
	  TpcRefResidualCollection * resCol = (*blka)[ires];
	  for (int ii = 0 ; ii < resCol->getSize(); ii++)
	    {
	      TpcResidual * miaou = new TpcResidual(*resCol->getResidual(ii));
	      TpcCluster* cl = (TpcCluster*)tpcArr->At(miaou->getHitIndex());
	      if(miaou->getResXYZ().Mag()>MAXRES) continue; //CUT ON RESIDALS
	      vres.push_back(miaou);
	      hresx->Fill(miaou->getResXYZ().X());
	      hresy->Fill(miaou->getResXYZ().Y());
	      hresz->Fill(miaou->getResXYZ().Z());
	    }
	}
    }//end event loop 
  std::cout << vres.size()<<" TpcResiduals, Starting sattelites alignement..." << std::endl; \
  
  TMinuit *gMinuit;
  gMinuit = new TMinuit(5);
  gMinuit->SetFCN(fcn);
  gMinuit->DefineParameter(0, "x", 0., 0.00001, -0.5,0.5);
  gMinuit->DefineParameter(1, "y", 0., 0.00001, -0.5,0.5);
  gMinuit->DefineParameter(2, "phi", 0., 0.000001, -30., 30.);//-3.
  gMinuit->DefineParameter(3, "theta", 0., 0.00001, -30., 30.);
  gMinuit->DefineParameter(4, "psi", 0., 0.00001, -30., 30.);
  gMinuit->DefineParameter(5, "drift", 1., 0.00001, 0.7, 1.3);
  
  //First XY
  std::cout <<"##############################################################\n"
	    <<"------------------ First XY  Alignment ------------------------" << std::endl;
  gMinuit->FixParameter(2);
  gMinuit->FixParameter(3);
  gMinuit->FixParameter(4);
  gMinuit->FixParameter(5);
  // Now ready for minimization step
  gMinuit->Migrad();
  // Print results
  Double_t amin,edm,errdef;
  Int_t nvpar,nparx,icstat;
  gMinuit->mnstat(amin,edm,errdef,nvpar,nparx,icstat);
  
  //Then Drift
  std::cout <<"##############################################################\n"
	    <<"------------------ DRIFT  Alignment ------------------------" << std::endl;
  gMinuit->FixParameter(0);
  gMinuit->FixParameter(1);
  gMinuit->Release(5);
  // Now ready for minimization step
  gMinuit->Migrad();
  // Print results
  gMinuit->mnstat(amin,edm,errdef,nvpar,nparx,icstat);
  
  /*
  std::cout <<"##############################################################\n"
	    <<"------------------ XYDRIFT  Alignment ------------------------" << std::endl;
  gMinuit->Release(0);
  gMinuit->Release(1);
  // Now ready for minimization step
  gMinuit->Migrad();
  // Print results
  gMinuit->mnstat(amin,edm,errdef,nvpar,nparx,icstat);
  
  //Then angles
  std::cout <<"##############################################################\n"
	    <<"------------------ Theta.Phi Alignment -----------------------" << std::endl;
  gMinuit->Release(2);
  //gMinuit->Release(3);
  //gMinuit->Release(4);
  gMinuit->FixParameter(0);
  gMinuit->FixParameter(1);
  gMinuit->FixParameter(5);
  
  // Now ready for minimization step
  gMinuit->Migrad();
  gMinuit->mnstat(amin,edm,errdef,nvpar,nparx,icstat);
  
  //Again XY
  std::cout <<"##############################################################\n"
	    <<"------------------ Second XY  Alignment ----------------------" << std::endl;
  gMinuit->Release(0);
  gMinuit->Release(1);
  gMinuit->Release(5);
  gMinuit->FixParameter(2);
  gMinuit->FixParameter(3);
  gMinuit->FixParameter(4);
  
  // Now ready for minimization step
  gMinuit->Migrad();
  gMinuit->mnstat(amin,edm,errdef,nvpar,nparx,icstat);
  
  gMinuit->FixParameter(0);
  gMinuit->FixParameter(1);
  gMinuit->FixParameter(5);
  gMinuit->Release(3);
  gMinuit->Migrad();
  
  gMinuit->FixParameter(3);
  gMinuit->Release(0);
  gMinuit->Release(1);
  gMinuit->Release(5);
  gMinuit->Migrad();
  
  gMinuit->FixParameter(0);
  gMinuit->FixParameter(1);
  gMinuit->FixParameter(5);
  gMinuit->Release(4);
  gMinuit->Migrad();
  */
  // Final Released
  std::cout <<"##############################################################\n"
	    <<"------------------ All released Alignment --------------------" << std::endl;
  gMinuit->Release(0);
  gMinuit->Release(1);
  gMinuit->Release(2);
  gMinuit->Release(3);
  gMinuit->Release(4);
  gMinuit->Release(5);
  
  // Now ready for minimization step
  gMinuit->Migrad();
  gMinuit->mnstat(amin,edm,errdef,nvpar,nparx,icstat);
  
  //The end
  double x,y,xerr,yerr,ntheta,nthetaerr,nphi,nphierr, npsi, npsierr, ndr, ndrerr;
  gMinuit->GetParameter(0,x,xerr);std::cout << "x,xerr" << x<<", " <<xerr << std::endl;
  gMinuit->GetParameter(1,y,yerr);std::cout << "y,yerr" << y<<", " <<yerr << std::endl;
  gMinuit->GetParameter(2,nphi,nphierr);std::cout << "nphi,nphierr" << nphi<<", " <<nphierr << std::endl;
  gMinuit->GetParameter(3,ntheta,nthetaerr);std::cout << "theta,thetaerr" << ntheta<<", " <<nthetaerr << std::endl;
  gMinuit->GetParameter(4,npsi,npsierr);std::cout << "psi,psierr" << npsi<<", " <<npsierr << std::endl;
  gMinuit->GetParameter(5,ndr,ndrerr);std::cout << "Drift,Drifterr" << ndr<<", " <<ndrerr << std::endl;
  TVector3 ntpcpos = tpcpos + TVector3(x,y,0.);
  //Final transforamtion
  alMan->setEulerAngles(TPCDET, nphi+phi,ntheta+theta,psi+npsi);  //and we rotate
  alMan->setTransformation(TPCDET, ntpcpos);        //then translate
  alMan->getEulerAngles(TPCDET, phi, theta, psi);
  tpcpos = alMan->getTransformationVector(TPCDET);
  //Modify all TpcCluster with the new Alignment
  for(unsigned int k=0; k<MAX && k<EvMax; k++) 
    {
      std::cout << "." << std::flush;
      recotree->GetEvent(k);
      int ntpccl = tpcArr->GetEntries();
      for(int c=0; c<ntpccl; c++) 
	{
	  TpcCluster* cl = (TpcCluster*)tpcArr->At(c);
	  TVector3 newxyz = alMan->localToMaster(TPCDET,cl->posuvw());
	  newxyz.SetZ(newxyz.Z()*ndr);
	  cl->SetPos(newxyz);
	}
      resgft->calc();
      std::vector<TpcRefResidualCollection*>* blka =  resgft->getResiduals() ;
      int nRes = blka->size();
      for(int ires=0; ires<nRes; ires++) 
	{
	  TpcRefResidualCollection * resCol = (*blka)[ires];
	  for (int ii = 0 ; ii < resCol->getSize(); ii++)
	    {
	      TpcResidual * miaou = new TpcResidual(*resCol->getResidual(ii));
	      TpcCluster* cl = (TpcCluster*)tpcArr->At(miaou->getHitIndex());
	      if(miaou->getResXYZ().Mag()>MAXRES) continue; //cut 
	      hnesx->Fill(miaou->getResXYZ().X());
	      hnesy->Fill(miaou->getResXYZ().Y());
	      hnesz->Fill(miaou->getResXYZ().Z());
	    }
	}
    }//end event loop 
  
  std::cout << "DRAWING " << std::endl;
  
  char tmpChar[500];
  //Drawing()
  TCanvas * cx = new TCanvas("cx","cx",1200,600);
  cx->Divide(3);
  cx->cd(1);
  sprintf(tmpChar,"X: Old RMS = %f, New RMS = %f",hresx->GetRMS(),hnesx->GetRMS());
  hnesx->SetTitle(tmpChar);
  hnesx->Draw();      hresx->Draw("same");
  cx->cd(2);
  sprintf(tmpChar,"Y: Old RMS = %f, New RMS = %f",hresy->GetRMS(),hnesy->GetRMS());
  hnesy->SetTitle(tmpChar);
  hnesy->Draw();hresy->Draw("same");
  cx->cd(3);
  sprintf(tmpChar,"Z: Old RMS = %f, New RMS = %f",hresz->GetRMS(),hnesz->GetRMS());
  hnesz->SetTitle(tmpChar);
  hnesz->Draw();hresz->Draw("same");
  
  TLegend * tl1 = new TLegend(0.1,0.7,0.48,0.9);
  tl1->SetHeader("Residuals in lab coordinates");
  tl1->AddEntry(hresx,"Before Alignment","l");
  tl1->AddEntry(hnesx,"After Alignment","l");
  tl1->Draw();
  std::cout << oldtpcxyz << std::endl;
  std::cout << "NEW TPC xyz("<<tpcpos.X()<< ", "<<tpcpos.Y()<< ", "<<tpcpos.Z()<< "). Rot. angles("
	    << phi<<", " << theta << ", " << psi<< ")"<< " And Drift = " << ndr << "+-" << ndrerr<< std::endl;
  std::cout << "OLD RMS:" << hresx->GetRMS() << ", "<< hresy->GetRMS() << ", "<< hresz->GetRMS() << std::endl;
  std::cout << "NEW RMS:" << hnesx->GetRMS() << ", "<< hnesy->GetRMS() << ", "<< hnesz->GetRMS() << std::endl;
  std::cout << "Saving alignment file in ./newal.txt"<<std::endl;
  alMan->write("newal.txt");
}//end of macro

  
//Minuit fonction to minized
void fcn(Int_t &npar, Double_t *gin, Double_t &f, Double_t *par, Int_t iflag)
{
  Int_t i;
  TVector3 vadd;
  double phipar = par[2];
  double thetapar =  par[3];
  double psipar= par[4];
  double corrdrift = par[5];
  vadd = TVector3(par[0], par[1], 0);
  Double_t chisq = 0;
  Double_t sum =0.;
  double phi,theta,psi, tmp;  alMan->getEulerAngles(TPCDET, phi, theta, psi);
  TVector3 tpcpos = alMan->getTransformationVector(TPCDET);
  Int_t nbins = vres.size();
  for (i=0;i<nbins; i++) 
    {
      TVector3 ref = vres[i]->getRefPos();
      TVector3 hit = vres[i]->getHitPos();
      TVector3 hituvw = alMan->masterToLocal(TPCDET,hit);
      hituvw.SetZ((hituvw.Z()-MEANZ)*corrdrift+MEANZ);
      alMan->setEulerAngles(TPCDET, phi+phipar,theta+thetapar,psi+psipar);  //and we rotate
      alMan->setTransformation(TPCDET, tpcpos+vadd);                 //then translate
      TVector3 newxyz = alMan->localToMaster(TPCDET,hituvw);
      alMan->setTransformation(TPCDET, tpcpos);      //back translation
      alMan->setEulerAngles(TPCDET, phi,theta,psi);  //and we rotate it back
      TVector3 res = ref-newxyz;
      chisq += res.Mag()*res.Mag();
      sum+=hituvw.Z();
    }
  MEANZ = sum/(1.*nbins);
  f = chisq;
}

  
  /*



TpcAlignmentManager read file: macro/tpc/TestBench/COMPASS_alignment.latest.txt with 9 entries
TPC xyz(-4.7, -1.7, 110.3). Rot. angles(90, 72, -90)
NEW TPC xyz(-4.66663, -1.70941, 110.3). Rot. angles(89.9969, 71.9785, -90.0214)
TpcAlignmentManager read file: ./newal.txt with 9 entries
TPC xyz(-4.66663, -1.70941, 110.3). Rot. angles(89.9969, 71.9785, -90.0214)
NEW TPC xyz(-4.63325, -1.71881, 110.3). Rot. angles(89.9936, 71.9568, -90.0428)

   */