// DTrackLSFitter.cc
//

#include <iostream>
#include <vector>
#include <string>
#include <CLHEP/Matrix/Matrix.h>
#include <CLHEP/Matrix/Vector.h>

#include "FDC/DFDCHit.h"
#include "FDC/DFDCPseudo_factory.h"
#include "TRACKING/DTrack.h"
#include "TRACKING/DMCTrackHit.h"
#include "DANA/DApplication.h"

#include "MyTrajectory.h"
#include "MyTrajectoryBfield.h"
#include "chisqMin.h"
#include "combinedResidFunc.h"
#include "globalGslFuncs.h"
#include "DTrackLSFitter.h"

#define TARGET_POSITION 65.0
#define THETA_FORWARD_CUT 1.178097
#define THETA_BACKWARD_CUT 1.963495
#define THREEPIOVER4 2.356194
#define MAX_TRAJS 1000
#define MAX_POINTS 10000

//------------------------------------------------------------------
// DTrackLSFitter 
//------------------------------------------------------------------
DTrackLSFitter::DTrackLSFitter():debug_level(1), ppEnd(5)
{
  configFile = new ifstream("config.txt");
  char hddmFileName[128] = "fitter.hddm";
  ios = init_fitter_HDDM(hddmFileName);
  int debug_level_in;
  *configFile >> debug_level_in;
  cout << "debug_level_in = " << debug_level_in << endl;
  if (debug_level_in >= -10 && debug_level_in <= 10) {
    debug_level = debug_level_in;
  }
}

//------------------------------------------------------------------
// ~DTrackLSFitter 
//------------------------------------------------------------------
DTrackLSFitter::~DTrackLSFitter()
{
  cout << "DTrackLSFitter destructor called\n";
}

//------------------------------------------------------------------
// init 
//------------------------------------------------------------------
jerror_t DTrackLSFitter::init(void)
{
  signatureFile = new ofstream("fitter_signature.txt");
  return NOERROR;
}

//------------------------------------------------------------------
// fini
//------------------------------------------------------------------
jerror_t DTrackLSFitter::fini(void)
{
  signatureFile->close();
  delete signatureFile;
  close_fitter_HDDM(ios);
  return NOERROR;
}

//------------------------------------------------------------------
// brun 
//------------------------------------------------------------------
jerror_t DTrackLSFitter::brun(JEventLoop *eventLoop, int runnumber)
{
  DApplication* dapp=dynamic_cast<DApplication*>(eventLoop->GetJApplication());
  bfield = dapp->GetBfield();
  lorentz_def = dapp->GetLorentzDeflections();
  JFactory_base *base = eventLoop->GetFactory("DFDCSegment");
  segment_factory = dynamic_cast<DFDCSegment_factory*>(base);
  return NOERROR;
}

//------------------------------------------------------------------
// erun 
//------------------------------------------------------------------
jerror_t DTrackLSFitter::erun()
{
 return NOERROR;
}

//------------------------------------------------------------------
// evnt 
//------------------------------------------------------------------

jerror_t DTrackLSFitter::evnt(JEventLoop *eventLoop, int eventnumber)
{

  if (debug_level > 1) cout << "----------- New Event "
			     << eventnumber << " -------------" << endl;

  fitPtr = NULL;
  fitter_HDDM_t* hddm;
  hddm = make_fitter_HDDM();
  fitter_Events_t* events;
  events = make_fitter_Events(1);
  hddm->events = events;
  events->mult = 1;
  events->in[0].run = -1;
  events->in[0].number = eventnumber;
  fitter_Trajectorys_t* trajsHddm = make_fitter_Trajectorys(MAX_TRAJS);
  events->in[0].trajectorys = trajsHddm;

  vector<const DTrack*>tracks;
  eventLoop->Get(tracks, "THROWN");
  if (debug_level > 1) cout << "number of tracks = " << tracks.size() << endl;

  pseudopoints.clear();
  //eventLoop->Get(pseudopoints);  
  tracks[0]->Get(pseudopoints);  

  trackhits.clear();
  //eventLoop->Get(trackhits);  
  tracks[0]->Get(trackhits);  
  
  thrown.clear();
  eventLoop->Get(thrown);

  size_fdc =  pseudopoints.size();
  if (debug_level > 1) cout << "number of pseudopoints = " << size_fdc << endl;

  size_cdc = trackhits.size();
  if (debug_level > 1) cout << "number of trackhits = " << size_cdc << endl;

  if (size_fdc + size_cdc < 6) {
    flush_fitter_HDDM(hddm, ios);
    return NOERROR;
  }

  fdcDataFile = new ofstream("fdcData.txt");
  cdcDataFile = new ofstream("cdcData.txt");
  cdcRawDataFile = new ofstream("cdcRawData.txt");
  trajFile = new ofstream("traj.txt");

  for (int i = 0; i < size_fdc; i++) {
    *fdcDataFile << i + 1 << " " << (pseudopoints[i])->x << " "
	      << (pseudopoints[i])->y << " "
	      << (pseudopoints[i])->wire->origin(2) << endl;
  }

  for (int i = 0; i < size_cdc; i++) {
    const DCDCTrackHit *trkhit = trackhits[i];
    const DCDCWire *wire = trkhit->wire;
    double x = wire->origin.X();
    double y = wire->origin.Y();
    double phi_wire = atan2(y, x);
    double phi = phi_wire + PIOVER2;
    double theta = wire->stereo;
    *cdcRawDataFile << i + 1 << " "
		    << x << " "
		    << y << " "
		    << TARGET_POSITION << " "
		    << theta << " " << phi << " " << trkhit->dist << endl;
  }

  setFitterStartParams();

  MyTrajectoryBfield helixTraj(bfield, debug_level);
 
  // instantiate the residual function class
  combinedResidFunc prf(&pseudopoints, &trackhits, &helixTraj, lorentz_def,
			debug_level);

  residFuncPtr = &prf; // set the global residual point to refer to this class
  fitPtr = new chisqMin(&prf, debug_level);
  HepVector ppStart(5);

  try {

    ppStart(1) = xpInitial;
    ppStart(2) = zInitial;
    ppStart(3) = thetaInitial;
    ppStart(4) = phiInitial;
    ppStart(5) = ptinvInitial;
    helixTraj.swim(ppStart); // swim with initial parameters
    int tag_initial = 0;
    helixTraj.dump_ascii(trajFile, tag_initial);
    writeTrajectoryHddm(helixTraj, tag_initial, trajsHddm);
    helixTraj.clear();

    for (double df = 0.0; df < 1.0001; df += 1.0) {
      prf.setInnerResidFrac(df);
      fitPtr->setStartParams(ppStart);
      fitPtr->fit();
      fitPtr->getParams(ppEnd);
      // cout << df << ' ' << fitPtr->getChi2() << endl;
      ppStart = ppEnd;
    }

    helixTraj.swim(ppEnd); // swim with final parameters
    int tag_final = 1;
    helixTraj.dump_ascii(trajFile, tag_final);
    writeTrajectoryHddm(helixTraj, tag_final, trajsHddm);
    *signatureFile << eventnumber << ' '
		   << size_fdc << ' '
		   << size_cdc << ' '
		   << thetaInitial << ' '
		   << phiInitial << ' '
		   << ppEnd(1) << ' ' 
		   << ppEnd(2) << ' '
		   << ppEnd(3) << ' '
		   << ppEnd(4) << ' '
		   << ppEnd(5) << ' '
		   << fitPtr->getChi2() << endl;
    // get details of hits on track
    helixTraj.clear();
    prf.setStoreDetails(true);
    // get some details on the fit
    int dummy_data;
    HepVector dummy_vector(size_fdc + size_cdc);
    prf.resid(&ppEnd, &dummy_data, &dummy_vector);
    vector<CDCHitDetails*> *CDCDetailsPtr;
    CDCDetailsPtr = prf.getCDCDetails();
    writeCDCDetailsHddm(CDCDetailsPtr, trajsHddm);
    for (unsigned int i = 0; i < CDCDetailsPtr->size(); i++) {
      const DCDCTrackHit *trkhit = trackhits[i];
      const DCDCWire *wire = trkhit->wire;
      double x = wire->origin.X();
      double y = wire->origin.Y();
      double phi_wire = atan2(y, x);
      double phi = phi_wire + PIOVER2;
      double theta = wire->stereo;
      *cdcDataFile << i + 1 << " "
		   << ((*CDCDetailsPtr)[i])->posWire(1) << " "
		   << ((*CDCDetailsPtr)[i])->posWire(2) << " "
		   << ((*CDCDetailsPtr)[i])->posWire(3) << " "
		   << theta << " " << phi << " " << trkhit->dist << endl;
    }

    // clean up
    prf.clearDetails();
    prf.setStoreDetails(false);
    helixTraj.clear();
  } catch (int code) {

    cout << "==========fit error = " << code << "===========" << endl;
    cout << "= at event " << eventnumber << endl;
    if (debug_level >= 3) {
      cout << "= trajectory" << endl;
      helixTraj.print();
    }

  }

  fdcDataFile->close();
  cdcDataFile->close();
  cdcRawDataFile->close();
  trajFile->close();
  delete fdcDataFile;
  delete cdcDataFile;
  delete cdcRawDataFile;
  delete trajFile;

  if (debug_level >= 3) {
    cout << "FDCHitDetails: " << FDCHitDetails::getInstanceCount() << endl;
    cout << "CDCHitDetails: " << CDCHitDetails::getInstanceCount() << endl;
  }

  flush_fitter_HDDM(hddm, ios);

  return NOERROR;
 
}

HepVector DTrackLSFitter::getParams() {
  return ppEnd;
}

double DTrackLSFitter::getChiSquared() {
  double chisq = -1.0;
  if (fitPtr) chisq = fitPtr->getChi2();
  return chisq;
}

int DTrackLSFitter::getSizeFDC() {
  return size_fdc;
}

int DTrackLSFitter::getSizeCDC() {
  return size_cdc;
}

double radialDist2(const DCDCTrackHit *trkhit) {
  double x, y, r2;
  const DCDCWire *wire;
  wire = trkhit->wire;
  x = wire->origin.X();
  y = wire->origin.Y();
  r2 = x*x + y*y;
  return r2;
}

class compareTrackHits {
public:
  int operator() (const DCDCTrackHit* const &hit1, const DCDCTrackHit* const &hit2) const
  {return radialDist2(hit1) < radialDist2(hit2);}
};

void DTrackLSFitter::setFitterStartParams() {
  zInitial = TARGET_POSITION;
  ptinvInitial = 0.0;
  if (size_fdc > 0) {
    xpInitial = 0.0;
    double xfirst, yfirst, zfirst;
    xfirst = (pseudopoints[0])->x;
    yfirst = (pseudopoints[0])->y;
    zfirst = (pseudopoints[0])->wire->origin(2);
    double delta_z = zfirst - zInitial;
    double rPerp = sqrt(xfirst*xfirst + yfirst*yfirst);
    thetaInitial = atan2(rPerp, delta_z);
    phiInitial = atan2(yfirst, xfirst);
  } else if (size_cdc > 0) {

    float thetaInitialTrue = thrown[0]->momentum().Theta();
    if (thetaInitialTrue < PIOVER2) {
      thetaInitial = PIOVER4;
    } else {
      thetaInitial = THREEPIOVER4;
    }

    stable_sort(trackhits.begin(), trackhits.end(), compareTrackHits());

    const DCDCTrackHit *trkhit_0 = trackhits[0];
    const DCDCWire *wire_0 = trkhit_0->wire;
    const DCDCTrackHit *trkhit_n = trackhits[size_cdc - 1];
    const DCDCWire *wire_n = trkhit_n->wire;
    double deltaX = wire_n->origin.X() - wire_0->origin.X();
    double deltaY = wire_n->origin.Y() - wire_0->origin.Y();
    phiInitial = atan2(deltaY, deltaX);
    double alpha = phiInitial - PIOVER2;
    xpInitial = wire_0->origin.X()*cos(alpha) + wire_0->origin.Y()*sin(alpha); 
  } else {
    int code = 32;
    throw code;
  }
}

void DTrackLSFitter::writeTrajectoryHddm
(MyTrajectory &trajectory, int tag, fitter_Trajectorys_t *trajsHddm) {
  fitter_Trajectory_t* trajHddm = &(trajsHddm->in[trajsHddm->mult++]);
  trajHddm->index = tag;
  fitter_Points_t* pointsHddm = make_fitter_Points(MAX_POINTS);
  trajHddm->points = pointsHddm;
  vector<HepVector*> *traj;
  traj = trajectory.getTrajectory();
  HepVector trajPoint(3);
  fitter_Point_t* pointHddm;
  for (unsigned int i = 0; i < traj->size(); i++) {
    trajPoint = *((*traj)[i]);
    pointHddm = &(pointsHddm->in[pointsHddm->mult++]);
    pointHddm->x = trajPoint(1);
    pointHddm->y = trajPoint(2);
    pointHddm->z = trajPoint(3);
  }
}

void DTrackLSFitter::writeCDCDetailsHddm(vector<CDCHitDetails*> *CDCDetailsPtr, fitter_Trajectorys_t *trajsHddm) {
  int nCDC = CDCDetailsPtr->size();
  fitter_Trajectory_t* trajHddm = &(trajsHddm->in[trajsHddm->mult - 1]);
  fitter_Residuals_t* residsHddm = make_fitter_Residuals(nCDC);
  trajHddm->residuals = residsHddm;
  double dist, doca, resid;
  for (int i = 0; i < nCDC; i++) {
    dist = ((*CDCDetailsPtr)[i])->dist;
    doca = ((*CDCDetailsPtr)[i])->doca;
    resid = doca - dist;
    residsHddm->in[i].detector_index = i;
    residsHddm->in[i].value = resid;
  }
  residsHddm->mult = nCDC;
}