#include <string>
#include <iostream>
#include <iomanip>
#include <cmath>
#include <vector>

#include "JMath/JMath.hh"
#include "JDetector/JDetector.hh"
#include "JDetector/JDetectorToolkit.hh"
#include "JDetector/JModuleRouter.hh"
#include "JDetector/JTripod.hh"

#include "JFit/JSimplex.hh"
#include "JLang/JException.hh"
#include "JSupport/JMeta.hh"
#include "JTools/JRange.hh"

#include "Jeep/JContainer.hh"
#include "Jeep/JParser.hh"
#include "Jeep/JMessage.hh"


namespace {

  using namespace JPP;

  typedef JRange<int> JRange_t;


  /**
   * Model for detector alignment.
   *
   * The model allows for a global offset of the detector position, a rotation around the z-axis and a tilt of the (x,y) plane.
   */
  struct JModel_t :
    JMath<JModel_t>
  {
    /**
     * Default constructor.
     */
    JModel_t() :
      x  (0.0),
      y  (0.0),
      z  (0.0),
      phi(0.0),
      tx (0.0),
      ty (0.0)
    {}


    /**
     * Constructor.
     *
     * \param  x                       x
     * \param  y                       y
     * \param  z                       z
     * \param  phi                     phi
     * \param  tx                      Tx
     * \param  ty                      Ty
     */
    JModel_t(const double x,
	     const double y,
	     const double z,
	     const double phi,
	     const double tx,
	     const double ty) :
      x  (x),
      y  (y),
      z  (z),
      phi(phi),
      tx (tx),
      ty (ty)
    {}

    
    /**
     * Negate model.
     *
     * \return                          this model
     */
    JModel_t& negate()
    {
      x   = -x;
      y   = -y;
      z   = -z;
      phi = -phi;
      tx  = -tx;
      ty  = -ty;

      return *this;
    }


    /**
     * Add model.
     *
     * \param  model                    model
     * \return                          this model
     */
    JModel_t& add(const JModel_t& model)
    {
      x   += model.x;
      y   += model.y;
      z   += model.z;
      phi += model.phi;
      tx  += model.tx;
      ty  += model.ty;

      return *this;
    }


    /**
     * Subtract model.
     *
     * \param  model                    model
     * \return                          this model
     */
    JModel_t& sub(const JModel_t& model)
    {
      x   -= model.x;
      y   -= model.y;
      z   -= model.z;
      phi -= model.phi;
      tx  -= model.tx;
      ty  -= model.ty;

      return *this;
    }


    /**
     * Scale model.
     *
     * \param  factor                   multiplication factor
     * \return                          this model
     */
    JModel_t& mul(const double factor)
    {
      x   *= factor;
      y   *= factor;
      z   *= factor;
      phi *= factor;
      tx  *= factor;
      ty  *= factor;

      return *this;
    }


    /**
     * Scale model.
     *
     * \param  factor                   division factor
     * \return                          this model
     */
    JModel_t& div(const double factor)
    {
      x   /= factor;
      y   /= factor;
      z   /= factor;
      phi /= factor;
      tx  /= factor;
      ty  /= factor;

      return *this;
    }


    double x;                      //!< x-offset
    double y;                      //!< y-offset
    double z;                      //!< z-offset
    double phi;                    //!< rotation around z-axis [rad]
    double tx;                     //!< tilt angle [rad]
    double ty;                     //!< tilt angle [rad]
  };


  /**
   * Fit function.
   */
  class JFit_t :
    public JModuleRouter
  {
  public:
    /**
     * Constructor.
     *
     * \param  detector                detector (reference)
     * \param  sigma_m                 resolution [m]
     * \param  option                  option to keep strings vertical
     * \param  range                   range of floors used in fit
     */
    JFit_t(const JDetector& detector,
	   const double     sigma_m,
	   const bool       option,
	   const JRange_t   range) :
      JModuleRouter(detector),
      sigma (sigma_m),
      option(option),
      range (range)
    {
      using namespace JPP;

      center = getAverage(make_array(detector.begin(), detector.end(), &JModule::getPosition));

      center.setPosition(JPosition3D(center.getX(), center.getY(), 0.0));
    }


    /**
     * Get chi2.
     *
     * \param  model                   model
     * \param  module                  module
     * \return                         chi2
     */
    double operator()(const JModel_t& model, const JModule& module) const
    {
      if (this->hasModule(module.getID()) && range(module.getFloor())) {

	const JPosition3D p1 = this->getModule(module.getID()).getPosition();
	const JPosition3D p2 = this->getPosition(model, module.getPosition());

	return p2.getDistanceSquared(p1) / (sigma*sigma);
      }

      return 0.0;
    }


    /**
     * Get position according model.
     *
     * \param  model                   model
     * \param  pos                     position
     * \return                         position
     */
    JPosition3D getPosition(const JModel_t& model, const JPosition3D& pos) const
    {
      using namespace std;
      using namespace JPP;

      JPosition3D p1 = pos;

      p1 -= center;

      p1.rotate(JRotation3Z(model.phi));

      {
	const double tx  =  model.tx;
	const double tz  =  sqrt((1.0 + tx) * (1.0 - tx));

	if (!option)
	  p1  = JPosition3D(tz * p1.getX() + tx * p1.getZ(), p1.getY(), tz * p1.getZ() - tx * p1.getX());	
	else
	  p1  = JPosition3D(tz * p1.getX(),                  p1.getY(), p1.getZ()      - tx * p1.getX());
      }
      {
	const double ty  =  model.ty;
	const double tz  =  sqrt((1.0 + ty) * (1.0 - ty));

	if (!option)
	  p1  = JPosition3D(p1.getX(), tz * p1.getY() + ty * p1.getZ(), tz * p1.getZ() - ty * p1.getY());
	else
	  p1  = JPosition3D(p1.getX(), tz * p1.getY(),                  p1.getZ()      - ty * p1.getY());
      }
      
      p1 += JPosition3D(model.x, model.y, model.z);
      p1 += center;

      return p1;
    }
    
    double      sigma;
    bool        option;
    JRange_t    range;
    JPosition3D center;
  };
}


/**
 * \file
 *
 * Auxiliary program to align two detectors.\n
 * \author mjongen
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;

  string             detectorFile_a;
  string             detectorFile_b;
  bool               overwriteDetector;
  string             tripodFile;
  double             sigma_m;
  bool               option;
  JRange_t           range;
  int                debug;

  try {

    JParser<> zap("Auxiliary program to align two detectors.");

    zap['a'] = make_field(detectorFile_a,      "detector - subject to alignment (option -A)");
    zap['b'] = make_field(detectorFile_b,      "detector - reference for alignment");
    zap['A'] = make_field(overwriteDetector,   "overwrite detector file provided through '-a' with modified positions.");
    zap['T'] = make_field(tripodFile,          "tripods")                               = "";
    zap['s'] = make_field(sigma_m)               = 0.2;
    zap['O'] = make_field(option,              "keep strings vertical");
    zap['r'] = make_field(range,               "range of floors used in fit")           = JRange_t();
    zap['d'] = make_field(debug)                 = 2;

    zap(argc, argv);
  }
  catch(const exception &error) {
    FATAL(error.what() << endl);
  }


  if (overwriteDetector) {
    if (tripodFile == "") {
      FATAL("No tripod file.");
    }
  }
  
  
  JDetector detector_a;
  JDetector detector_b;

  try {
    load(detectorFile_a, detector_a);
  }
  catch(const JException& error) {
    FATAL(error);
  }

  try {
    load(detectorFile_b, detector_b);
  }
  catch(const JException& error) {
    FATAL(error);
  }


  const JFit_t fit(detector_b, sigma_m, option, range);

  vector<JModule> data;

  for (JDetector::const_iterator module = detector_a.begin(); module != detector_a.end(); ++module) {
    if (fit.hasModule(module->getID())) {
      data.push_back(*module);
    }
  }


  JSimplex<JModel_t> simplex;
  
  JSimplex<JModel_t>::debug              = debug;
  JSimplex<JModel_t>::MAXIMUM_ITERATIONS = 100000;

  simplex.value = JModel_t(0.0, 0.0, 0.0, 0.0, 0.0, 0.0);

  simplex.step.resize(6);

  simplex.step[0] = JModel_t(0.01, 0.00, 0.00, 0.0,    0.0, 0.0);
  simplex.step[1] = JModel_t(0.00, 0.01, 0.00, 0.0,    0.0, 0.0);
  simplex.step[2] = JModel_t(0.00, 0.00, 0.01, 0.0,    0.0, 0.0);
  simplex.step[3] = JModel_t(0.00, 0.00, 0.00, 5.0e-4, 0.0, 0.0);
  simplex.step[4] = JModel_t(0.00, 0.00, 0.00, 0.0, 1.0e-4, 0.0);
  simplex.step[5] = JModel_t(0.00, 0.00, 0.00, 0.0, 0.0, 1.0e-4);

  const double chi2 = simplex(fit, data.begin(), data.end());

  cout << "Number of iterations " << simplex.numberOfIterations << endl;
  cout << "chi2/NDF " << FIXED(7,3) << chi2 << '/' << (detector_a.size() - simplex.step.size()) << endl;

  cout << "model:" << endl;

  cout << "x    " << FIXED(7,3) << simplex.value.x   << endl;
  cout << "y    " << FIXED(7,3) << simplex.value.y   << endl;
  cout << "z    " << FIXED(7,3) << simplex.value.z   << endl;
  cout << "phi  " << FIXED(9,5) << simplex.value.phi << endl;
  cout << "Tx   " << FIXED(9,5) << simplex.value.tx  << endl;
  cout << "Ty   " << FIXED(9,5) << simplex.value.ty  << endl;


  if (overwriteDetector) {

    NOTICE("Store alignment data on files " << detectorFile_a << " and " << tripodFile << endl);

    detector_a.comment.add(JMeta(argc,argv));

    for (JDetector::iterator module = detector_a.begin(); module != detector_a.end(); ++module) {

      const JPosition3D p1 = module->getPosition();
      const JPosition3D p2 = fit.getPosition(simplex.value, p1);

      module->add(p2 - p1);
    }

    try {
      store(detectorFile_a, detector_a);
    }
    catch(const JException& error) {
      FATAL(error);
    }

    JContainer< vector<JTripod> > tripods;

    tripods.load(tripodFile.c_str());

    tripods.comment.add(JMeta(argc,argv));
    
    for (vector<JTripod>::iterator tripod = tripods.begin(); tripod != tripods.end(); ++tripod) {
      
      const JPosition3D p1 = tripod->getUTMPosition() - detector_a.getUTMPosition();
      const JPosition3D p2 = fit.getPosition(simplex.value, p1);

      tripod->add(p2 - p1);
    }
    
    tripods.store(tripodFile.c_str());
  }
}