// DTrackLSFitter.cc
//

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

#include "DTrackLSFitter.h"
#include "FDC/DFDCHit.h"
#include "FDC/DFDCPseudo_factory.h"
#include "FDC/DFDCSegment_factory.h"
#include "CDC/DCDCTrackHit.h"
#include "TRACKING/DMCThrown.h"
#include "TRACKING/DMCTrackHit.h"
#include "DANA/DApplication.h"

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

#define TARGET_POSITION 65.0
#define THETA_FORWARD_CUT 1.178097
#define THETA_BACKWARD_CUT 1.963495
#define THREEPIOVER4 2.356194

//------------------------------------------------------------------
// DTrackLSFitter 
//------------------------------------------------------------------
DTrackLSFitter::DTrackLSFitter():debug_level(1), ppEnd(5)
{
  configFile = new ifstream("config.txt");
  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;
  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;
  vector<const DFDCHit*>fdchits;
  //eventLoop->Get(fdchits);  

  vector<const DFDCPseudo*>pseudopoints;
  eventLoop->Get(pseudopoints);  

  vector<const DCDCTrackHit*>trackhits;
  eventLoop->Get(trackhits);  
  
  vector<const DMCTrackHit*>truepoints;
  eventLoop->Get(truepoints);

  vector<const DMCThrown*>thrown;
  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) return NOERROR;

  fdcDataFile = new ofstream("fdcData.txt");
  cdcDataFile = new ofstream("cdcData.txt");
  trajFile = new ofstream("fdcTraj.txt");

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

  HepVector B(3);
  B(1) = 0.0; B(2) = 0.0; B(3) = 4.0;

  double xpInitial, zInitial = TARGET_POSITION, thetaInitial, phiInitial,
    pinvInitial = 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) {
    thetaInitial = 0.35; // about 20 degrees
    // overwrite with very rough theta guess
    float thetaInitialTrue = thrown[0]->momentum().Theta();
    if (0.0 <= thetaInitialTrue && thetaInitialTrue < THETA_FORWARD_CUT) {
      thetaInitial = PIOVER4;
    } else if (thetaInitialTrue <= THETA_BACKWARD_CUT) {
      thetaInitial = PIOVER2;
    } else {
      thetaInitial = THREEPIOVER4;
    }
    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); 
    /*
    // overwrite with true values
    pinvInitial = pinvTrue;
    thetaInitial = thetaTrue;
    phiInitial = phiTrue;
    xpInitial = xp0True;
    */
  } else {
    int code = 32;
    throw code;
  }

  vector<HepVector*>* trajPtr;
  HepVector trajPoint(3);
  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) = pinvInitial;
    helixTraj.swim(ppStart); // swim with initial parameters
    trajPtr = helixTraj.getTrajectory();
    for (unsigned int i = 0; i < trajPtr->size(); i++) {
      trajPoint = *((*trajPtr)[i]);
      *trajFile << "0 " << i + 1 << " " << trajPoint(1) << " " << trajPoint(2)
		<< " " << trajPoint(3) << endl;
    }
    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
    trajPtr = helixTraj.getTrajectory();
    for (unsigned int i = 0; i < trajPtr->size(); i++) {
      trajPoint = *((*trajPtr)[i]);
      *trajFile << "1 " << i + 1 << " " << trajPoint(1) << " " << trajPoint(2)
		<< " " << trajPoint(3) << endl;
    }
    *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();
    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();
  trajFile->close();
  delete fdcDataFile;
  delete cdcDataFile;
  delete trajFile;

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

  return NOERROR;
 
}

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

double DTrackLSFitter::getChiSquared() {
  return fitPtr->getChi2();
}

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

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