/*
 * main.cpp
 *
 *  Created on: Jan 11, 2012
 *      Author: Anastasia Karavdina
 */

#include "ResonanceP.h"
#include "StableP.h"
#include "Particle.h"
#include "TwoParticleState.h"
#include "Chain.h"
#include "Amplitude.h"
#include "Intensity.h"
#include "Util.h"
#include "TLorentzVector.h"
#include "TGenPhaseSpace.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TCanvas.h"
#include <iostream>
#include <list>
#include <string>
#include <sstream>
#include <utility>
#include <fstream>
int main(int __argc,char *__argv[]){
	///Read input from command line -----------------------
	int sizemass = 20000000;//number of events to be generated
	std::string nStr="";

	// decode arguments
	if( __argc>1 && ( strcmp( __argv[1], "-help" ) == 0
			|| strcmp( __argv[1], "--help" ) == 0 ) ){

		std::cout << "This is generator based on amplitude calculation in helicity formalism with parameters\n"
				<<"-n Number of events \n"
				<<"Have fun! \n"
				<< std::endl;
		return 0;
	}
	while ((optind < (__argc-1) ) && (__argv[optind][0]=='-')) {
		bool found=false;
		std::string sw = __argv[optind];
		if (sw=="-n"){
			optind++;
			nStr = __argv[optind];
			found=true;
		}

		if (!found){
			std::cout<< "Unknown switch: "
					<< __argv[optind] <<std::endl;
			optind++;
		}

		while ( (optind < __argc ) && __argv[optind][0]!='-' ) optind++;
	}

	std::stringstream nSStr(nStr);
	nSStr >> sizemass;
	///----------------------------------------------------

	///External parameters--------------------------------
	double Esys=1;//Initial energy
	double  Jsys;//Initial momenta
	list<Particle> lParticles;//list of final particles
	vector<int> pParticles;
	list<ResonanceP> lResonances;//list of resonances
	vector<int> pResonances;//list of resonances IDs
	vector<int> chResonances;//number of element in chain (1=last decay)
	///----------------------------------------------------

	///Read input ASCII file ------------------------------
	// decode arguments from input file
	string idStr="";
	string pStr="";
	int ip=0;
	ifstream input("HelGenParams.txt");
	string sout;
	while(!input.eof()){
		input>>sout;
		if(sout=="Input"){ //TODO: read\take into account more parameters from initial state
			while(sout!="}"){
				input>>sout;
				if(sout=="*"){
					input>>sout;
					if(sout=="energy"){
						input>>sout;
						idStr = sout;
						stringstream eSStr(idStr);
						eSStr >> Esys;
					}
					if(sout=="J"){
						input>>sout;
						idStr = sout;
						stringstream jSStr(idStr);
						jSStr >> Jsys;
					}
				}
			}
		}

		if (sout=="Output"){
			int fpID;// IDs of final particle
			double fpM;//mass of final particle
			double fpS;//spin of final particle
			int fpP;//P parity of final particle
			int fpC;//C parity of final particle //TODO: take into account C parity!

			while(sout!="}"){
				input>>sout;
				if(sout=="*"){
					input>>sout;
					input>>sout;
					idStr = sout;
					stringstream idSStr(idStr);
					idSStr >> fpID;
					pParticles.push_back(fpID);

					input>>sout;
					pStr=sout;
					stringstream pSStr(pStr);
					pSStr>>fpM;
					input>>sout;
					pStr=sout;
					stringstream pSStr1(pStr);
					pSStr1>>fpS;
					input>>sout;
					pStr=sout;
					stringstream pSStr2(pStr);
					pSStr2>>fpP;
					input>>sout;
					pStr=sout;
					stringstream pSStr3(pStr);
					pSStr3>>fpC;
					Particle *stab1 = new StableP(fpID,fpM,fpS,fpP,fpC);
					lParticles.push_back(*stab1);
				}
			}
		}
		//cout<<"size(list of particles)="<<lParticles.size()<<endl;
		if(sout=="Resonances"){
			int rID;//resonance ID
			double rM;//resonance mass
			double rW;//resonance width
			double rS;//resonance spin
			int rP;//resonance P-parity
			int rC;//resonance C-parity

			int fpID1;//ID of 1st particle for resonance decay
			int fpID2;//ID of 2nd particle for resonance decay
			while(sout!="}"){
				input>>sout;
				if(sout=="*"){
					input>>sout;
					input>>sout;
					idStr = sout;
					stringstream idSStr(idStr);
					idSStr >> rID;
					pResonances.push_back(rID);
					input>>sout;
					pStr = sout;
					stringstream pSStr(pStr);
					pSStr >> rM;
					input>>sout;
					pStr = sout;
					stringstream pSStr1(pStr);
					pSStr1 >> rW;
					input>>sout;
					pStr = sout;
					stringstream pSStr2(pStr);
					pSStr2 >> rS;
					input>>sout;
					pStr = sout;
					stringstream pSStr3(pStr);
					pSStr3 >> rP;
					input>>sout;
					pStr = sout;
					stringstream pSStr4(pStr);
					pSStr4 >> rC;
					input>>sout;
					input>>sout;
					idStr = sout;
					stringstream idSStrp1(idStr);
					idSStrp1 >> fpID1;
					input>>sout;
					input>>sout;
					idStr = sout;
					stringstream idSStrp2(idStr);
					idSStrp2 >> fpID2;
					ResonanceP res1(rID,rM,rS,rP,rC);
					res1.SetWidth(rW);//?
					res1.SetDecay(fpID1,fpID2);
					res1.Print();
					lResonances.push_back(res1);
				}//if(sout=="*")
			}//	while(sout!="}")
		}//if(sout=="Resonances")
	}//while(!input.eof())
	///----------------------------------------------------


	///Fill list of chain ------------------------------------
	//list<Chain> lchain = FillChain(lParticles, pParticles, lResonances, pResonances, chResonances);
	Util uhelp;
	list<Chain> lchain = uhelp.FillChain(lParticles, pParticles, lResonances, pResonances);
	///------------------------------------------------------


	/// Phase Space generation ------------------------------
	vector<TLorentzVector> pf_lab;//output 4-vectors
	list<TLorentzVector> lpf_lab;//output list of 4-vectors
	vector<int> lfpcode;//output particle code

	TLorentzVector W(0.0,0.0,0.0,Esys);
	cout<<"Esys = "<<Esys<<endl;
	const int sizem = lParticles.size();
	Double_t masses[sizem];
	int mi=0;
	cout<<"Order of particle:"<<endl;
	for (list<Particle>::iterator it = lParticles.begin(); it != lParticles.end(); it++){
		masses[mi]=(*it).Mass();
		lfpcode.push_back((*it).Code());
		cout<<(*it).Code()<<", ";
		mi++;
	}
	cout<<""<<endl;

	/// Amplitude initialization ----------------------------
	list<Amplitude> lAmpl;
	for (list<Chain>::iterator it = lchain.begin(); it != lchain.end(); it++){
		//cout<<"size of chain for ampl calc="<<(*it).GetStates().size()<<endl;
		Amplitude ampl1((*it));
		ampl1.Build(lfpcode);
		lAmpl.push_back(ampl1);
	}
	///------------------------------------------------------

	/// Intensity initialization ----------------------------
	Intensity intens(lAmpl);
	///------------------------------------------------------

	/// Test output -----------------------------------------
	/// Test Histogramm ------------------------------------------
	double minx = TMath::Power((masses[0]+masses[1]),2)-0.1;
	double maxx = TMath::Power((Esys-masses[0]),2)+0.1;
	double miny = TMath::Power((masses[0]+masses[1]),2)-0.1;
	double maxy = TMath::Power((Esys-masses[0]),2)+0.1;
	//  std::cout<<minx<<", "<<maxx<<", "<<miny<<", "<<maxy<<std::endl;
	//TH2F *h2 = new TH2F("h2","Uniform phase space di stribution;m^{2}_{123}, GeV^{2};m^{2}_{234}, GeV^{2}", 100,minx,maxx,100,miny,maxy); //test
	TH2F *h2 = new TH2F("h2","Uniform phase space di stribution;m^{2}_{12}, GeV^{2};m^{2}_{23}, GeV^{2}", 100,minx,maxx,100,miny,maxy); //test
	TH1F *h1m12i = new TH1F("h1m12i",";m_{123};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1m23i = new TH1F("h1m23i",";m_{234};Events",100,TMath::Sqrt(miny+0.1)-0.1,TMath::Sqrt(maxy));
//	TH2F *h2r = new TH2F("h2r","Result phase space distribution;m^{2}_{123}, GeV^{2};m^{2}_{234}, GeV^{2}", 100, minx,maxx, 100,miny,maxy); //test
	TH2F *h2r = new TH2F("h2r","Result phase space distribution;m^{2}_{12}, GeV^{2};m^{2}_{23}, GeV^{2}", 100, minx,maxx, 100,miny,maxy); //test
	TH1F *h1m12 = new TH1F("h1m12",";m_{123};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1m23 = new TH1F("h1m23",";m_{234};Events",100,TMath::Sqrt(miny+0.1)-0.1,TMath::Sqrt(maxy));

	//TH2F *h2r2 = new TH2F("h2r2","Result phase space distribution;m^{2}_{134}, GeV^{2};m^{2}_{234}, GeV^{2}", 100, minx,maxx, 100,miny,maxy); //test
	TH2F *h2r2 = new TH2F("h2r2","Result phase space distribution;m^{2}_{23}, GeV^{2};m^{2}_{12}, GeV^{2}", 100, minx,maxx, 100,miny,maxy); //test
	TH1F *h1m13 = new TH1F("h1m13",";m_{134};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *hinttest = new TH1F("h1r","Intensity;I", 1e4,0,1.); //test
	TH1F *hthetatest = new TH1F("hTheta","#theta;#theta, rad", 1e3,0,3.15); //test
	TH1F *hcostheta = new TH1F("hcosTheta","cos#theta;cos#theta",1e3,-1.1,1.1);
	TH1F *hphitest = new TH1F("hPhi","#phi;#phi, rad", 1e3,3.15,3.15); //test

	TCanvas *c1 = new TCanvas("test");
	c1->Divide(3,3);

	//input
	TH1F *h1t12 = new TH1F("h1t12",";m_{12};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1t13 = new TH1F("h1t13",";m_{13};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1t14 = new TH1F("h1t14",";m_{14};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1t23 = new TH1F("h1t23",";m_{23};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1t24 = new TH1F("h1t24",";m_{24};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	//result
	TH1F *h1tr12 = new TH1F("h1tr12",";m_{12};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1tr13 = new TH1F("h1tr13",";m_{13};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1tr14 = new TH1F("h1tr14",";m_{14};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1tr23 = new TH1F("h1tr23",";m_{23};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));
	TH1F *h1tr24 = new TH1F("h1tr24",";m_{24};Events",100,TMath::Sqrt(minx+0.1)-0.1,TMath::Sqrt(maxx));

	TCanvas *cIm = new TCanvas("PairInvarMass");
	cIm->Divide(5,2);
	///---------------------------------------------------------

	///------------------------------------------------------


	TGenPhaseSpace event;
	event.SetDecay(W, sizem, masses);
	for (Int_t n=0;n<sizemass;n++) {
		cout<<"n="<<n<<endl;
		double weight = event.Generate();
		pf_lab.erase(pf_lab.begin(),pf_lab.end());
		lpf_lab.erase(lpf_lab.begin(),lpf_lab.end());
		for(int mi=0;mi<sizem;mi++){
			pf_lab.push_back(*(event.GetDecay(mi)));
			lpf_lab.push_back(*(event.GetDecay(mi)));
		//	(event.GetDecay(mi))->Print();
		}
		///----------------------------------

		/*TLorentzVector p123 = pf_lab[0]+pf_lab[1]+pf_lab[2];
		TLorentzVector p234 = pf_lab[1]+pf_lab[2]+pf_lab[3];
		TLorentzVector p134 = pf_lab[0]+pf_lab[2]+pf_lab[3];*/
		TLorentzVector p123 = pf_lab[0]+pf_lab[1];
		TLorentzVector p234 = pf_lab[1]+pf_lab[2];
		TLorentzVector p134 = pf_lab[0]+pf_lab[2];
		//h2->Fill(p123.M2(), p234.M2() ,weight);
		h2->Fill((pf_lab[0]+pf_lab[1]).M2(), (pf_lab[1]+pf_lab[2]).M2() ,weight);
		h1m12i->Fill(p123.M(),weight);
		h1m23i->Fill(p234.M(),weight);
		hthetatest->Fill(pf_lab[0].Theta());
		hcostheta->Fill(pf_lab[0].CosTheta());
		hphitest->Fill(pf_lab[0].Phi());

		h1t12->Fill((pf_lab[0]+pf_lab[1]).M(),weight);
		h1t13->Fill((pf_lab[0]+pf_lab[2]).M(),weight);
		h1t14->Fill((pf_lab[0]+pf_lab[3]).M(),weight);
		h1t23->Fill((pf_lab[1]+pf_lab[2]).M(),weight);
		h1t24->Fill((pf_lab[1]+pf_lab[3]).M(),weight);
		///---------------------------------

		double new_weight = intens.Calc(lpf_lab,lfpcode);
		//cout<<"new_weight = "<<new_weight<<endl;
		new_weight *=weight;

		hinttest->Fill(new_weight);
	//	h2r->Fill(p123.M2(), p234.M2(),new_weight);
		//	h2r2->Fill(p134.M2(),p234.M2(),new_weight);
		h2r->Fill((pf_lab[0]+pf_lab[1]).M2(), (pf_lab[1]+pf_lab[2]).M2(),new_weight);
		h2r2->Fill((pf_lab[1]+pf_lab[2]).M2(),(pf_lab[0]+pf_lab[1]).M2(),new_weight);
		h1m12->Fill(p123.M(),new_weight);
		h1m23->Fill(p234.M(),new_weight);
		h1m13->Fill(p134.M(),new_weight);

		h1tr12->Fill((pf_lab[0]+pf_lab[1]).M(),new_weight);
		h1tr13->Fill((pf_lab[0]+pf_lab[2]).M(),new_weight);
		h1tr14->Fill((pf_lab[0]+pf_lab[3]).M(),new_weight);
		h1tr23->Fill((pf_lab[1]+pf_lab[2]).M(),new_weight);
		h1tr24->Fill((pf_lab[1]+pf_lab[3]).M(),new_weight);
	}
	/// 	----------------------------------------------------------

	cIm->cd(1);
	h1t12->Draw();
	cIm->cd(2);
	h1t13->Draw();
	cIm->cd(3);
	h1t14->Draw();
	cIm->cd(4);
	h1t23->Draw();
	cIm->cd(5);
	h1t24->Draw();
	cIm->cd(6);
	h1tr12->Draw();
	cIm->cd(7);
	h1tr13->Draw();
	cIm->cd(8);
	h1tr14->Draw();
	cIm->cd(9);
	h1tr23->Draw();
	cIm->cd(10);
	h1tr24->Draw();
	cIm->SaveAs("testPaitInvarMass.pdf");

	c1->cd(1);
	h2->Draw("COLZ");
	c1->cd(2);
	h1m12i->Draw();
	c1->cd(3);
	h1m23i->Draw();
	c1->cd(4);
	h2r->Draw("COLZ");
	c1->cd(5);
	h1m12->Draw();
	c1->cd(6);
	h1m23->Draw();
	c1->cd(7);
	h2r2->Draw("COLZ");
	c1->cd(8);
	h1m13->Draw();
	c1->cd(9);
	h1m23->Draw();
	c1->SaveAs("testDalitzPlot.pdf");

	TCanvas *c2 = new TCanvas("DalitzPlot");
	h2r->Draw("COLZ");
	c2->SaveAs("testDalitzPlot_onlyDP.root");

	TCanvas *c3 = new TCanvas("IntencityTest");
	hinttest->Draw();
	c3->SaveAs("testIntenc.pdf");

	TCanvas *c4 = new TCanvas("ThetaTest");
	c4->Divide(3,1);
	c4->cd(1);
	hthetatest->Draw();
	c4->cd(2);
	hphitest->Draw();
	c4->cd(3);
	hcostheta->Draw();
	c4->SaveAs("testAngles.pdf");

	cout<<"well done! "<<endl;
	return 0;
}