// -------------------------------------------------------------------------
// -----                CbmSttMinuitVertexFitter source file            -----
// -----                  Created 29/03/06  by R. Castelijns           -----
// -------------------------------------------------------------------------
#include <string>
#include <sstream>

#include "CbmSttMinuitVertexFitter.h"
#include "CbmRootManager.h"
#include "CbmTask.h"
#include "CbmSttTrack.h"
#include "CbmSttGeomCircle.h"
#include "CbmSttGeomPoint.h"
 
#include "TArc.h"
#include "TF1.h"
#include "TH2.h"
#include "TH3.h"
#include "TLine.h"
#include "TMarker.h"
#include "TGraphErrors.h"
#include "TMinuit.h"
    
#define rootoutput kFALSE
   using namespace std;
CbmSttMinuitVertexFitter::CbmSttMinuitVertexFitter()
{
  fNumberOfTracks = 0;
}

CbmSttMinuitVertexFitter::~CbmSttMinuitVertexFitter()
{
}

void CbmSttMinuitVertexFitter::Init()
{
  // Get and check CbmRootManager
  CbmRootManager
    *ioman = CbmRootManager::Instance();

  if (! ioman) 
    {
      cout << "-E- CbmSttMinuitVertexFitter::Init: "
	   << "RootManager not instantised!" << endl;
    }

  fSttTracks = (TClonesArray*) ioman->GetObject("STTTrack");
  if (! fSttTracks) 
    {
      cout << "-E- " << GetName() << "::Init: No SttTrack array!" << endl;
      return;
    }
}

CbmSttGeomPoint CbmSttMinuitVertexFitter::GetStartVertex(CbmSttTrack* pTrack1, CbmSttTrack* pTrack2)
{
  Double_t
    x1 = (pTrack1->GetParamLast()->GetTx() + pTrack1->GetParamLast()->GetX()) * cos(pTrack1->GetParamLast()->GetY()),
    x2 = (pTrack2->GetParamLast()->GetTx() + pTrack2->GetParamLast()->GetX()) * cos(pTrack2->GetParamLast()->GetY()),
    y1 = (pTrack1->GetParamLast()->GetTx() + pTrack1->GetParamLast()->GetX()) * sin(pTrack1->GetParamLast()->GetY()),
    y2 = (pTrack2->GetParamLast()->GetTx() + pTrack2->GetParamLast()->GetX()) * sin(pTrack2->GetParamLast()->GetY()),
    r1 = pTrack1->GetParamLast()->GetTx(),
    r2 = pTrack2->GetParamLast()->GetTx(),
    z1 = pTrack1->GetParamLast()->GetZ(),
    z2 = pTrack2->GetParamLast()->GetZ(),
    a1 = pTrack1->GetParamLast()->GetTy(),
    a2 = pTrack2->GetParamLast()->GetTy();

  CbmSttGeomCircle
    firstCircle(x1, y1, r1),
    secondCircle(x2, y2, r2);

  CbmSttGeomPoint
    vertexIntersection,
    firstIntersection,
    secondIntersection;

  Int_t
    numberOfIntersections = firstCircle.IntersectWith(secondCircle, firstIntersection, secondIntersection);

  Double_t
      dy1Int1, dx1Int1, dy2Int1, dx2Int1, phi1Int1, phi2Int1,
      dy1Int2, dx1Int2, dy2Int2, dx2Int2, phi1Int2, phi2Int2;

  if (numberOfIntersections > 1)
    {
      // get the z-difference at the first intersection
      dy1Int1 = firstIntersection.GetY() - firstCircle.GetYcenter();
      dx1Int1 = firstIntersection.GetX() - firstCircle.GetXcenter();
      dy2Int1 = firstIntersection.GetY() - secondCircle.GetYcenter();
      dx2Int1 = firstIntersection.GetX() - secondCircle.GetXcenter();

      phi1Int1 = atan(dy1Int1 / dx1Int1);
      phi2Int1 = atan(dy2Int1 / dx2Int1);

      if (dx1Int1 < 0.)
	{
	  if (dy1Int1 < 0.)
	    phi1Int1 -= dPi;
	  else
	    phi1Int1 += dPi;
	}
      if (dx2Int1 < 0.)
	{
	  if (dy2Int1 < 0.)
	    phi2Int1 -= dPi;
	  else
	    phi2Int1 += dPi;
	}

      
      Double_t
	zdiffFirstIntersection = (z1 + phi1Int1 * a1) - (z2 + phi2Int1 * a2);

      cout << "intersection1: " << (z1 + phi1Int1 * a1) << "   " << (z2 + phi2Int1 * a2) << endl;

      // get the z-difference at the second intersection
      dy1Int2 = secondIntersection.GetY() - firstCircle.GetYcenter();
      dx1Int2 = secondIntersection.GetX() - firstCircle.GetXcenter();
      dy2Int2 = secondIntersection.GetY() - secondCircle.GetYcenter();
      dx2Int2 = secondIntersection.GetX() - secondCircle.GetXcenter();

      phi1Int2 = atan(dy1Int2 / dx1Int2);
      phi2Int2 = atan(dy2Int2 / dx2Int2);

      if (dx1Int2 < 0.)
	{
	  if (dy1Int2 < 0.)
	    phi1Int2 -= dPi;
	  else
	    phi1Int2 += dPi;
	}
      if (dx2Int2 < 0.)
	{
	  if (dy2Int2 < 0.)
	    phi2Int2 -= dPi;
	  else
	    phi2Int2 += dPi;
	}

      cout << "intersection2: " << (z1 + phi1Int2 * a1) << "   " << (z2 + phi2Int2 * a2) << endl;
      Double_t
	zdiffSecondIntersection = (z1 + phi1Int2 * a1) - (z2 + phi2Int2 * a2);

      if (zdiffFirstIntersection < zdiffSecondIntersection)
	{
	  vertexIntersection = firstIntersection;
	  vertexIntersection.SetZ(((z1 + phi1Int1 * a1) + (z2 + phi2Int1 * a2)) / 2.);
	}    
      else
	{
	  vertexIntersection = secondIntersection;
	  vertexIntersection.SetZ(((z1 + phi1Int2 * a1) + (z2 + phi2Int2 * a2)) / 2.);
	}
    }
  else
    {
      // get the z-difference at the first intersection
      dy1Int1 = firstIntersection.GetY() - firstCircle.GetYcenter();
      dx1Int1 = firstIntersection.GetX() - firstCircle.GetXcenter();
      dy2Int1 = firstIntersection.GetY() - secondCircle.GetYcenter();
      dx2Int1 = firstIntersection.GetX() - secondCircle.GetXcenter();

      phi1Int1 = atan(dy1Int1 / dx1Int1);
      phi2Int1 = atan(dy2Int1 / dx2Int1);

      if (dx1Int1 < 0.)
	{
	  if (dy1Int1 < 0.)
	    phi1Int1 -= dPi;
	  else
	    phi1Int1 += dPi;
	}
      if (dx2Int1 < 0.)
	{
	  if (dy2Int1 < 0.)
	    phi2Int1 -= dPi;
	  else
	    phi2Int1 += dPi;
	}

      vertexIntersection = firstIntersection;
      vertexIntersection.SetZ(((z1 + phi1Int1 * a1) + (z2 + phi2Int1 * a2)) / 2.);
    }

  return vertexIntersection;
}


CbmSttTrack *CbmSttMinuitVertexFitter::GetTrackPtr(Int_t index) const
{
  return tracks[index]; 
}


Int_t CbmSttMinuitVertexFitter::GetNumberOfTracks() const
{
  return fNumberOfTracks;
}


Int_t CbmSttMinuitVertexFitter::DoFit(CbmSttVertex* pVertex)
{   
    tracks[0] = (CbmSttTrack*) fSttTracks->At(pVertex->GetTrack(0)),
    tracks[1] = (CbmSttTrack*) fSttTracks->At(pVertex->GetTrack(1));

    // get the first approximation to the vertex
    CbmSttGeomPoint
	firstGuessVertex = GetStartVertex(tracks[0], tracks[1]);

    cout << "pre-vertex: " << firstGuessVertex.GetX() << "   " << firstGuessVertex.GetY() << "   " << firstGuessVertex.GetZ() << endl;

    // minimize over 2 tracks * 5 parameters + 1 vertex * 3 parameters = 13
    TMinuit 
      minimizer(13);

    // minimization setup:
    minimizer.SetObjectFit(this);
    //    minimizer.SetPrintLevel(-1);
    minimizer.SetMaxIterations(500);
    minimizer.SetFCN(fcnVertex);
    minimizer.SetErrorDef(1);

    // set the vertex parameters:
    minimizer.DefineParameter(0, "Vx", firstGuessVertex.GetX(), 0.05, 0., 0.);
    minimizer.DefineParameter(1, "Vy", firstGuessVertex.GetY(), 0.05, 0., 0.);
    minimizer.DefineParameter(2, "Vz", firstGuessVertex.GetZ(), 0.20, 0., 0.);

    fNumberOfTracks = pVertex->GetNofTracks();
    // set the track parameters:
    for (Int_t trackCounter = 0; trackCounter < fNumberOfTracks; trackCounter++)
    {
	minimizer.DefineParameter(3 + trackCounter * 5, "d0", tracks[trackCounter]->GetParamLast()->GetX(), 
				  sqrt(tracks[trackCounter]->GetParamLast()->GetCovariance(0, 0)), 0., 0.); 
	minimizer.FixParameter(3 + trackCounter * 5);

	minimizer.DefineParameter(4 + trackCounter * 5, "phi", tracks[trackCounter]->GetParamLast()->GetY(), 
				  sqrt(tracks[trackCounter]->GetParamLast()->GetCovariance(1, 1)), 0., 0.); 
	minimizer.FixParameter(4 + trackCounter * 5);

	minimizer.DefineParameter(5 + trackCounter * 5, "r", tracks[trackCounter]->GetParamLast()->GetTx(), 
				  sqrt(tracks[trackCounter]->GetParamLast()->GetCovariance(2, 2)), 0., 0.); 
	minimizer.FixParameter(5 + trackCounter * 5);

	minimizer.DefineParameter(6 + trackCounter * 5, "z0", tracks[trackCounter]->GetParamLast()->GetZ(), 
				  sqrt(tracks[trackCounter]->GetParamLast()->GetCovariance(3, 3)), 0., 0.); 
	minimizer.FixParameter(6 + trackCounter * 5);

	minimizer.DefineParameter(7 + trackCounter * 5, "theta", tracks[trackCounter]->GetParamLast()->GetTy(), 
				  sqrt(tracks[trackCounter]->GetParamLast()->GetCovariance(4, 4)), 0., 0.); 
	minimizer.FixParameter(7 + trackCounter * 5);
    }

    // minimize:
    minimizer.Migrad();

    Double_t
      vertexX, vertexY, vertexZ,
      errorX, errorY, errorZ;

    minimizer.GetParameter(0, vertexX, errorX);
    minimizer.GetParameter(1, vertexY, errorY);
    minimizer.GetParameter(2, vertexZ, errorZ);

//    pVertex->SetVertex(firstGuessVertex.GetX(), firstGuessVertex.GetY(), firstGuessVertex.GetZ());
    pVertex->SetVertex(vertexX, vertexY, vertexZ);

    cout << "final vertex: " << vertexX << "   " << vertexY << "   " << vertexZ << endl;
    return 0;
}

void fcnVertex(Int_t &npar, Double_t *gin, Double_t &f, Double_t *par, Int_t iflag)
{
  Bool_t
    verbose = 0;

  const CbmSttMinuitVertexFitter 
    *mama = (CbmSttMinuitVertexFitter *)gMinuit->GetObjectFit();

  Double_t
    chisq = 0.;

  CbmSttGeomPoint
    vertexPoint(par[0], par[1], par[2]);

  if (verbose)
      cout << "vertexpoint: " << par[0] << "   " << par[1] << "   " << par[2] << endl;

  for (Int_t trackCounter = 0; trackCounter < mama->GetNumberOfTracks(); trackCounter++)
    {
      Double_t
	x = (par[3 + trackCounter * 5] + par[5 + trackCounter * 5]) * cos(par[4 + trackCounter * 5]),
	y = (par[3 + trackCounter * 5] + par[5 + trackCounter * 5]) * sin(par[4 + trackCounter * 5]),
	r = par[5 + trackCounter * 5];
      
      CbmSttGeomCircle
	trackCircle(x, y, r);
            
      Double_t
	Dr = trackCircle.DistanceTo(vertexPoint) / 0.0050;

      Double_t
	dy = vertexPoint.GetY() - trackCircle.GetYcenter(),
	dx = vertexPoint.GetX() - trackCircle.GetXcenter(),
	phi = atan(dy / dx);
      
      if (dx < 0.)
	{
	  if (dy < 0.)
	    phi -= dPi;
	  else
	    phi += dPi;
	}

      Double_t
	Dz = (vertexPoint.GetZ() - (par[6 + trackCounter * 5] + phi * par[7 + trackCounter * 5])) / 0.0100;

      // chi2 due to distance to vertex
      chisq += (Dr * Dr + Dz * Dz) * 1000.;

      if (verbose)
      {
	  cout << " Dr: " << trackCircle.DistanceTo(vertexPoint) 
	       << " Dz: " << vertexPoint.GetZ() - (par[6 + trackCounter * 5] + phi * par[7 + trackCounter * 5])
	       << endl;

	  char
	      waitchar;

	  cout << "press key to continue" << endl;

	  cin >> waitchar;
      }

      /*
      // chi2 due to change of helix parameters
      for (Int_t rowCounter; rowCounter < 5; rowCounter++)
	{
	  for (Int_t colCounter; colCounter < 5; colCounter++)
	    {
	      chisq += ((par[3 + trackCounter * 5 + rowCounter] - mama->GetTrackPar(trackCounter, rowCounter)) *
			(par[3 + trackCounter * 5 + colCounter] - mama->GetTrackPar(trackCounter, colCounter)) *
			mama->GetTrackPtr(trackCounter)->GetParamLast()->GetCovariance(rowCounter, colCounter) * 0.5);
	    }
	}
      */
    }

  f = chisq;
  if (verbose)
      cout << "chisq: " << f << endl;
}

Double_t CbmSttMinuitVertexFitter::GetTrackPar(Int_t trackNo, Int_t index) const
{
  CbmSttTrack
    *myTrack = tracks[trackNo];

  Double_t
    retval = 0.;

  if (index == 0)
    retval = myTrack->GetParamLast()->GetX();
  if (index == 1)
    retval = myTrack->GetParamLast()->GetY();
  if (index == 2)
    retval = myTrack->GetParamLast()->GetTx();
  if (index == 3)
    retval = myTrack->GetParamLast()->GetZ();
  if (index == 4)
    retval = myTrack->GetParamLast()->GetTy();

  return retval;
}

ClassImp(CbmSttMinuitVertexFitter)