/* This file is part of MAUS: http://micewww.pp.rl.ac.uk:8080/projects/maus
 *
 * MAUS is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * MAUS is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MAUS.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#ifndef _SRC_MAP_MAPCPPTrackerTOFCombinedFit_H_
#define _SRC_MAP_MAPCPPTrackerTOFCombinedFit_H_

// #define KALMAN_TEST

// C headers
#include <assert.h>
#include <json/json.h>
#include <CLHEP/Units/PhysicalConstants.h>

// C++ headers
#include <string>
#include <sstream>
#include <vector>
#include <map>

// Other headers
#include "Utils/Exception.hh"
#include "Utils/Squeak.hh"
#include "Config/MiceModule.hh"
#include "src/common_cpp/Utils/CppErrorHandler.hh"
#include "src/common_cpp/Utils/JsonWrapper.hh"
#include "src/common_cpp/Utils/Globals.hh"
#include "src/common_cpp/Globals/GlobalsManager.hh"
#include "src/common_cpp/JsonCppProcessors/SpillProcessor.hh"
#include "src/common_cpp/DataStructure/ReconEvent.hh"

#include "src/common_cpp/DataStructure/SciFiEvent.hh"
#include "src/common_cpp/DataStructure/TOFEvent.hh"
#include "src/common_cpp/DataStructure/Spill.hh"
#include "src/common_cpp/DataStructure/SciFiSeed.hh"
#include "src/common_cpp/Recon/SciFi/PatternRecognition.hh"
#include "src/common_cpp/Recon/Kalman/KalmanTrackFit.hh"
#include "src/common_cpp/Recon/Kalman/KalmanTrack.hh"
#include "src/common_cpp/Recon/Kalman/MAUSSciFiPropagators.hh"
#include "src/common_cpp/Recon/Kalman/MAUSSciFiMeasurements.hh"
#include "src/common_cpp/API/MapBase.hh"

namespace MAUS {

struct SciFiPlaneGeometry;

/** @class MapCppTrackerTOFCombinedFit
 *  Perform the final tracker track fit using the Kalman filter
 *
 */
class MapCppTrackerTOFCombinedFit : public MapBase<Data> {

 public:
  /** @brief Constructor - initialises pointers to NULL */
  MapCppTrackerTOFCombinedFit();

  /** @brief  Constructor - deletes any allocated memory */
  ~MapCppTrackerTOFCombinedFit();

 private:
  /** @brief Sets up the worker
   *
   *  @param 	argJsonConfigDocument 	a JSON document with the configuration.
   */
  void _birth(const std::string& argJsonConfigDocument);

  /** @brief Shutdowns the worker
   *
   *  This takes no arguments and deletes the Kalman fitter object
   */
  void _death();

  /** @brief Process MAUS data object
   *
   *  Receive a data object with digits (either MC or real) and then call the higher level
   *  reconstruction algorithms
   *
   * @param	data			Pointer to the data object
   */
  void _process(Data* data) const;

  /** @brief Calculates the longitudinal momentum from the TOF measurement, in the muon assumption
   *
   *  Takes the time-of-flight and reconstructs the longitudinal momentum
   *
   *  @param 	tof_event		Event that contains the TOF information
   *  @param	tof_p			TOF momentum
   *  @param 	tof_dp 			Uncertainty on the TOF momentum
   */
  bool calculate_tof_momentum(TOFEvent* tof_event, double& tof_p, double& tof_dp) const;

  /** @brief Propagates the momentum to from A to B, in the muon assumption
   *
   *  Takes the momentum and calculates the energy loss step-by-step
   *  from the starting position to the target position
   *
   *  @param 	p 			Initial momentum
   *  @param 	dp 			Uncertainty on the propagated momentum
   *  @param 	start 			Starting z position
   *  @param 	end 			Target z position
   */
  bool propagate_momentum(double& p, double& dp, double start, double end) const;

  /** @brief Checks whether the track needs refitting or not
   *
   *  Track fit takes the seed built upon the TOF information to produce a new track
   *
   *  @param 	seed 			Current tracker Kalman track seed
   */
  bool needs_refitting(MAUS::SciFiSeed* seed) const;

  /** @brief Updates the Kalman seed for failed helical tracks
   *
   *  @param 	seed 			Current tracker Kalman track seed
   *  @param 	tof_p 			TOF momentum
   */
  void update_seed(MAUS::SciFiSeed* seed, double tof_p) const;

  /** @brief Updates the Kalman seed for straight helical tracks
   *
   *  @param 	seed 			Current tracker Kalman track seed
   *  @param 	tof_p 			TOF momentum
   */
  void update_straight_seed(MAUS::SciFiSeed* seed, double tof_p) const;

  /** @brief Performs the new track fit
   *
   *  Track fit takes the seed built upon the TOF information to produce a new track
   *
   *  @param 	seed 			New tracker Kalman track seed
   */
  SciFiTrack* track_fit_helix(MAUS::SciFiSeed* seed) const;

  /** @brief Updates the Kalman track by taking the weighted average with the TOF information
   *
   *  @param 	track 			Kalman track
   *  @param 	tof_p 			TOF momentum
   *  @param 	tof_dp 			Uncertainty on the TOF momentum
   */
  bool update_track(MAUS::SciFiTrack* track, double tof_p, double tof_dp) const;

 private:

  /// This will contain the configuration
  Json::Value _configJSON;

  // Various options and parameters
  double _radius_cut;
  double _tof_resolution;
  bool _has_tof;

  // Position of reference planes
  double _tof0_z;
  double _tof1_z;
  double _tku_z;
  double _tkd_z;

  /// Kalman Track Fitter Object
  Kalman::TrackFit* _helical_track_fitter;

  SciFiHelicalMeasurements* _helical_measurement;

  int SciFiRunRecon;
}; // Don't forget this trailing colon!!!!

} // ~namespace MAUS

#endif