#include <iostream>
#include <string>

#include "TGenPhaseSpace.h"
#include "TRandom3.h"

#include "CLHEP/Vector/LorentzVector.h"

#include "IUAmpTools/Kinematics.h"
#include "IUAmpTools/ConfigFileParser.h"
#include "IUAmpTools/ConfigurationInfo.h"
#include "IUAmpTools/AmpToolsInterface.h"

#include "particleType.h"
#include "AMPTOOLS_DATAIO/ROOTDataWriter.h"
#include "AMPTOOLS_DATAIO/HDDMDataWriter.h"
#include "AMPTOOLS_AMPS/Uniform.h"

#include "HCA_Amplitudes/HCAmplitude_Chain.h"
#include "HCA_Amplitudes/BreitWigner.h"
#include "HCA_Amplitudes/TChannel_Basic.h"

using namespace std;
using namespace CLHEP;

int main(int argc, char** argv)
{
	unsigned int locRandomNumberSeed = 0;
	bool locWeightedOutputFlag = false; //false for accept/reject, true to save all but save weights
	bool locHDDMOutputFlag = true;
	bool locROOTOutputFlag = true;
	bool locPhaseSpaceFlag = false; //if true, generate as flat phase space

	// For Phase-Space Generator //SET ON INPUT!!!
	double locBeamEnergy = 2.0;
	double locTargetMass = 1.87561;
	Double_t* locProductMasses = NULL;
	unsigned int locNumProducts = 0;

	/************************************* PARSE COMMAND LINE *************************************/

	// Parse Command Line Arguments
	if((argc != 4) && (argc != 5))
	{
		cout << "Usage:" << endl;
		cout << "	generatePhysics <config file name> <output file name (no extension!)> <number of events> <random # seed (optional)>" << endl;
		return 0;
	}

	string cfgname(argv[1]);
	string locOutputFileName(argv[2]);
	int nevents(atoi(argv[3]));
	if(argc == 5)
		locRandomNumberSeed = (atoi(argv[4]));

	cout << "Config file name = " << cfgname << endl;
	cout << "Output file name = " << locOutputFileName << endl;
	cout << "Number of events = " << nevents << endl;

	/************************************** SETUP GENERATOR ***************************************/

	//save particle types for HDDM output //NO WAY TO SET TARGET OR WEIGHT YET!!!
	vector<int> locParticleTypes;
	locParticleTypes.push_back(1);
	locParticleTypes.push_back(8);
	locParticleTypes.push_back(9);
	locParticleTypes.push_back(45);

	AmpToolsInterface* locAmpToolsInterface = NULL;
	ConfigurationInfo* locConfigurationInfo = NULL;
	ReactionInfo* locReactionInfo = NULL;
	if(!locPhaseSpaceFlag)
	{
		// Parse Config File
		ConfigFileParser parser(cfgname);
		locConfigurationInfo = parser.getConfigurationInfo();
		locConfigurationInfo->display();
		locReactionInfo = locConfigurationInfo->reactionList()[0];

		// Register Amplitudes
		AmpToolsInterface::registerAmplitude(BreitWigner());
		AmpToolsInterface::registerAmplitude(HCAmplitude_Chain());
		AmpToolsInterface::registerAmplitude(TChannel_Basic());
		AmpToolsInterface::registerAmplitude(Uniform());
		locAmpToolsInterface = new AmpToolsInterface(locConfigurationInfo, AmpToolsInterface::kMCGeneration);
	}

	// Create Data Writer(s)
	HDDMDataWriter* locHDDMDataWriter = locHDDMOutputFlag ? new HDDMDataWriter(locOutputFileName + string(".hddm")) : NULL;
	ROOTDataWriter* locROOTDataWriter = locROOTOutputFlag ? new ROOTDataWriter(locOutputFileName + string(".root")) : NULL;

	/******************************** GENERATE PHASE SPACE EVENTS *********************************/

	TLorentzVector locBeamP4(0.0, 0.0, locBeamEnergy, locBeamEnergy);
	TLorentzVector locTargetP4(0.0, 0.0, 0.0, locTargetMass);
	TLorentzVector locInitP4 = locBeamP4 + locTargetP4;
	if(locReactionInfo != NULL)
	{
		vector<string> locParticleList = locReactionInfo->particleList();
		locNumProducts = locParticleList.size();
		locProductMasses = new Double_t[locNumProducts];
		//strip everything
		for(size_t loc_i = 0; loc_i < locParticleList.size(); ++loc_i)
		{
			size_t locUnderscoreIndex = locParticleList[loc_i].rfind("_");
			string locParticleName = (locUnderscoreIndex != string::npos) ? locParticleList[loc_i].substr(0, locUnderscoreIndex) : locParticleList[loc_i];
			Particle_t locPID = ParticleEnum(locParticleName.c_str());
			locProductMasses[loc_i] = ParticleMass(locPID);
		}
	}

	// use ROOT to generate phase space
	TRandom3 locRandomNumberGenerator(locRandomNumberSeed);
	TGenPhaseSpace locPhaseSpaceGenerator;
	locPhaseSpaceGenerator.SetDecay(locInitP4, locNumProducts, locProductMasses);
	double locMaxWeight = locPhaseSpaceGenerator.GetWtMax();

	// Generate
	for(int loc_i = 0; loc_i < nevents; ++loc_i)
	{
		double locRandomNumber = locRandomNumberGenerator.Rndm();
		double locWeight = locPhaseSpaceGenerator.Generate();
		if(locWeight < locRandomNumber * locMaxWeight)
		{
			--loc_i;
			continue;
		}

		// pack the decay products into a Kinematics object
		vector<HepLorentzVector> locProductP4s;
		for(unsigned int loc_j = 0; loc_j < locNumProducts; ++loc_j)
		{
			TLorentzVector* locP4 = locPhaseSpaceGenerator.GetDecay(loc_j);
			locProductP4s.push_back(HepLorentzVector(locP4->Px(), locP4->Py(), locP4->Pz(), locP4->E()));
		}
		Kinematics* locKinematics = new Kinematics(locProductP4s);

		if(!locPhaseSpaceFlag)
		{
			// load events into the AmpToolsInterface
			locAmpToolsInterface->loadEvent(locKinematics, loc_i, nevents);
		}
		else //save output
		{
			if(locHDDMOutputFlag)
				locHDDMDataWriter->writeEvent(*locKinematics, locParticleTypes);
			if(locROOTOutputFlag)
				locROOTDataWriter->writeEvent(*locKinematics);
		}
		delete locKinematics;
	}
	cout << "Phase space events generated." << endl;

	if(locPhaseSpaceFlag)
	{
		if(locHDDMOutputFlag)
			delete locHDDMDataWriter;
		if(locROOTOutputFlag)
			delete locROOTDataWriter;
		return 0;
	}

	// calculate intensities for all events // unless generating pure phase space
	cout << "calculating intensities..." << endl;
	string locReactionName = locReactionInfo->reactionName();
	double locMaxIntensity = 1.5*locAmpToolsInterface->processEvents(locReactionName);
	cout << "... finished calculating intensities" << endl;

	/************************************ REJECT / SAVE EVENTS ************************************/

	// loop over all events again
	for(int loc_i = 0; loc_i < nevents; ++loc_i)
	{
		double locIntensity = locAmpToolsInterface->intensity(loc_i);
		double locRandomNumber = locRandomNumberGenerator.Rndm();
		double locRandomIntensity = locRandomNumber * locMaxIntensity;
		if((locIntensity < locRandomIntensity) && (!locWeightedOutputFlag))
			continue; //rejected

		Kinematics* locKinematics = locAmpToolsInterface->kinematics(loc_i);
		if(!locWeightedOutputFlag)
			locKinematics->setWeight(1.0); // we want to save events with weight 1

		if(locHDDMOutputFlag)
			locHDDMDataWriter->writeEvent(*locKinematics, locParticleTypes);
		if(locROOTOutputFlag)
			locROOTDataWriter->writeEvent(*locKinematics);

		delete locKinematics;
	}

	if(locHDDMOutputFlag)
		delete locHDDMDataWriter;
	if(locROOTOutputFlag)
		delete locROOTDataWriter;

	return 0;
}