//////////////////////////////ALIGNMENT OF TPC PROTOTYPE WITH RESPECT TO THE CDC TRACKS/////////////////////////
//Minized residuals between TpcClusters and GFTracks
//Residuals are calculated via TpcRefGFTrkResCalc
//Translations/Rotations of the TPC is done via the TpcAlignment Manager
//Name of the TPC ("tpc") is hardcoded here
//
//Plots in world coordinates, blue is corrected, red is original distributions
//
//New alignement saved in ./newal.txt
//
//Needs to be compiled!
//gROOT->ProcessLine(".x $VMCWORKDIR/gconfig/rootlogon.C");
//Typical command in root : .x macro/tpc/FOPI/alignFopiPrototype.C++("/nfs/hicran/scratch/user/fcusanno/rootfiles/tmp/run_3282.reco.root","tpc/parfiles/june11_alignment.txt",5000)
//Author: Tiger. Adapted to the GEM-TPC prototype in FOPI by Francesco Cusanno, T.U.M.

#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 "TpcAlignmentManager.h"
//#include "tpc/TestBench/TBStripCluster.h"
#include "TpcCluster.h"
//#include "tpc/TestBench/TBGEMCluster.h"
//#include "tpc/TestBench/TBSICluster.h"
#include "TpcResidual.h"
#include "TpcRefResidualCollection.h"
#include "TpcResidualCollection.h"
#include "TpcRefGFTrkResCalc.h"
#include "TpcCdcFit2DResCalc.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 "RKTrackRep.h"
#include "/nfs/hicran/project/panda/SIM/fcusanno/dst2root/src/fopi2root/CdcTrack.h"

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

#define DEBUG 1
//
const double MAXRES = 8.; //IMPORTANT: cut on maximal distance bet clusters and tracks taking into account for al.
const float MaxResXY = 1.;//IMPORTANT: cut on maximal distance bet clusters and tracks taking into account for al.
const float MaxResZ = 5.;//IMPORTANT: cut on maximal distance bet clusters and tracks taking into account for al.
const char* TPCDET = "tpc";
//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 alignFopiPrototype(TString filename, std::string alfile = "tpc/parfiles/june11_alignment.txt", unsigned int EvMax=0)
{
  //NOT REALLY NEEDED
  //TGeoManager* geom = new TGeoManager("Geometry", "geometry manager");
  //TGeoManager::Import("/nfs/hicran/project/panda/SIM/fcusanno/devfopiroot/tpc/FOPI/geometry/FopiGeom_data.root");
  //
  //VERY PROBABLY NOT NEEDED IF TpcCdcFit2DResCalc CALCULATOR IS USED
  GFFieldManager::getInstance()->init(new GFConstField(0.,0.,6.16));
  
  while(gROOT->GetListOfCanvases()->GetSize())
    delete gROOT->GetListOfCanvases()->At(0);
  gStyle->SetOptStat(10);
  gStyle->SetPalette(1);
  // int a = 0;
  //Opening input file
  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;  
  }
  //Store branches in map
  std::map<TString, TClonesArray*> Tbranches;
  std::cout << "ARRAYS..." << std::endl;
  //I CAN HAZ ARRAYS: 3
  TClonesArray* trCdcArr=new TClonesArray("CdcTrack");
  recotree->SetBranchAddress("CdcTrack", &trCdcArr);

  Tbranches["CdcTrack"] = trCdcArr;
  if(trCdcArr==0) {
    std::cout<< "Init CDC track 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;
  }
  // NOT REALLY USED:
  /*
  TClonesArray* tpcResArr=new TClonesArray("TpcRefResidualCollection");
  recotree->SetBranchAddress("CdcTpcResiduals", &tpcResArr);
  if(tpcResArr==0) {
    std::cout<< "Init TpcResidualCollection array not found!"<< std::endl;
    return;
  }
  */
  //
  //Initialization of alignment manager
  std::cout << "ALMAN..." << std::endl;
  TpcAlignmentManager::init(alfile.c_str());
  alMan = TpcAlignmentManager::getInstance();
  if(alMan==0) {
    std::cout<< "Init TpcAlignmentManager:"<< alfile.c_str()<<" not found!"<< std::endl;
    return;
  }
  
  //Getting initial state of TPC
  double phi,theta,psi;   
  std::cout << "EULER..." << std::endl;
  alMan->getEulerAngles(TPCDET, phi, theta, psi);
  std::cout << "TRANS..." << std::endl;
  TVector3 tpcpos=  alMan->getTranslation(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 << "Tuning test passed." << std::endl;
  
  //Quality assurance histograms
  TH1F * hresx = new TH1F("hresx","hresx",800,-MaxResXY,MaxResXY);
  hresx->SetLineColor(kRed);
  TH1F * hresy = new TH1F("hresy","hresy",800,-MaxResXY,MaxResXY);
  hresy->SetLineColor(kRed);
  TH1F * hresxy = new TH1F("hresxy","Resolution on XY plane",400,0,MaxResXY);
  hresxy->SetLineColor(kRed);
  TH1F * hresz = new TH1F("hresz","hresz",800,-MaxResZ,MaxResZ);hresz->SetLineColor(kRed);
  TH1F * hnesx = new TH1F("hnesx","hnesx",800,-MaxResXY,MaxResXY);hnesx->SetLineColor(kBlue);
  TH1F * hnesy = new TH1F("hnesy","hnesy",800,-MaxResXY,MaxResXY);hnesy->SetLineColor(kBlue);
  TH1F * hnesxy = new TH1F("hnesxy","Resolution on XY plane",400,0,MaxResXY);
  hnesxy->SetLineColor(kBlue);
  TH1F * hnesz = new TH1F("hnesz","hnesz",800,-MaxResZ,MaxResZ);hnesz->SetLineColor(kBlue);

  //********************************************************************************
  //Starting calculating of residuals
  //********************************************************************************
  std::cout << "CALCULATOR : " << std::endl;
  //Initializing the RefTrackResCalc
  //TpcRefGFTrkResCalc* resgft = new TpcRefGFTrkResCalc();
  TpcCdcFit2DResCalc* resgft = new TpcCdcFit2DResCalc();
  resgft->setAlignment(kTRUE);
  //Adding branch names
  int test =   resgft->addBranchName("TpcCluster");
  if(test) {
    std::cout << "resgft->addBranchName(\"TpcCluster\"); is inconsistent" << std::endl;
  }
  test = resgft->addBranchName("CdcTrack");
  if(test) {
    std::cout << "resgft->addBranchName(\"CdcTrack\"); is inconsistent" << std::endl;
  }
  //arbitrary vertical plane downstream of the detector, probably not used by the calc
  TVector3 uu(1.,0.,0.);  TVector3 vv(0.,1.,0.);  TVector3 oo(0.,0.,105.);
  //resgft->setProjectionPlane(oo,uu,vv);
  //resgft->setTrackRepId(0);
  resgft->setDistanceCut(MaxResXY);
  //Adding Branches to the residual calculator
  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;
  } 
  //Initializing, searching for set branches
  std::cout << "Self awareness ... " << std::endl;
  resgft->init();
  //
  unsigned int MAX =recotree->GetEntries();
  std::cout << "Number of events in the input file: " << MAX  <<std::endl;
  if(!EvMax)EvMax=MAX;
  
  //Calculating residuals event by event
  vres.clear();
  for(unsigned int k=0; k<MAX && k<EvMax; k++)   {
      //get the event
      recotree->GetEvent(k);
      //int ntpccl = tpcArr->GetEntries();
      
      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 RESIDUALS
	  if(miaou->getResXYZ().Perp()>MaxResXY) continue; //CUT ON RESIDUALS
	  if(miaou->getResXYZ().Z()>MaxResZ) continue; //CUT ON RESIDUALS
          vres.push_back(miaou);
          hresx->Fill(miaou->getResXYZ().X());
          hresy->Fill(miaou->getResXYZ().Y());
	  hresxy->Fill(miaou->getResXYZ().Perp());
          hresz->Fill(miaou->getResXYZ().Z());
        }
      }
  }//end event loop 
  std::cout << vres.size()<<" TpcResiduals, Starting sattelites alignment..." << std::endl; \
  
  // CHECK THE PARAMETER LIMITS. ANGLES ARE IN DEGREES!
  TMinuit *gMinuit;
  gMinuit = new TMinuit(6); //currently a max of 6 parameters is set
  gMinuit->SetFCN(fcn);
  gMinuit->DefineParameter(0, "x", 0., 0.00001, -1.5,1.5);
  gMinuit->DefineParameter(1, "y", 0., 0.00001, -1.5,1.5);
  gMinuit->DefineParameter(2, "phi", 0., 0.000001, -3., 3.);//-3.
  gMinuit->DefineParameter(3, "theta", 0., 0.00001, -3., 3.);
  gMinuit->DefineParameter(4, "psi", 0., 0.00001, -3., 3.);
  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);
  */  
  std::cout <<"##############################################################\n"
	    <<"------------ Params all released but drift --------------------" << std::endl;
  gMinuit->FixParameter(5);
  gMinuit->Migrad();
  Double_t amin,edm,errdef;
  Int_t nvpar,nparx,icstat;
  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.);
  //********************************************************************************
  //Set final alignment for TPC and recalculate residuals
  //********************************************************************************

  alMan->setEulerAngles(TPCDET, nphi+phi,ntheta+theta,psi+npsi);  //we rotate
  alMan->setTranslation(TPCDET, ntpcpos);        //then translate
  alMan->getEulerAngles(TPCDET, phi, theta, psi); 
  tpcpos = alMan->getTranslation(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);

      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 residuals 
	      if(miaou->getResXYZ().Perp()>MaxResXY) continue; //cut on residuals
	      if(miaou->getResXYZ().Z()>MaxResZ) continue; //cut on residuals
	      hnesx->Fill(miaou->getResXYZ().X());
	      hnesy->Fill(miaou->getResXYZ().Y());
	      hnesxy->Fill(miaou->getResXYZ().Perp());
	      hnesz->Fill(miaou->getResXYZ().Z());
	    }
	}
    }//end event loop 
  
  std::cout << "PLOTTING RESULTS... " << std::endl;
  
  char tmpChar[500];
  //
  TCanvas * cx = new TCanvas("cx","cx",1200,600);
  cx->Divide(3);
  cx->cd(1);
  sprintf(tmpChar,"X: Initial RMS = %f, Final RMS = %f",hresx->GetRMS(),hnesx->GetRMS());
  hnesx->SetTitle(tmpChar);
  hnesx->GetXaxis()->SetTitle("X residuals [cm]");
  hnesx->Draw();      hresx->Draw("same");
  cx->cd(2);
  sprintf(tmpChar,"Y: Initial RMS = %f, Final RMS = %f",hresy->GetRMS(),hnesy->GetRMS());
  hnesy->SetTitle(tmpChar);
  hnesy->GetXaxis()->SetTitle("Y residuals [cm]");
  hnesy->Draw();hresy->Draw("same");
  cx->cd(3);
  sprintf(tmpChar,"Z: Initial RMS = %f, Final RMS = %f",hresz->GetRMS(),hnesz->GetRMS());
  hnesz->SetTitle(tmpChar);
  hnesz->GetXaxis()->SetTitle("Z residuals [cm]");
  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();

  TCanvas * cxy = new TCanvas("cxy","cxy",700,600);
  cxy->cd();
  sprintf(tmpChar,"XY: Initial Mean = %f, Final Mean = %f",hresxy->GetMean(),hnesxy->GetMean());
  hnesxy->SetTitle(tmpChar);
  hnesxy->GetXaxis()->SetTitle("XY residuals [cm]");
  hnesxy->Draw(); hresxy->Draw("same");
  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");
 std:cout << "Saving the correction factor to drift velocity in ./DriftVel.txt" <<std::endl;
  TString OutParFileName = "DriftVel.txt";
  std::ofstream OutParFile(OutParFileName.Data());
  OutParFile<<"correction factor for drift velocity: " << ndr << std::endl;
}//end of macro

  
//Minuit function to be minimized
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); //translation params
  Double_t chisq = 0;
  Double_t sum =0.;
  double phi,theta,psi;  
  alMan->getEulerAngles(TPCDET, phi, theta, psi);
  TVector3 tpcpos = alMan->getTranslation(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); // adjust Z coordinate
    alMan->setEulerAngles(TPCDET, phi+phipar,theta+thetapar,psi+psipar);  //and we rotate
    alMan->setTranslation(TPCDET, tpcpos+vadd);                 //then translate
    TVector3 newxyz = alMan->localToMaster(TPCDET,hituvw);
    alMan->setTranslation(TPCDET, tpcpos);      //back translation
    alMan->setEulerAngles(TPCDET, phi,theta,psi);  //and we rotate it back
    TVector3 res = ref-newxyz;
    //chisq += res.Mag()*res.Mag();
    //chisq += pow(res.Perp(),2)+pow(res.Z(),2)/10. ;
    chisq += pow(res.Perp(),2);
    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)

   */