// main112.cc is a part of the PYTHIA event generator.
// Copyright (C) 2017 Torbjorn Sjostrand.

// Simple program Heavy Ion Collisions. 
// Modified by Andre Vieira da Silva in 2018-05-30

// This file contains the main class for event generation.
// Pythia: provide the main user interface to everything else.
#include "Pythia8/Pythia.h"

// You need to include this to get access to the HIInfo object for HeavyIons.
#include "Pythia8/HeavyIons.h"

// ROOT, for saving Pythia events as trees in a file.
#include "TTree.h"
#include "TFile.h"

//Library Math
#include "TMath.h"

#include "sstream"

//load Library
#include "MyParticle.cxx"
#include "MyEvent.cxx"

using namespace Pythia8;

    int main() {
        
        Pythia pythia;
       
        //define variables:
        Double_t px, py, pz, pt, E, eta, phi;
        int id;
        
        // Set up the ROOT TFile and TTree.
        TFile *file = TFile::Open("pythia_heavy_ion.root","recreate");
        
        //Define TTRee "Tree" with the variable of the event
        TTree *Tree = new TTree("Tree","Tree");
        MyEvent *event = new MyEvent;
        
        Tree->Branch("event", &event);
        
        // Setting the beams.
        pythia.readString("Beams:idA = 2212");        // proton
        pythia.readString("Beams:idB = 1000822080");  // Lead Ion
        pythia.readString("Beams:eA = 4000");         // Energy beam A
        pythia.readString("Beams:eB = 1570");         // Energy beam B
        pythia.readString("Beams:frameType = 2");     // LAB Frame
        
        // Initialize the Angantyr model to fit the total and semi-includive
        // cross sections in Pythia within some tolerance.
        pythia.readString("HeavyIon:SigFitErr = "
        "0.02,0.02,0.1,0.05,0.05,0.0,0.1,0.0");
        
        // These parameters are typicall suitable for sqrt(S_NN)= 5.02 TeV
        pythia.readString("HeavyIon:SigFitDefPar = "
        "17.24,2.15,0.33,0.0,0.0,0.0,0.0,0.0");
        
        // A simple genetic algorithm is run for 20 generations 
        // to fit the parameters.
        pythia.readString("HeavyIon:SigFitNGen = 20");
        
        
        // Initialise Pythia.
        pythia.init();
       
        // Number of events!
        int nEvents = 10;  
    
        // Loop over events.
        for ( int iEvent = 0; iEvent < nEvents; ++iEvent ) {
            
            // make sure the event object is empty
            event->Clear();
          
            if ( !pythia.next() ) continue;
            
                //Loop over the particles of the event
                for (int i = 0; i < pythia.event.size(); ++i) {
                                  
                    if ( pythia.event[i].isFinal()) {
                    
                        // Create a particle and set its properties
                        MyParticle *part = new MyParticle();
                        
                        px = pythia.event[i].px();
                        py = pythia.event[i].py();
                        pz = pythia.event[i].pz();
                        
                        part->SetPx(px);
                        part->SetPy(py);
                        part->SetPz(pz);
                        part->PrintInfo();
                        
                        // add this particle to the current event 
                        event->AddParticle(part);
                          
                    }
                    
                }
                //Fill the TTRee "Tree" with this event
                Tree->Fill();
        }
        
        // Print the screen (terminal) the information about the TTRee particle
        // Tree->Print(); 
        
        // Record Tree Object in the file
        Tree->Write();
        
        // The run is over, so we write out some statistics.
        // I moved the delete statement here.
        file->Close();
        delete file;
        
        cout<< " And we're done!" << endl;
        
        return 0;
    }