/* 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/>.
 *
 */

#include "src/map/MapCppTrackerTOFReFit/MapCppTrackerTOFReFit.hh"

#include <sstream>

#include "src/common_cpp/API/PyWrapMapBase.hh"
#include "src/common_cpp/Recon/Kalman/MAUSTrackWrapper.hh"
#include "src/common_cpp/Recon/Global/MaterialModelAxialLookup.hh"
#include "TMatrixD.h"

#include "Geant4/G4Material.hh"

namespace MAUS {

PyMODINIT_FUNC init_MapCppTrackerTOFReFit(void) {
  PyWrapMapBase<MapCppTrackerTOFReFit>::PyWrapMapBaseModInit
                                        ("MapCppTrackerTOFReFit", "", "", "", "");
}


MapCppTrackerTOFReFit::MapCppTrackerTOFReFit() : MapBase<Data>("MapCppTrackerTOFReFit"),
                                           _helical_track_fitter(NULL),
                                           _radius_cut(10.0),
                                           _has_tof(false) {
}


MapCppTrackerTOFReFit::~MapCppTrackerTOFReFit() {
}


void MapCppTrackerTOFReFit::_birth(const std::string& argJsonConfigDocument) {
  // Pull out the global settings
  if (!Globals::HasInstance()) {
    GlobalsManager::InitialiseGlobals(argJsonConfigDocument);
  }
  Json::Value* json = Globals::GetConfigurationCards();
  std::vector<const MiceModule*> tof_modules = Globals::GetReconstructionMiceModules()->
                                           findModulesByPropertyString("SensitiveDetector", "TOF");
  _radius_cut          = (*json)["SciFiHelicalRadiusLimit"].asDouble();
  _has_tof             = ( tof_modules.size() > 0 ? true : false );
  // Values used to set the track rating:

  if (!_has_tof) {
    Squeak::mout(Squeak::warning) << "Tracker-TOF Refit Mapper is initialised but there is no TOF"
                                  << "in the reconstruction geometry";
  }

  _tof1_z = 0.0;
  bool found = false;
  for (std::vector<const MiceModule*>::const_iterator it = tof_modules.begin();
                                                                   it != tof_modules.end(); ++it) {
    const MiceModule* parent = (*it)->mother();
    if (parent->propertyInt("Station") == 1) {
      _tof1_z = parent->globalPosition().z();
      found = true;
      break;
    }
  }
  if (!found) {
    throw Exceptions::Exception(Exceptions::nonRecoverable, "MapCppTrackerTOFReFit::_birth()",
                                                           "Could not locate the TOF1 Detetector");
  }
  double lower, upper;
  MaterialModelAxialLookup::GetBounds(_tof1_z, lower, upper);
  _tof1_z = lower; // Find the start of the TOF detector.

  SciFiTrackerMap& geo_map = Globals::GetSciFiGeometryHelper()->GeometryMap();

  // Set up final track fit (Kalman filter)
  HelicalPropagator* helical_prop = new HelicalPropagator(Globals::GetSciFiGeometryHelper());
  helical_prop->SetCorrectPz(true);
  helical_prop->SetIncludeMCS(true);
  helical_prop->SetSubtractELoss(true);
  _helical_track_fitter = new Kalman::TrackFit(helical_prop);

  // Each measurement plane has a unique alignment and rotation => they all need their own
  //  measurement object.
  for (SciFiTrackerMap::iterator track_it = geo_map.begin();
                                                           track_it != geo_map.end(); ++track_it) {
    int tracker_const = (track_it->first == 0 ? -1 : 1);
    for (SciFiPlaneMap::iterator plane_it = track_it->second.Planes.begin();
                                           plane_it != track_it->second.Planes.end(); ++plane_it) {

      int id = plane_it->first * tracker_const;
      _helical_track_fitter->AddMeasurement(id,
                                             new MAUS::SciFiHelicalMeasurements(plane_it->second));

      if (id == -15) {
        _tracker_z = plane_it->second.GlobalPosition.Z();
      }
    }
  }
}


void MapCppTrackerTOFReFit::_death() {
  if (_helical_track_fitter) {
    delete _helical_track_fitter;
    _helical_track_fitter = NULL;
  }
}


void MapCppTrackerTOFReFit::_process(Data* data) const {
  if (!_has_tof) {
    return;
  }

  Spill& spill = *(data->GetSpill());
  /* return if not physics spill */
  if (spill.GetDaqEventType() != "physics_event")
    return;

  if (spill.GetReconEvents()) {
    for (unsigned int k = 0; k < spill.GetReconEvents()->size(); k++) {
      SciFiEvent *event = spill.GetReconEvents()->at(k)->GetSciFiEvent();
      if (!event) {
        continue;
      }

      TOFEvent *tof_event = spill.GetReconEvents()->at(k)->GetTOFEvent();
      if (!tof_event) {
        continue;
      }

      SciFiTrackPArray tracks = event->scifitracks();
      SciFiTrackPArray new_tracks;

      for (SciFiTrackPArray::iterator it = tracks.begin(); it != tracks.end(); ++it) {
        try {
          SciFiSeed* seed = (*it)->scifi_seed();
          switch (seed->getAlgorithm()) {
            case 1 :
              if ( needs_refitting(seed) || (*it)->is_dud() ) {

                if ( calculate_pz_seed(seed, tof_event) ) {
                  SciFiTrack* new_track = track_fit_helix(seed);
                  new_track->SetWasRefit(1);
                  new_tracks.push_back(new_track);
                  delete *it;
                  continue;
                }
              }
              new_tracks.push_back((*it));
              break;
            case 0 :
              if ( (*it)->tracker() == 0 ) {
                if ( calculate_straight_pz_seed(seed, tof_event) ) {
                  SciFiTrack* new_track = track_fit_helix(seed);
                  new_track->SetWasRefit(2);
                  new_tracks.push_back(new_track);
                  delete *it;
                  continue;
                }
              }
              new_tracks.push_back((*it));
              break;
            default:
              break;
          }
        } catch (Exceptions::Exception& e) {
          std::cerr << "Track Fit Failed: " << e.what();
        }
      }

      event->set_scifitracks(new_tracks);
    }
  } else {
    std::cout << "No recon events found\n";
  }
}

SciFiTrack* MapCppTrackerTOFReFit::track_fit_helix(SciFiSeed* seed) const {

  SciFiHelicalPRTrack* helical = static_cast<SciFiHelicalPRTrack*>(seed->getPRTrackTobject());

  Kalman::Track data_track = BuildTrack(helical, Globals::GetSciFiGeometryHelper(), 5);
  Kalman::State kalman_seed(seed->getStateVector(), seed->getCovariance());

  _helical_track_fitter->SetTrack(data_track);
  _helical_track_fitter->SetSeed(kalman_seed);

  _helical_track_fitter->Filter(false);
  _helical_track_fitter->Smooth(false);

  SciFiTrack* track = ConvertToSciFiTrack(_helical_track_fitter,
                                                       Globals::GetSciFiGeometryHelper(), helical);
  track->set_scifi_seed_tobject(seed);

  ThreeVector seed_pos = track->GetSeedPosition();
  ThreeVector seed_mom = track->GetSeedMomentum();

  return track;
}

bool MapCppTrackerTOFReFit::needs_refitting(MAUS::SciFiSeed* seed) const {
  if ( seed->getTracker() != 0 ) {
    return false;
  } else if ( static_cast<SciFiHelicalPRTrack*>(seed->getPRTrackTobject())->get_R() >
                                                                              this->_radius_cut ) {
    return false;
  }
  return true;
}

bool MapCppTrackerTOFReFit::calculate_pz_seed(MAUS::SciFiSeed* seed, TOFEvent* tof_event) const {
  TOFEventSpacePoint* spacepoints = tof_event->GetTOFEventSpacePointPtr();

  if (spacepoints->GetTOF0SpacePointArraySize() != 1 )
    return false;
  if (spacepoints->GetTOF1SpacePointArraySize() != 1 )
    return false;

  TOFSpacePoint tof0 = spacepoints->GetTOF0SpacePointArrayElement(0);
  TOFSpacePoint tof1 = spacepoints->GetTOF1SpacePointArrayElement(0);

  double delta_T = tof1.GetTime() - tof0.GetTime();
  double delta_Z = tof1.GetGlobalPosZ() - tof0.GetGlobalPosZ();

  double beta = delta_Z / ( delta_T * CLHEP::c_light);
  double gamma = 1.0 / sqrt(1.0 - beta*beta);
  double pz = gamma*beta*Recon::Constants::MuonMass;

  // Propagate from TOF to tracker
  pz = calculate_new_pz(pz);

  TMatrixD vector = seed->getStateVector();

  if (vector(4, 0) >= 0.0) {
    vector(4, 0) =  1.0 / pz;
  } else {
    vector(4, 0) = -1.0 / pz;
  }

  seed->setStateVector(vector);

  return true;
}

bool MapCppTrackerTOFReFit::calculate_straight_pz_seed(MAUS::SciFiSeed* seed,
                                                                       TOFEvent* tof_event) const {
  TOFEventSpacePoint* spacepoints = tof_event->GetTOFEventSpacePointPtr();

  if (spacepoints->GetTOF0SpacePointArraySize() != 1 )
    return false;
  if (spacepoints->GetTOF1SpacePointArraySize() != 1 )
    return false;

  TOFSpacePoint tof0 = spacepoints->GetTOF0SpacePointArrayElement(0);
  TOFSpacePoint tof1 = spacepoints->GetTOF1SpacePointArrayElement(0);

  double delta_T = tof1.GetTime() - tof0.GetTime();
  double delta_Z = tof1.GetGlobalPosZ() - tof0.GetGlobalPosZ();

  double beta = delta_Z / (delta_T * CLHEP::c_light);
  double gamma = 1.0 / sqrt(1.0 - beta*beta);
  double pz = gamma*beta*Recon::Constants::MuonMass;

  // Propagate from TOF to tracker
  pz = calculate_new_pz(pz);

  TMatrixD vector = seed->getStateVector();
  TMatrixD new_vector(5, 1);

  new_vector(0, 0) = vector(0, 0);
  new_vector(1, 0) = 0.0;
  new_vector(2, 0) = vector(2, 0);
  new_vector(3, 0) = 0.0;

  if (seed->getTracker() == 0) {
    new_vector(4, 0) =  1.0 / pz;
  } else {
    new_vector(4, 0) = -1.0 / pz;
  }
  seed->setStateVector(new_vector);


  TMatrixD covariance = seed->getCovariance();
  TMatrixD new_covariance(5, 5);

  for (unsigned int i = 0; i < 5; ++i) {
    for (unsigned int j = 0; j < 5; ++j) {
      new_covariance(i, j) = 0.0;
    }
  }
  new_covariance(0, 0) = 100.0;
  new_covariance(1, 1) = 100.0;
  new_covariance(2, 2) = 100.0;
  new_covariance(3, 3) = 100.0;
  new_covariance(4, 4) = 2.5e-7;

  seed->setCovariance(new_covariance);

  return true;
}

double MapCppTrackerTOFReFit::calculate_new_pz(double pz) const {
  double energy = std::sqrt(Recon::Constants::MuonMass*Recon::Constants::MuonMass + pz*pz);

  double lower;
  double upper;

  double start_position = _tof1_z;
  double position = _tof1_z;

  MaterialModelAxialLookup lookup(0, 0, position);

  do {
    lookup.SetMaterial(0, 0, position);
    MaterialModelAxialLookup::GetBounds(position, lower, upper);

    if (lower < _tof1_z) {
      start_position = _tof1_z;
    } else {
      start_position = lower;
    }

    if (upper > _tracker_z) {
      position = _tracker_z;
    } else {
      position = upper;
    }

    double dEdx = lookup.dEdx(energy, Recon::Constants::MuonMass, +1);
    energy += dEdx*(position-start_position);

    position += MaterialModelAxialLookup::GetZTolerance();
  } while (position < _tracker_z);

  return std::sqrt( energy*energy - Recon::Constants::MuonMass*Recon::Constants::MuonMass);
}
} // ~namespace MAUS