/* 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 <string>
#include <vector>

#include <sstream>

#include "Utils/Squeak.hh"

#include "src/map/MapCppReconCuts/ReconTrackerCuts.hh"

namespace MAUS {

    ReconTrackerCuts::~ReconTrackerCuts() {
    }

	ReconTrackerCuts::ReconTrackerCuts(Json::Value cutParams) {
		if (!cutParams.isNull()) {
			_min_p_value = cutParams["min_p_value"].asFloat();
			_TKU_min_mom = cutParams["TKU_min_mom"].asFloat();
			_TKU_max_mom = cutParams["TKU_max_mom"].asFloat();
		    _TKD_min_mom = cutParams["TKD_min_mom"].asFloat();
			_TKD_max_mom = cutParams["TKD_max_mom"].asFloat();
        	_min_mom_loss = cutParams["min_mom_loss"].asFloat();
			_max_mom_loss = cutParams["max_mom_loss"].asFloat();
	 		_mom_correction = cutParams["mom_correction"].asFloat();
			_min_mass = cutParams["min_mass"].asFloat();
			_max_mass = cutParams["max_mass"].asFloat();
        } else {
			// default cut values
        	_min_p_value = 0.02;
			_TKU_min_mom = 130.;
			_TKU_max_mom = 150.;
			_TKD_min_mom = 50.;
			_TKD_max_mom = 100.;
			_min_mom_loss = 5;
			_max_mom_loss = 43.4;
			_mom_correction = 0;
			_min_mass = 99.;
			_max_mass = 110.;
        }
	}

	/** Copy constructor */
	ReconTrackerCuts::ReconTrackerCuts(ReconTrackerCuts& copy) {
		_min_p_value = copy._min_p_value;
		_TKU_min_mom = copy._TKU_min_mom;
		_TKU_max_mom = copy._TKU_max_mom;
		_TKD_min_mom = copy._TKD_min_mom;
		_TKD_max_mom = copy._TKD_max_mom;
		_min_mom_loss = copy._min_mom_loss;
		_max_mom_loss = copy._max_mom_loss;
		_mom_correction = copy._mom_correction;
		_min_mass = copy._min_mass;
		_max_mass = copy._max_mass;
	}

	/** operator= */
	ReconTrackerCuts& ReconTrackerCuts::operator=(ReconTrackerCuts& rhs) {
		_min_p_value = rhs._min_p_value;
		_TKU_min_mom = rhs._TKU_min_mom;
		_TKU_max_mom = rhs._TKU_max_mom;
		_TKD_min_mom = rhs._TKD_min_mom;
		_TKD_max_mom = rhs._TKD_max_mom;
		_min_mom_loss = rhs._min_mom_loss;
		_max_mom_loss = rhs._max_mom_loss;
		_mom_correction = rhs._mom_correction;
		_min_mass = rhs._min_mass;
		_max_mass = rhs._max_mass;
		return *this;
	}

	void ReconTrackerCuts::DoCuts(ReconEvent* anEvent) {

		MAUS::SciFiEvent *sciFiEvent = anEvent->GetSciFiEvent();

		if(sciFiEvent == NULL) return;

		std::vector<MAUS::SciFiTrack*> tracks = sciFiEvent->scifitracks();

		bool single_track_cut = false;
		bool p_value_cut = false;
		float track_p_value = 0;

		bool TKU_hits_cut = false;
		int TKU_hits = -1;
		int tracker = -1;
		bool TKD_hits_cut = false;
		int TKD_hits = -1;

		bool TKU_mom_cut = false;
		float TKU_mom = 0;
		float TKU_mom_saved = 0;
		bool TKD_mom_cut = false;
		float TKD_mom = 0;
		float TKD_mom_saved = 0;

		float min_mom_cut = 0;
		float max_mom_cut = 0;
		bool mom_loss_cut = false;
		bool mass_cut = false;

		// Single Track Cut and P-value cut
		if (tracks.size() == 1) {
			single_track_cut = true;

			track_p_value = tracks.at(0)->P_value();
			if (track_p_value > _min_p_value) {
				p_value_cut = true;
			}
		}

		std::vector<MAUS::SciFiHelicalPRTrack*> pr_tracks = sciFiEvent->helicalprtracks();

		// Check there is only 1 pr track and that it hit 5 stations
		if (pr_tracks.size() == 1) {
			int pr_hits = pr_tracks.at(0)->get_num_points();
			tracker = pr_tracks.at(0)->get_tracker();
			if (tracker == 0) {
				TKU_hits = pr_hits;
				if (pr_hits == 5) {
					TKU_hits_cut = true;
				}
			} else if (tracker == 1) {
				TKD_hits = pr_hits;
				if (pr_hits == 5) {
					TKD_hits_cut = true;
				}
			}
		}

		// TKU and TKD momentum cuts
		std::vector<MAUS::SciFiTrack*>::iterator track_it;
		if (tracks.size() != 0) {
			for (track_it = tracks.begin(); track_it != tracks.end(); track_it++) {
				std::vector<MAUS::SciFiTrackPoint*> tr_point = (*track_it)->scifitrackpoints();
				std::vector<MAUS::SciFiTrackPoint*>::iterator tp_iter;
				for (tp_iter = tr_point.begin(); tp_iter != tr_point.end(); tp_iter++) {
					MAUS::SciFiTrackPoint* point = (*tp_iter);
					if(point->plane() == 2 && point->station() == 5) {
						if (point->tracker() == 0) {
							// point is in TKU
							TKU_mom = point->mom().mag();
							if (TKU_mom >= _TKU_min_mom && TKU_mom <= _TKU_max_mom) {
								TKU_mom_cut = true;
								TKU_mom_saved = TKU_mom;
							}
						} else if (point->tracker() == 1) {
							// point is in TKD
							TKD_mom = point->mom().mag();
							if (TKD_mom >= _TKD_min_mom && TKD_mom <= _TKD_max_mom) {
								TKD_mom_cut = true;
								TKD_mom_saved = TKD_mom;
							}
						}
					}
				}
    		}
		}

		// Momentum loss and mass cuts
		float dt = 0;
		float beta_gamma_tof = 0;
		float mass = 0;

		MAUS::TOFEvent * tofEvent = anEvent->GetTOFEvent();
		if (tofEvent != NULL) {
			MAUS::TOFEventSpacePoint SpacePoint = tofEvent->GetTOFEventSpacePoint();
			if (SpacePoint.GetTOF0SpacePointArray().size() == 1 &&
			SpacePoint.GetTOF1SpacePointArray().size() == 1) {
				dt = SpacePoint.GetTOF1SpacePointArray()[0].GetTime() -
					SpacePoint.GetTOF0SpacePointArray()[0].GetTime();
			}
		}

		if (dt != 0) {
			double m = 105.6583715;
			float dt_e = 25.48;
		    double beta_tof = dt_e/dt;
			// float mom_corr = 18.82;

    		if (beta_tof <= 1.0 && beta_tof >= 0) {
				min_mom_cut = (TKU_mom + _min_mom_loss) / m;
				max_mom_cut = (TKU_mom + _max_mom_loss) / m;
				double gamma_tof = 1.0/(sqrt(1.0 - beta_tof*beta_tof));
				beta_gamma_tof = beta_tof*gamma_tof;

				// Momentum loss cut
				if (beta_gamma_tof >= min_mom_cut && beta_gamma_tof <= max_mom_cut) {
				     mom_loss_cut = true;
				}

				// Mass cut
				mass = (TKU_mom + _mom_correction)/beta_gamma_tof;
				if (mass >= _min_mass && mass <= _max_mass) {
					mass_cut = true;
				}
			}
		}

		std::stringstream cutDesc;
		// Add cut objects to the event
		anEvent->GetCutsList()->push_back(new ReconCutDataBase("Single_track",
		"tracks==1", single_track_cut, tracks.size()));

        cutDesc << "track_p_value>" << _min_p_value;
		anEvent->GetCutsList()->push_back(new ReconCutDataBase("Track_p_value",
		cutDesc.str(), p_value_cut, track_p_value));

		anEvent->GetCutsList()->push_back(new ReconCutDataBase("TKU_hits",
		"pr_tracks==1&&pr_track_hits==5&&tracker==0", TKU_hits_cut, TKU_hits));

		anEvent->GetCutsList()->push_back(new ReconCutDataBase("TKD_hits",
		"pr_tracks==1&&pr_track_hits==5&&tracker==1", TKD_hits_cut, TKD_hits));

		cutDesc.str("");
		cutDesc << "TKU_mom>="<< _TKU_min_mom << "&&TKU_mom<=" << _TKU_max_mom;
		// Save momentum of track that passed the cut, if there is one.
		if(TKU_mom_saved>0) TKU_mom= TKU_mom_saved;
		anEvent->GetCutsList()->push_back(new ReconCutDataBase("TKU_mom_cut",
		cutDesc.str(), TKU_mom_cut, TKU_mom));

		cutDesc.str("");
		cutDesc << "TKD_mom>=" << _TKD_min_mom << "&&TKD_mom<=" << _TKD_max_mom;
		// Save momentum of track that passed the cut, if there is one.
		if(TKD_mom_saved>0) TKD_mom= TKD_mom_saved;
		anEvent->GetCutsList()->push_back(new ReconCutDataBase("TKD_mom_cut",
		cutDesc.str(), TKD_mom_cut, TKD_mom));

		cutDesc.str("");
		cutDesc << "beta_gamma_tof>=" << min_mom_cut << "&&beta_gamma_tof<=" << max_mom_cut;
		anEvent->GetCutsList()->push_back(new ReconCutDataBase("Mom_loss_cut",
		cutDesc.str(), mom_loss_cut, beta_gamma_tof));

		cutDesc.str("");
		cutDesc << "mass>=" << _min_mass << "&&mass<=" << _max_mass;
		anEvent->GetCutsList()->push_back(new ReconCutDataBase("Mass_cut",
		cutDesc.str(), mass_cut, mass));

// Squeak::mout(Squeak::warning) << "in ReconTrackerCuts::DoCuts" << std::endl;
	}
} // namespace MAUS