#include <CLHEP/Matrix/Matrix.h>
#include <CLHEP/Matrix/Vector.h>

using namespace std;

#include "FDC/DFDCPseudo_factory.h"
#include "MyTrajectory.h"
#include "residFunc.h"
#include "combinedResidFunc.h"
#include "HDGEOMETRY/DLorentzDeflections.h"

// constructor, takes pointer to vector of pseudopoints and a pointer
// to a trajectory
combinedResidFunc::combinedResidFunc(vector<const DFDCPseudo*> *pseudopoints,
				     vector<const DCDCTrackHit*> *trackHits,
				     MyTrajectory *trajectory, const DLorentzDeflections *lorentz_def_in, int level) : 
  n_fdc(pseudopoints->size()), n_cdc(trackHits->size()), ppPtr(pseudopoints),
  trkhitPtr(trackHits), trajPtr(trajectory), delta(trajPtr->getDelta()),
  debug_level(level), lorentz_def(lorentz_def_in)
{}

void combinedResidFunc::resid(const HepVector *x, void *data, HepVector *f){
// input parameters
//   x: pointer to a vector of fit parameters
//   data: ???
// output parameters
//   f: pointer to vector of residuals
  double doca;
  if (debug_level > 2) {
    cout << "combinedResidFunc::resid: resid called\n";
    cout << "                          params: " << *x;
  }
  // do a swim with the input parameters
  trajPtr->swim(*x);
  // populate f vector with residuals
  HepVector point(3);
  for (unsigned int i = 0; i < n_fdc; i++) {
    point = pseudo2HepVector(*((*ppPtr)[i]));
    (*f)(i + 1) = trajPtr->doca(point);
  }
  DLine line;
  double dist;
  for (unsigned int j = 0; j < n_cdc; j++) {
    line = trackhit2line(*((*trkhitPtr)[j]));
    doca = trajPtr->doca(line);
    dist = (*trkhitPtr)[j]->dist;
    if (debug_level > 2) cout << " dist = " << dist << " doca = " << doca << " resid(prelim) = " << dist - doca << endl;
    if (isnan(dist)) {
      (*f)(n_fdc + j + 1) = 0.0;
    } else {
      (*f)(n_fdc + j + 1) = dist - doca;
    }
  }
  if (debug_level > 2) cout << "combinedResidFunc::resid: resids:" << *f;
  trajPtr->clear();
};

void combinedResidFunc::deriv(const HepVector *x, void *data, HepMatrix *J){
  if (debug_level > 2) {
    cout << "combinedResidFunc::deriv: deriv called\n";
    cout << "                          params: " << *x << endl;
  }
  // save base parameters
  HepVector xBase = *x;
  // do central swim
  trajPtr->swim(xBase);
  // store pseudo points as three vectors
  vector<HepVector *>pPoints;
  HepVector *thisPointPtr;
  for (unsigned int i = 0; i < n_fdc; i++) {
    thisPointPtr = new HepVector(3);
    *thisPointPtr = pseudo2HepVector(*((*ppPtr)[i]));
    pPoints.push_back(thisPointPtr);
  }
  // store track hits as lines
  vector<DLine *> linePtrs;
  DLine *thisLinePtr;
  for (unsigned int j = 0; j < n_cdc; j++) {
    thisLinePtr = new DLine();
    *thisLinePtr = trackhit2line(*((*trkhitPtr)[j])); // copy constructor needed??????
    linePtrs.push_back(thisLinePtr);
  }
  // store base residuals
  HepVector residBase(n_fdc + n_cdc);
  // base resids for FDC
  for (unsigned int i = 0; i < n_fdc; i++) {
    residBase(i + 1) = trajPtr->doca(*(pPoints[i]));
  }
  // base resids for CDC
  double docaThis, distThis;
  for (unsigned int j = 0; j < n_cdc; j++) {
    distThis = (*trkhitPtr)[j]->dist;
    docaThis = trajPtr->doca(*(linePtrs[j]));
    residBase(n_fdc + j + 1) = distThis - docaThis;
  }
  if (debug_level > 2) cout << "base resids:" << residBase;
  trajPtr->clear();
  // calculate Jacobian
  unsigned int p = trajPtr->getNumberOfParams();
  HepVector xThis(p);
  int iHep, jHep; // index for HepVector () notation
  for (unsigned int i = 0; i < p; i++) {
    iHep = i + 1;
    xThis = xBase; // set params back to base
    xThis(iHep) = xBase(iHep) + delta[i];
    if (debug_level > 2) cout << "perturbed params: iHep = " << iHep << ", values:" << xThis << endl;
    // do the perturbed swim
    trajPtr->swim(xThis);
    // calculate derivatives for FDC points
    for (unsigned int j = 0; j < n_fdc; j++) {
      jHep = j + 1;
      docaThis = trajPtr->doca(*(pPoints[j]));
      if (debug_level > 2) cout << "resid " << j << " = " << docaThis << endl;
      (*J)(jHep, iHep) = (docaThis - residBase(jHep))/delta[i];
    }
    // calculate derivatives for CDC points
    for (unsigned int j = 0; j < n_cdc; j++) {
      jHep = n_fdc + j + 1;
      distThis = (*trkhitPtr)[j]->dist;
      docaThis = trajPtr->doca(*(linePtrs[j]));
      if (debug_level > 2) cout << j << " dist = " << distThis << " doca = " << docaThis << " resid  = " << distThis - docaThis << endl;
      if (isnan(distThis)) {
	(*J)(jHep, iHep) = 0;
      } else {
	(*J)(jHep, iHep) = (distThis - docaThis - residBase(jHep))/delta[i];
      }
    }
    trajPtr->clear();
  }
  if (debug_level >= 3) {
    for (unsigned int i = 0; i < p; i++) {
      iHep = i + 1;
      cout << iHep;
      for (unsigned int j = 0; j < n_fdc + n_cdc; j++) {
	jHep = j + 1;
	cout << ' ' << (*J)(jHep, iHep);
      }
      cout << endl;
    }
  }
  for (unsigned int i = 0; i < n_fdc; i++) {
    delete pPoints[i];
  }
  for (unsigned int j = 0; j < n_cdc; j++) {
    delete linePtrs[j];
  }
};

void combinedResidFunc::residAndDeriv(const HepVector *x, void *data, HepVector *f,
		   HepMatrix *J){};

HepVector combinedResidFunc::pseudo2HepVector(const DFDCPseudo &ppoint) {
  double x;
  double y;
  double z = ppoint.wire->origin(2);
  trajPtr->get_xy(z, x, y); // on trajectory
  bool ispos = get_correction_sign(ppoint, x, y, z);
  double delta_x, delta_y;
  get_correction_value(ppoint, x, y, z, delta_x, delta_y);
  HepVector point(3);
  if (ispos) {
    point(1) = ppoint.x + delta_x;
    point(2) = ppoint.y + delta_y;
  } else {
    point(1) = ppoint.x - delta_x;
    point(2) = ppoint.y - delta_y;
  }
  point(3) = z;
  return point;
}

bool combinedResidFunc::get_correction_sign(const DFDCPseudo &ppoint, double x, double y, double z) {
  double cosangle = ppoint.wire->udir(1);
  double sinangle = ppoint.wire->udir(0);
  double delta_w = x*cosangle - y*sinangle - ppoint.w;
  bool ispos;
  if (delta_w > 0.0) {
    ispos = true;
  } else {
    ispos = false;
  }
  return ispos;
}

void combinedResidFunc::get_correction_value(const DFDCPseudo &ppoint, double x, double y, double z, double &delta_x, double &delta_y) {
  double driftDist = ppoint.time*DRIFT_VELOCITY;
  double cosangle = ppoint.wire->udir(1);
  double sinangle= ppoint.wire->udir(0);
  double alpha = 0.0; // for now
  double lorentzShift = lorentz_def->GetLorentzCorrection(x, y, z, alpha, driftDist);
  double ds = -lorentzShift;
  double dw = driftDist;
  delta_x = dw*cosangle + ds*sinangle;
  delta_y = -dw*sinangle + ds*cosangle;
  return;
}

DLine combinedResidFunc::trackhit2line(const DCDCTrackHit &trkhit) {
  const DCDCWire *wire = trkhit.wire;
  double x = wire->origin.X();
  double y = wire->origin.Y();
  double z = wire->origin.Z();
  double phi_wire = atan2(y, x);
  double phi = phi_wire + PIOVER2;
  double theta = wire->stereo;
  DLine line(x, y, z, theta, phi);
  return line;
}