//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class PndRiemannTrack
//      see PndRiemannTrack.hh for details
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sebastian Neubert    TUM            (original author)
//
//
//-----------------------------------------------------------

// Panda Headers ----------------------

// This Class' Header ------------------
#include "PndRiemannTrack.h"

// C/C++ Headers ----------------------
#include <iostream>
#include <assert.h>

// Collaborating Class Headers --------
#include "PndRiemannHit.h"
#include "TGraph.h"
#include "TMath.h"
#include "TMatrixTSym.h"
// Class Member definitions -----------

void MatrixOutput(TMatrixD mat)
{
	for (int col = 0; col < mat.GetNcols(); col++)
	{
		for (int row = 0; row < mat.GetNrows(); row++)
		{
			std::cout << mat[row][col] << " ";
		}
		std::cout << std::endl;
	}
}

ClassImp(PndRiemannTrack)

PndRiemannTrack::PndRiemannTrack()
  : fn(3),fav(3), fc(0), fcovPlane(4,4), fjacRXY(3,4), fcovRXY(3,3)
{}


void
PndRiemannTrack::init(double x0_, double y0_, double R_, 
			 double dip, double z0){
  double x0=x0_/100;
  double y0=y0_/100;
  double R=R_/100;

  double R2=R*R;
  double lambda=R2-x0*x0-y0*y0;
  double D=lambda/sqrt(4*R2+(1-lambda)*(1-lambda));
  double DC=-D/lambda;
  double C=DC-D;
  double A=-2*x0*DC;
  double B=-2*y0*DC;
  
  fn[0]=A;
  fn[1]=B;
  fn[2]=C;
  fc=D;
}


void
PndRiemannTrack::addHit(PndRiemannHit* hit){
  int nbefore=fHits.size();
  fHits.push_back(hit);
  //fav*=(double)nbefore;
  fav[0]+=hit->x().X() /(hit->sigmaXY()*hit->sigmaXY());
  fav[1]+=hit->x().Y() /(hit->sigmaXY()*hit->sigmaXY());;
  fav[2]+=hit->x().Z() /(hit->sigmaXY()*hit->sigmaXY());;
  //fav*=1./(double)(nbefore+1);
  fweight += 1/(hit->sigmaXY()*hit->sigmaXY());
  
}

double
PndRiemannTrack::dist(PndRiemannHit* hit){
  double d2=fc;
  d2+=hit->x().X()*fn[0];
  d2+=hit->x().Y()*fn[1];
  d2+=hit->x().Z()*fn[2];
  return d2;
}


void
PndRiemannTrack::refit()
{
  
  TMatrixT<double> av(3,1);
  std::cout << "fweight: " << fweight << std::endl;
  fav *= TMath::Power(fweight,-1); //TS
  av[0][0]=fav[0];// *fweight; //TS *fweight added
  av[1][0]=fav[1];// *fweight;
  av[2][0]=fav[2];// *fweight;
  
  std::cout << "av: " << fav[0] << " " << fav[1] << " " << fav[2] << std::endl;

  TMatrixD sampleCov(3,3);
  
  int nh=fHits.size();
  for(int i=0;i<nh;++i){
    TMatrixD h(3,1);
    h[0][0]=fHits[i]->x().X();
    h[1][0]=fHits[i]->x().Y();
    h[2][0]=fHits[i]->x().Z();
    TMatrixD d(3,1);
    d=h-av;
    TMatrixD dt(TMatrixD::kTransposed,d);
    TMatrixD ddt(d,TMatrixD::kMult,dt);
    ddt *= 1/(fHits[i]->sigmaXY()*fHits[i]->sigmaXY());  //TS
    sampleCov+=ddt;  
  }

  sampleCov*=1./(double)nh;
  
  std::cout << "sampleCov: " << std::endl;
  MatrixOutput(sampleCov);
  
  TVectorD eigenValues(3);
  TMatrixD eigenVec=sampleCov.EigenVectors(eigenValues);
    
  // find smallest eigenvalue
  double val=eigenValues[0]; int g=0;
  for(int k=1;k<3;++k){
    if(eigenValues[k]<val){
      val=eigenValues[k];
      g=k;
    }
  }
  
  fn=TMatrixDColumn(eigenVec,g);
  
  std::cout << "eigenVec: " << std::endl;
  MatrixOutput(eigenVec);
  
  std::cout << "eigenValues: ";
  for (int i = 0; i < 3; i++){
	  std::cout << eigenValues[i] << " ";
  }
  std::cout << std::endl;
 
  double norm=1./TMath::Sqrt(fn.Norm2Sqr());
  fn*=norm;
  fc=-1.*fn*fav;
  
  std::cout << "fn: " << fn[0] << " " << fn[1] << " " << fn[2] << std::endl;
  std::cout << "fc: " << fc << std::endl;
  //------------ Calculation of the full covariance matrix -----------
  
  // 1. Calculation of the covarianz matrix of the normal vector
  
  TMatrixD  CovNorm(3,3);
  
  for (int i = 0; i < 3; i++){
	if (i != g){
		TMatrixD res = eigenVec.GetSub(0,2,i,i)*eigenVec.GetSub(0,2,i,i).T();
		CovNorm += res * (val * eigenValues[i]/TMath::Power(val-eigenValues[i],2));
	}
  }
  if (fHits.size() > 5)
	  CovNorm *= TMath::Power(fHits.size()-5,-1);  //-5 is very strange. According to the paper this value has to be fixed with simulation???
  else 
	  CovNorm *= TMath::Power(fHits.size(),-1);
  std::cout << "CovNorm: " << std::endl;
  MatrixOutput(CovNorm);

  
  // 2. Calculation of the covarianz matrix of fav
  
  TMatrixD covAv(3,3);
  for (int i = 0; i < fHits.size();i++){
	  covAv[0][0] += fHits[i]->covX(0,0)/(TMath::Power(fHits[i]->sigmaXY(),4));
	  covAv[1][1] += fHits[i]->covX(1,1)/(TMath::Power(fHits[i]->sigmaXY(),4));
	  covAv[1][0] += fHits[i]->covX(1,0)/(TMath::Power(fHits[i]->sigmaXY(),4));
	  covAv[2][0] += 2 * (fHits[i]->covX(0,0) * fHits[i]->x().X() + fHits[i]->covX(1,0) * fHits[i]->x().Y())
	  				 /(TMath::Power(fHits[i]->sigmaXY(),4));
	  covAv[2][1] += 2 * (fHits[i]->covX(1,1) * fHits[i]->x().Y() + fHits[i]->covX(1,0) * fHits[i]->x().X())
		             /(TMath::Power(fHits[i]->sigmaXY(),4));
	  covAv[2][2] += 4 * (fHits[i]->covX(0,0) * fHits[i]->x().X() * fHits[i]->x().X() +
			  			  2 * fHits[i]->covX(1,0) * fHits[i]->x().X() * fHits[i]->x().Y() +
			  			  fHits[i]->covX(1,1) * fHits[i]->x().Y() * fHits[i]->x().Y())
			  		 /(TMath::Power(fHits[i]->sigmaXY(),4));
  }
  covAv[0][1] = covAv[1][0];
  covAv[0][2] = covAv[2][0];  
  covAv[1][2] = covAv[2][1];
  
  covAv *= TMath::Power(fweight,-2);

  
  std::cout << "covAv: " << std::endl;
  MatrixOutput(covAv);
  
  
  // 3. Calculation of var(fc)
  
  double nCrn = (fn * (covAv * fn));
  double rCnr = (fav * (CovNorm * fav));
  
  TMatrixD CnCr(CovNorm,TMatrixD::kMult,covAv);
  double trCnCr = 0;
  for (int i = 0; i < CnCr.GetNcols(); i++){
	  trCnCr += CnCr[i][i];
  }
  std::cout << "nCrn: " << nCrn << " rCnr: " << rCnr << " trCnCr: " << trCnCr << std::endl;
  double varc = nCrn + rCnr + trCnCr;

  TVectorD corr_cn = CovNorm * fav;
  corr_cn *= -1;
  
  // 4. Calculation of the covarianz matrix of c,n
  
  fcovPlane[0][0] = varc;
  fcovPlane[1][1] = CovNorm[0][0];
  fcovPlane[2][1] = CovNorm[1][0];
  fcovPlane[2][2] = CovNorm[1][1];
  fcovPlane[3][1] = CovNorm[2][0];
  fcovPlane[3][2] = CovNorm[2][1];
  fcovPlane[3][3] = CovNorm[2][2];
  fcovPlane[1][0] = corr_cn[0];
  fcovPlane[2][0] = corr_cn[1];
  fcovPlane[3][0] = corr_cn[2];
  
  //5. Convert plane covariance in start parameter(r,x0,y0) covariances
  calcJacRXY();
  
  std::cout << "jacRXY: " << std::endl;
  MatrixOutput(fjacRXY);
  
  //fcovRXY = fjacRXY * fcovPlane * fjacRXY.T();
  
  TMatrixD temp(fjacRXY,TMatrixD::kMult, fcovPlane);
  std::cout << "temp: " << temp.GetNrows() << "x" << temp.GetNcols() << std::endl;
  
  std::cout << "temp: " << std::endl;
  MatrixOutput(temp);
  
  fcovRXY = temp * fjacRXY.T();  // J * fcovPlane * J.T()

  std::cout << "covPlane: " << std::endl;
  MatrixOutput(fcovPlane);
  
  std::cout << "covRXY: " << std::endl;
  MatrixOutput(fcovRXY);
}


TVectorD
PndRiemannTrack::orig() const {
  TVectorD o(2);
  double den=TMath::Power((2 * fn[2]),-1);  // modified for paraboloid
  o[0]=-fn[0]*den;
  o[1]=-fn[1]*den;
  return o;
}


double 
PndRiemannTrack::r() const {
  if(fc==0)return 0;
  if(fabs(fc)>100)return 0;
  double a=2.*fn[2];  // modified for paraboloid
  //if(a==0){
  //  std::cout<<"PndRiemannTrack:: a==0 cannot calc r! set r=1E4"<<std::endl;
  //  return 1.E-4;
  //}
  //std::cout<<fn[0]<<"   "<<fn[1]<<"   "<<fc<<"   "<<a<<std::endl;
  double nom = 1-fn[2]*fn[2]-4*fc*fn[2];  // modified for paraboloid
  if(nom<0.0){
    nom=1.E-4;
    std::cout<<"PndRiemannTrack::nom<0! set r=1E-4!"<<std::endl;
  }
  return sqrt(nom)/TMath::Abs(a);
}

void
PndRiemannTrack::szFit(){
  TGraph g(getNumHits());
  // get s'es and zs
  unsigned int n=getNumHits();
  for(unsigned int i=0;i<n;++i){
    fHits[i]->calcPosOnTrk(this);
    g.SetPoint(i,fHits[i]->s(),fHits[i]->z());
    //std::cout<<fHits[i]->z()<<"|"<<fHits[i]->s()<<std::endl;
  }
  int errorcode;
  g.LeastSquareLinearFit(n,ft,fm,errorcode,-999,999);
  //std::cout<<"szFit Error Code:"<<errorcode<<std::endl;
  return;
}

double
PndRiemannTrack::szDist(PndRiemannHit* hit){
  szFit();
  hit->calcPosOnTrk(this);
  double hits=hit->s();
  double predz=hits*fm+ft;
  return predz-hit->z();
}

// only after szFit!
double
PndRiemannTrack::dip() const {
  return cos(atan(fm));
}


// Todo: check for backward going tracks!!!
double
PndRiemannTrack::sign() const {
  double s=fHits[3]->s(); 
  return s<0 ? -1. : 1;
}

void PndRiemannTrack::calcJacRXY()
{
	double val = (TMath::Sqrt(1-fn[2]*fn[2]-4*fc*fn[2]));
	fjacRXY[0][0] = -1/val;
	fjacRXY[0][3] = -1/(2*fn[2]) * (fn[2]+2*fc)/val - val/(2*fn[2]*fn[2]);
	fjacRXY[1][1] = -fn[0]/(2*fn[2]);
	fjacRXY[1][3] = fn[0]/(2*fn[2]*fn[2]);
	fjacRXY[2][2] = -fn[1]/(2*fn[2]);
	fjacRXY[2][3] = fn[1]/(2*fn[2]*fn[2]);
}