#include "user/TTrackFitter.hh"
#include "user/TAlignmentManager.hh"

#include <TMath.h>
#include <Math/MinimizerOptions.h>
#include <Math/WrappedMultiTF1.h>

TrackMinimizer::TrackMinimizer(const TWireHitInfoCollection* wireHits,
			       const TFibreHitInfoCollection*   fibreHits,
			       const TF1* positionResolution) :
  fWireHits(wireHits),
  fFibreHits(fibreHits),
  fPositionResolution(positionResolution)
{
}

TrackMinimizer::~TrackMinimizer()
{
}

Int_t TrackMinimizer::GetDataSize() const
{
  return
    (fWireHits? fWireHits->size() : 0) +
    (fFibreHits?   fFibreHits->size()   : 0);
}

Double_t TrackMinimizer::CalcDistance(const Int_t ch, 
				      const Double_t slope,
				      const Double_t offset,
				      Double_t* err,
				      const Double_t slopeErr,
				      const Double_t offsetErr)
{
  TVector2 pos = TAlignmentManager::Get().Get2dPos(TAlignmentManager::kWire, ch);
  Double_t z = pos.X();
  Double_t x = pos.Y();
  Double_t d = std::sqrt(slope*slope+1.0);
  Double_t r = (z*slope + offset - x) / d;
  if(err){
    Double_t dOffset = 1.0/d;
    Double_t dSlope  = z/d - slope*r/(d*d);
    *err = std::sqrt(dSlope*dSlope*slopeErr*slopeErr + dOffset*dOffset*offsetErr*offsetErr);
  }    
  return r;
}

void TrackMinimizer::GetSlopeAndOffsetCandidates(TWireHitInfo hit1, TWireHitInfo hit0, 
						 Double_t *slope, Double_t *offset)
{
  // Referred http://hg.hatenablog.jp/entry/2016/02/08/210906
  TAlignmentManager& aman = TAlignmentManager::Get();
  TVector2 pos0 = aman.Get2dPos(TAlignmentManager::kWire, hit0.strCh);
  TVector2 pos1 = aman.Get2dPos(TAlignmentManager::kWire, hit1.strCh);
  Double_t x[2] = {pos0.X(), pos1.X()};
  Double_t y[2] = {pos0.Y(), pos1.Y()};
  Double_t p  = x[1] - x[0];
  Double_t q  = y[1] - y[0];
  Double_t ds = p*p + q*q;
  Double_t r1 = hit0.distance;
  Double_t r2 = hit1.distance;

#define SQ(a) ((a)*(a))
  Double_t alpha[4] = {
    ( p*r1*(r1+r2) + q*r1*std::sqrt( ds - SQ(r1+r2) ) ) / ds,
    ( p*r1*(r1+r2) - q*r1*std::sqrt( ds - SQ(r1+r2) ) ) / ds,
    ( p*r1*(r1-r2) + q*r1*std::sqrt( ds - SQ(r1-r2) ) ) / ds,
    ( p*r1*(r1-r2) - q*r1*std::sqrt( ds - SQ(r1-r2) ) ) / ds  };
  Double_t beta[4] = {
    ( q*r1*(r1+r2) - p*r1*std::sqrt( ds - SQ(r1+r2) ) ) / ds,
    ( q*r1*(r1+r2) + p*r1*std::sqrt( ds - SQ(r1+r2) ) ) / ds,
    ( q*r1*(r1-r2) - p*r1*std::sqrt( ds - SQ(r1-r2) ) ) / ds,
    ( q*r1*(r1-r2) + p*r1*std::sqrt( ds - SQ(r1-r2) ) ) / ds };

  for(Int_t i=0; i<4; i++){
    slope[i]  = -alpha[i] / beta[i];
    offset[i] = (SQ(r1) + alpha[i]*x[0]) / beta[i] + y[0];
  }
#undef SQ
}

Double_t TrackMinimizer::GetSigmaR(const Double_t x)
{
  if(fPositionResolution)
    return 1e-3 * fPositionResolution->Eval(std::abs(x)/10.0);
  else 
    return 0.200;
}

Double_t TrackMinimizer::operator () (const Double_t *par)
{
  const Double_t& slope  = par[0]; 
  const Double_t& offset = par[1];
  Double_t chi2 = 0.0;
  
  // For wire hits
  Double_t dR, sigmaR;
  for (TWireHitInfoCollection::const_iterator 
	 it=fWireHits->begin(); it!=fWireHits->end(); it++) {
    dR     = it->distance - std::abs(CalcDistance(it->strCh, slope, offset));
    sigmaR = GetSigmaR(it->distance);
    //if (TMath::Abs(dR)<8.0) ndf++;
    chi2 += (dR*dR)/(sigmaR*sigmaR);
  }

  // For Fibre hits
  Double_t dy;
  for (TFibreHitInfoCollection::const_iterator 
	 it=fFibreHits->begin(); it!=fFibreHits->end(); it++){
    dy = it->y - (slope*it->x + offset);
    chi2 += dy*dy/(it->err*it->err);
  }
  return chi2;
}




TTrackFitter::TTrackFitter(const Char_t* minimizerType,
			   const Char_t* algoType) :
  fPositionResolution(NULL)
{
  // Setup the fitter
  fFitter.Config().SetMinimizer(minimizerType, algoType);

  
  // Configure the minimizer
  ROOT::Math::MinimizerOptions& minOpt = fFitter.Config().MinimizerOptions();
  minOpt.SetMaxFunctionCalls(100000); /// For Minuit/Minuit2
  minOpt.SetTolerance(0.025);
  minOpt.SetPrintLevel(0);
  
  // Clear parameters
  Reset();
}

TTrackFitter::~TTrackFitter()
{
}

void TTrackFitter::Reset()
{
  fFitResult = ROOT::Fit::FitResult();
  fWireHits.clear();
  fFibreHits.clear();
}

Bool_t TTrackFitter::Fit(const Double_t initSlope,  const Double_t initOffset,
			 const Double_t limitSlope, const Double_t limitOffset)
{
  if(0 == fWireHits.size() && 0 == fFibreHits.size()) return kFALSE;

  TrackMinimizer trkMin(&fWireHits, &fFibreHits, fPositionResolution);

  Double_t par[2] = {initSlope, initOffset};
  fFitter.SetFCN(2, trkMin, par, trkMin.GetDataSize(), true);
  fFitter.Config().ParSettings(0).SetValue(initSlope);
  fFitter.Config().ParSettings(0).SetLimits(initSlope-limitSlope,
					    initSlope+limitSlope);
  fFitter.Config().ParSettings(1).SetValue(initOffset);
  fFitter.Config().ParSettings(1).SetLimits(initOffset-limitOffset,
					    initOffset+limitOffset);
  fFitter.Config().SetUpdateAfterFit(true);

  bool res = fFitter.FitFCN();
  fFitResult = fFitter.Result();
  
  return res? kTRUE:kFALSE;
}