/* 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/MapCppTrackerDownstreamReFit/MapCppTrackerDownstreamReFit.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_MapCppTrackerDownstreamReFit(void) {
  PyWrapMapBase<MapCppTrackerDownstreamReFit>::PyWrapMapBaseModInit
                                        ("MapCppTrackerDownstreamReFit", "", "", "", "");
}


MapCppTrackerDownstreamReFit::MapCppTrackerDownstreamReFit()
  : MapBase<Data>("MapCppTrackerDownstreamReFit"),
  _helical_track_fitter(NULL),
  _radius_cut(10.0) {
}


MapCppTrackerDownstreamReFit::~MapCppTrackerDownstreamReFit() {
}


void MapCppTrackerDownstreamReFit::_birth(const std::string& argJsonConfigDocument) {
  // Pull out the global settings
  if (!Globals::HasInstance()) {
    GlobalsManager::InitialiseGlobals(argJsonConfigDocument);
  }
  Json::Value* json = Globals::GetConfigurationCards();
  _radius_cut          = (*json)["SciFiHelicalRadiusLimit"].asDouble();

  _tracker_z = 0.0;
  _downstream_tracker_z = 0.0;

  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);

  // Find trackers and initialise measurement objects.
  // Each measurement plane has a unique alignment and rotation => they all need their own
  // measurement object.
  SciFiTrackerMap& geo_map = Globals::GetSciFiGeometryHelper()->GeometryMap();

  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) {
        _downstream_tracker_z = plane_it->second.GlobalPosition.Z();
      }
      if (id == -1) {
        _tracker_z = plane_it->second.GlobalPosition.Z();
      }
    }
  }
}


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


void MapCppTrackerDownstreamReFit::_process(Data* data) const {
  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;
      }

      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()) {
                std::cerr << "Found 1 needs refitting" << std::endl;

                if (calculate_pz_seed(seed, event)) {
                  std::cerr << "Calculated new pz" << std::endl;
                  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() == 1) {
                std::cerr << "Found 0 needs reffiting" << std::endl;
                if (calculate_straight_pz_seed(seed, event)) {
                  std::cerr << "Calculated new pz" << std::endl;
                  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* MapCppTrackerDownstreamReFit::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 MapCppTrackerDownstreamReFit::needs_refitting(MAUS::SciFiSeed* seed) const {
  if (seed->getTracker() != 1) { // Dowmnstream Tracker
    return false;
  } else if (static_cast<SciFiHelicalPRTrack*>(seed->getPRTrackTobject())->get_R() >
                                                                               this->_radius_cut) {
    return false; // With a small radius
  }
  return true; // Then refit it.
}

bool MapCppTrackerDownstreamReFit::calculate_pz_seed(MAUS::SciFiSeed* seed,
                                                                   SciFiEvent* scifi_event) const {

  std::vector<SciFiTrack*> found_tracks;
  SciFiTrackPArray tracks = scifi_event->scifitracks();
//  std::cerr << tracks.size() << " Tracks" << std::endl;

  for (SciFiTrackPArray::iterator it = tracks.begin(); it != tracks.end(); ++it) {
    if ((*it)->tracker() == 0) {
      found_tracks.push_back((*it));
    }
  }

  if (found_tracks.size() != 1) {
    return false;
  }

  if (found_tracks[0]->GetAlgorithmUsed() != 1) {
    return false;
  }

  double pz = 0.0;

  SciFiTrackPointPArray trackpoints = found_tracks[0]->scifitrackpoints();
  for (SciFiTrackPointPArray::iterator it = trackpoints.begin(); it != trackpoints.end(); ++it) {
    if ((*it)->station() == 1 && (*it)->plane() == 0) {
      pz = (*it)->mom().z();
      break;
    }
  }

  if (std::fabs(pz) < 0.001) {
    return false;
  }

  // Propagate from Upstream to downstream tracker
  pz = calculate_new_pz(pz);
  std::cerr << " -> " << pz << std::endl;

  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 MapCppTrackerDownstreamReFit::calculate_straight_pz_seed(MAUS::SciFiSeed* seed,
                                                                   SciFiEvent* scifi_event) const {
  std::vector<SciFiTrack*> found_tracks;
  SciFiTrackPArray tracks = scifi_event->scifitracks();

  for (SciFiTrackPArray::iterator it = tracks.begin(); it != tracks.end(); ++it) {
    if ((*it)->tracker() == 0) {
      found_tracks.push_back((*it));
    }
  }

  if (found_tracks.size() != 1) {
    return false;
  }

  if (found_tracks[0]->GetAlgorithmUsed() != 1) {
    return false;
  }

  double pz = 0.0;

  SciFiTrackPointPArray trackpoints = found_tracks[0]->scifitrackpoints();
  for (SciFiTrackPointPArray::iterator it = trackpoints.begin(); it != trackpoints.end(); ++it) {

    if ((*it)->station() == 1 && (*it)->plane() == 0) {
      pz = (*it)->mom().z();
    }
  }

  if (std::fabs(pz) < 0.001) {
    return false;
  }

  // Propagate from upstream to downstream tracker
  pz = calculate_new_pz(pz);
  std::cerr << " -> " << pz << std::endl;

  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 MapCppTrackerDownstreamReFit::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 = _tracker_z;
  double position = _tracker_z;

  MaterialModelAxialLookup lookup(0, 0, position);

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

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

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

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

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

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