#include <iostream>
#include <strstream>
#include <math.h>

#include "pythia_if.h"
#include "TROOT.h"
#include "TH1.h"
#include "TH2.h"
#include "TFile.h"
#include "TNtuple.h"
#include "TRandom.h"

#include <stdio.h>

using namespace std;

TH1F *hall, *hpip, *hpim, *hgam, *hKp, *hKm, *hEt, *hp, *hap, *hn, *han;
TH1F *hall2, *hpip2, *hpim2, *hgam2, *hKp2, *hKm2, *hp2, *hap2, *hn2, *han2;
TH1F *hypi, *hetapi;
TH1F *heta2, *hEt2;
TH1F *hetaEt, *hetaEt2;
TH1F *he, *hepi, *heeta;
TH1F *hec_ally, *hec, *hec2;
TH1F *hyD0, *hyDp, *hyDs,  *hyDst0, *hyDstp, *hyDsts ;
TH1F *hptD0, *hptDp, *hptDs,  *hptDst0, *hptDstp, *hptDsts ;
TH1F *hx1,*hx2;
TH2F *hx1x2, *hx1x2e, *hx1pte, *hx2pte, *hxpte;
TH1F *hee, *hee_like,  *hee_c, *hee_like_c;
TH1F *hee2,*hee2like, *hee2_c, *hee2like_c;
TH1F *heb_ally, *heb, *hee_b, *hee2_b, *hpt;
TH2F *hh;
TNtuple *C1;

int Ne,Ne5,Ne8,Ne10,Ne13,Ne15,Ne18,Ne20,Ne25;


void calc_mass(int i, int j, float &eta_i, float &eta_j, float &mass) {
  float pxi = pythia::Px(i);
  float pyi = pythia::Py(i);
  float pzi = pythia::Pz(i);
  float Ei  = pythia::E(i);
  float pi  = sqrt(pxi*pxi+pyi*pyi+pzi*pzi);
  eta_i = 0.5*log((pi+pzi)/(pi-pzi));

  float pxj = pythia::Px(j);
  float pyj = pythia::Py(j);
  float pzj = pythia::Pz(j);
  float Ej  = pythia::E(j);
  float pj  = sqrt(pxj*pxj+pyj*pyj+pzj*pzj);
  eta_j = 0.5*log((pj+pzj)/(pj-pzj));

  float px = pxi + pxj;
  float py = pyi + pyj;
  float pz = pzi + pzj;
  float E  = Ei  + Ej;

  mass = sqrt(E*E - px*px - py*py - pz*pz);
}
//----------------------------------------------------------------------------------

void fill_histo(void) {
    
  float pi = 3.14159;
  float theta, ptmu;
  //  float Etsum = 0.0;
  //float Etsum2 = 0.0;
  //const float Mp = 0.938;
  //const float Mn = 0.939;

  //float x1 = pypars_.PARI[32];
  //float x2 = pypars_.PARI[33];

  //int nc = 0;
                                           // loop over all records
  //cout<<" event xx, #records: "<<pythia::Nrecord()<<endl;
  for(int i=1; i <= pythia::Nrecord(); i++) {
    int status = pythia::Status(i);
    int abs_pid = abs(pythia::Pid(i));
                                           // if it is a D_anything
    if(abs_pid == pythia::D0   || abs_pid == pythia::Dp   || abs_pid == pythia::Ds || 
       abs_pid == pythia::Dst0 || abs_pid == pythia::Dstp || abs_pid == pythia::Dsts) {
      float px = pythia::Px(i);
      float py = pythia::Py(i);
      float pz = pythia::Pz(i);
      float pid = pythia::Pid(i);

      //float E = pythia::E(i);
      //float y  = 0.5*log((E+pz)/(E-pz));
      //float pt = sqrt(px*px + py*py);

      int parent2 = pythia::Parent(i);
      int parent_pid2 = pythia::Pid(parent2);
                                                // parent id = charm (=4) 
      while (abs(parent_pid2)!=4 && abs(parent_pid2)!=0 ) {
        parent2 = pythia::Parent(parent2);
        parent_pid2 = pythia::Pid(parent2);
        //cout<<parent_pid2<<" ";
      }
      if (parent_pid2 ==0 ) continue;         // parent = ??? - quit
      //cout<<parent_pid2<<endl;
      float pxparent  = pythia::Px(parent2);
      float pyparent  = pythia::Py(parent2);
      float pzparent  = pythia::Pz(parent2);
      float pidparent = pythia::Pid(parent2);

      //float Eparent = pythia::E(parent2); 
      //float yparent  = 0.5*log((Eparent+pzparent)/(Eparent-pzparent));
      //float ptparent = sqrt(pxparent*pxparent + pyparent*pyparent);
      int written = 0;
                                                // loop over D decay products
      for (int j = pythia::FirstChild(i); j<=pythia::LastChild(i);j++) {
	                                        // find a mu in the decay
        if (abs(pythia::Pid(j)) == pythia::mu ) {

//C1 = new TNtuple("C1","charm ntuple","
//               px_parent:py_parent:pz_parent:pid_parent:
//               px_d:py_d:pz_d:pid_d:
//               px_mu:py_mu:pz_mu:pid_mu:
//               emu:vxmu:vymu:vzmu:status");

	  float arr[17];
          int i=0;
          arr[i++] = pxparent;      // the c, cbar
          arr[i++] = pyparent;
          arr[i++] = pzparent;
          arr[i++] = pidparent;

          arr[i++] = px;            // the D+- D0, D* etc
          arr[i++] = py;
          arr[i++] = pz;
          arr[i++] = pid;

	  arr[i++] = pythia::Px(j);   // the mu+-
	  arr[i++] = pythia::Py(j);
	  arr[i++] = pythia::Pz(j);
          arr[i++] = pythia::Pid(j);

	  arr[i++] = pythia::E(j);	  		       
          arr[i++] = pythia::Vx(j);    // vertex
          arr[i++] = pythia::Vy(j);
          arr[i++] = pythia::Vz(j);
          arr[i++] = float(status);

          float pxmu = pythia::Px(j);
          float pymu = pythia::Py(j);
          float pzmu = pythia::Pz(j);

	  ptmu = sqrt(pxmu*pxmu + pymu*pymu);

          theta = atan2(ptmu,pzmu)*180.0/pi;

          if (theta>6.0 && theta<42.0) { 
            C1->Fill(arr);
	  }
	  written = 1;
	}
      }

      //if (written ==0) C1->Fill(ptparent,yparent,parent_pid2,pt,y,pythia::Pid(i),0,0,0,status);
    }           // end if any D
  }            // end loop over records
}             // end void fill_histo

  void ask_user(float &root_s, int &nev, char *FileName, float *pkt) {
    root_s = 200.0;
    //cout << " enter root_s: ";
    //cin >> root_s;
  cout << " enter number of events: ";
  cin >> nev;
  *pkt = 1.5;
  //cout << " enter intrinsic kt (negative use default): ";
  //cin >> *pkt;
  cout << " enter root file name :";
  cin >> FileName;
  cout << "root_s = "<<root_s<<" nev = "<< nev << " File = "<< FileName<<endl;
}
//------------------------------------------------------------------------------
int main() {
TROOT simple("ppjets","Store PYTHIA results in ROOT");

float root_s;
int nev;
char FileName[100];
float kt;

ask_user(root_s,nev,FileName,&kt);
 cout << "=========================== kt set to " << kt <<endl;
//Select HQ production
pythia::pygive("msel=4");       //HQ production  c-cbar

// Use PDFLIB
pythia::pygive("mstp(52) = 2");   // Use PDFLIB
cout << "Select PDF"<<endl;
//cout << " 1--> CTEQ5L   (4046)"<<endl;
//cout << " 2--> CTEQ4L   (4032)"<<endl;
//cout << " 3--> GRV94LO  (5005)"<<endl;
//cout << " 4--> GRV98LO  (5012)"<<endl;
//cout << " 5--> MRST(c-g)(3072)"<<endl;
int ipdf;
//cin>> ipdf;
 ipdf = 1;
if(ipdf==1) {
pythia::pygive("mstp(51) = 4046");// PDFset is CTEQ5LO
} else if(ipdf==2) {
pythia::pygive("mstp(51) = 4032");// PDFset is CTEQ4LO
} else if(ipdf==3) {
pythia::pygive("mstp(51) = 5005");// PDFset is GRV94LO
} else if(ipdf==4) {
pythia::pygive("mstp(51) = 5012");// PDFset is GRV98LO
} else if(ipdf==5) {
pythia::pygive("mstp(51) = 3072");// PDFset is MRST(c-g)
} else {//default is CTEQ5L
pythia::pygive("mstp(51) = 4046");// PDFset is CTEQ5LO
}
// set primordal <kt>
if(kt > 0) 
{
cout << "set kt value"<<endl;
pythia::pygive("mstp(91) = 1"); // switch for primordal kt. Default is 1.
// 0-->no kt, 1-->gaussian, 2-->exponetial
pythia::pygive("parp(91) = ",kt);// gaussian width of <kt> (GeV/c)
// default is 0.44. Note <kt**2> = PARJ(91)**2.
pythia::pygive("parp(93)=5.0"); // <kt> cut off limit (GeV/c)
}
else
{
    cout << "set kt value"<<endl;
    pythia::pygive("mstp(91) = 0"); // switch for primordal kt. Default is 1.
    // 0-->no kt, 1-->gaussian, 2-->exponential
    pythia::pygive("parp(91) = 0");// gaussian width of <kt> (GeV/c)
    // default is 0.44. Note <kt**2> = PARJ(91)**2.
    pythia::pygive("parp(93)=5.0"); // <kt> cut off limit (GeV/c)
    pythia::pygive("mstj(22)=2.0"); // only decay prompts, if z<10mm (no pi, k decay)
}

// change charm mass to 1.25 GeV
// This is motivated by the recent value of charm mass...
// PDG value is, for example, Mc(Mc)=1.25+-0.1 GeV
//
// Ralf's value is 1.35 <--> K=5.2
// if mc=1.25 --> K=3.5
// if mc=1.5  --> K=8
// if mc=1.15 --> K=?
//
int charmid = pythia::pycomp(4);
cout << "charm id="<<charmid<<endl;
pythia::pygive("pmas(4,1)=1.25");

// set K-factor
pythia::pygive("parp(31)=1.0");
pythia::pygive("mstp(33)=1");

// set Q^2 choice
pythia::pygive("mstp(32)=4");  // scale is at Q^2 = root(s^)
  
// derive a random number seed from the current time:
time_t curr=time(0);                // get current time, = large integer
cout << "current time is: " << curr <<endl;
int int_time = int(curr);
char string[50];                    // holds command string
sprintf(string,"mrpy(1)=,%d",int_time); // compose command string
pythia::pygive(string);             // issue command to pythia

//chose ramdom sequence
//cout<<"select random sequence (1 to 9)"<<endl;
//int seed_choice;
//cin >> seed_choice;
//seed_choice = 1;
//cout << "seed_choice="<<seed_choice<<endl;

//if(seed_choice == 1)      pythia::pygive("mrpy(1)=1019390");
//else if(seed_choice == 2) pythia::pygive("mrpy(1)=1029390");
//else if(seed_choice == 3) pythia::pygive("mrpy(1)=1039390");
//else if(seed_choice == 4) pythia::pygive("mrpy(1)=1049390");
//else if(seed_choice == 5) pythia::pygive("mrpy(1)=1059390");
//else if(seed_choice == 6) pythia::pygive("mrpy(1)=1069390");
//else if(seed_choice == 7) pythia::pygive("mrpy(1)=1079390");
//else if(seed_choice == 8) pythia::pygive("mrpy(1)=1089390");
//else if(seed_choice == 9) pythia::pygive("mrpy(1)=1099390");
//else cout <<"use default"<<endl;

//
// set min. value of root_s_hat
//
cout << " select min. value of root_s hat"<<endl;
//cout << " 1 --> ckin(3) = 0."<<endl;
//cout << " 2 --> ckin(3) = 1."<<endl;
//cout << " 3 --> ckin(3) = 2."<<endl;
//cout << " 4 --> ckin(3) = 3."<<endl;
//cout << " 6 --> ckin(3) = 5."<<endl;
int ckin3_choice;
//cin >> ckin3_choice;
ckin3_choice = 1;
if(ckin3_choice == 1) { 
    pythia::pygive("ckin(3) = 0."); // min. value of root_s
} else if(ckin3_choice == 2) {
    pythia::pygive("ckin(3) = 1."); // min. value of root_s
} else if(ckin3_choice == 3) {
    pythia::pygive("ckin(3) = 2."); // min. value of root_s
} else if(ckin3_choice == 4) {
    pythia::pygive("ckin(3) = 3."); // min. value of root_s
} else if(ckin3_choice == 5) {
    pythia::pygive("ckin(3) = 4."); // min. value of root_s
} else if(ckin3_choice == 6) {
    pythia::pygive("ckin(3) = 5."); // min. value of root_s
} else {
    pythia::pygive("ckin(3) = 0."); // min. value of root_s
}

pythia::pyinit("CMS","p","p",root_s);

TFile *tfile = new TFile(FileName,"recreate","pythia Et");

C1 = new TNtuple("C1","charm ntuple","px_parent:py_parent:pz_parent:pid_parent:px_d:py_d:pz_d:pid_d:px_mu:py_mu:pz_mu:pid_mu:emu:vxmu:vymu:vzmu:status");

cout << "now start event loop for "<<nev<<" events:"<<endl;
for(int event=0; event<nev; event++) {
  if (event%1000 == 0) cout<<"At event "<<event<<", "<<100*event/nev<<"%"<<endl;

  pythia::pyevnt();
  if(event<=1) {
    //pythia::pylist(1);   // list full pythia history for the first generated event
    //pythia::pylist(12);  // list of all decay modes for all defined particles
  }
  fill_histo();
}
pythia::pystat(1);

tfile->Write();
tfile->Close();
}