#include "StQGSMdstMaker.h"


ClassImp(StQGSMdstMaker)


//
// Set maximum file size to 1.9 GB (Root has a 2GB limit)
//
#define MAXFILESIZE 1900000000


//_________________
StQGSMdstMaker::StQGSMdstMaker(const Char_t *iFileName,
                               const Char_t *oFileName)
{
  std::cout << "[StQGSMdstMaker] Creating an instance... ";
  TDatime d;

  mOutFileName = oFileName;
  mInFileName  = iFileName;
  mTree = 0;
  mOutFile = 0;
  mInEventFile = 0;
  mInTrackFile = 0;
  mCompression = 9;
  mNEvents = 0;
  mStop = false;

  matrix  = new TMatrixTSym<double>(2); // matrix for calculation of transverse sphericity
  matrix2 = new TMatrixTSym<double>(2); // matrix for calculation of transverse sphericity 2

  std::cout << "[OK]" << std::endl;
}


//_________________
StQGSMdstMaker::~StQGSMdstMaker()
{
  /* nothing to do */
}


//_________________
Int_t StQGSMdstMaker::Init()
{
  std::cout << "[StQGSMdstMaker] ---=== Initializing maker ===---" << std::endl;
  mEvent = new StO97Event();
  
  //
  // Input file
  //
  std::cout << "[StQGSMdstMaker] Openning an input file... ";
  if (mInEventFile)   fclose(mInEventFile);
  if (mInTrackFile)   fclose(mInTrackFile);

  mInEventFile = fopen(Form("%s/B_MULT", mInFileName), "r");
  mInTrackFile = fopen(Form("%s/finalpr.data", mInFileName), "r");
  
  if (mInEventFile && mInTrackFile)
    std::cout << "[OK]" << std::endl;
  else
  {
    std::cout << "[FAIL]" << std::endl;
    return kStFatal;
  }
  
  //
  // Output file
  //
  std::cout << "[StQGSMdstMaker] Creating output file... ";
  if (mOutFile)
  {
    mOutFile->Close();
    mOutFile->Open(mOutFileName, "RECREATE");
  }
  else
    mOutFile = new TFile(mOutFileName, "RECREATE");

  if (mOutFile)
    std::cout << "[OK]" << std::endl;
  else
  {
    std::cout << "[FAIL]" << std::endl;
    return kStFatal;
  }
  SetCompressionlevel(mCompression); // default value from constructor
  
  //
  // TTree setup
  //
  std::cout << "[StQGSMdstMaker] Creating TTree... ";
  if (mTree)
    delete mTree;

  mTree = new TTree("StO97Dst", "StO97Dst");
  if (!mTree)
  {
    std::cout << "[FAIL]" << std::endl;
    return kStFatal;
  }
  
  mTree->SetMaxTreeSize(MAXFILESIZE);
  mTree->Branch("StO97Event", "StO97Event", &mEvent);
  mNBytes = 0;
  std::cout << "[OK]" << std::endl;
  //
  // Finalization of initialization :)
  //
  std::cout << "[StQGSMdstMaker] ---=== Initialization complite ===---" << std::endl;

  return StMaker::Init();
}


//_________________
void StQGSMdstMaker::SetCompressionlevel(Int_t comp)
{
  mCompression = comp;
  std::cout << "[StQGSMdstMaker] Compression level = " << mCompression << std::endl;
  mOutFile->SetCompressionLevel(mCompression);
}


//_________________
Int_t StQGSMdstMaker::Make()
{
  //
  // Read event
  //
  mNEvents++;
  mStop = !mStop;
  mStop = ReadEvent();
  if (mStop)
    return kStOk;

  //
  // Read tracks
  //
  mStop = ReadTracks();
  if (!mStop)
  {
    mEvent->SetTransverseSphericity(mSperp);
    mEvent->SetTransverseSphericity2(mSperp2);
    mNBytes += mTree->Fill();
  }

  return kStOk;
}

//_________________
Int_t StQGSMdstMaker::Finish()
{
  std::cout << "[StQGSMdstMaker] ---=== Finalization of maker ===---"
            << "[StQGSMdstMaker] total number of events = "
            << mNEvents << std::endl
            << "[StQGSMdstMaker] total bytes = "
            << mNBytes
            << std::endl;

  // close event file
  if (mInEventFile)
  {
    fclose(mInEventFile);
    mInEventFile = 0;
  }
  
  // close track file
  if (mInTrackFile)
  {
    fclose(mInTrackFile);
    mInTrackFile = 0;
  }

  // write tree to output file
  if (mTree)
  {
    if (mOutFile)
      mOutFile->Write();

    delete mTree;
    mTree = 0;
  }

  // close output file
  if (mOutFile)
  {
    mOutFile->Close();
    delete mOutFile;
    mOutFile = 0;
  }

  return kStOk;
}


//_________________
Bool_t StQGSMdstMaker::ReadEvent()
{
  Int_t eventNumber;
  Float_t impPar;
  Char_t buf[128];

  while (1)
  {
    Int_t ret = fscanf(mInEventFile, " %i %f %i",
                       &eventNumber, &impPar, &mNTracks);
    if (ret <= 0)               // if can't read a line
      fgets(buf, 127, mInEventFile); // skip it
    else
    {
      if (ret == 3) // 3 - total number of event parameters
      { 
        mEvent->Clear();
        mEvent->SetEventNumber(eventNumber);
        mEvent->SetImpactPar(impPar);
        return false;
      }
    }
    
    if (feof(mInEventFile)) // if EOF then return stop flag
      return true;
  }
}

//_________________
Bool_t StQGSMdstMaker::ReadTracks()
{
  Float_t px, py, pz;   // [GeV/c]
  Float_t energy;       // [GeV]
  Float_t mass;         // [GeV/c^2]
  Float_t xFreeze, yFreeze, zFreeze; // [fm]
  Float_t tFreeze;                   // [fm/c]
  Float_t xFormation, yFormation, zFormation; // [fm]
  Float_t tFormation;                         // [fm/c]
  Float_t tForm; /* [fm/c] this is life time of resonance (?)
                  * for primary particles this time equals 0
                  */
 
  Int_t ident; // particle code in ISAJET 7.79
  Int_t idiag; // Feynman diagram code
  Int_t iorig; // charge*|KP(see QGSM code, it is like orig)|*10^4 + IPAD (decayed particle number in event)
  Int_t ib, is, ich; // baryon, strangeness, charge

  TLorentzVector particle;

  Float_t pt_full  = 0.0;
  Float_t pt_full2 = 0.0;
  mSperp  = -999.;
  mSperp2 = -999.;

  matrix->Zero();
  matrix2->Zero();

  for (Int_t iTrack = 0; iTrack < mNTracks; iTrack++)
  {
    Int_t ret = fscanf(mInTrackFile,
                       " %f %f %f %f"    // tfr, xfr, yfr, zfr
                       " %f %f %f %f %f" // E, px, py, pz, mass
                       " %i %i %i %i %i" // ident, idiag, B, S, Q
                       " %f %f %f %f"    // tForm, xForm, yForm, zForm
                       " %i %f ",        // iorig and resonance life time
                       &tFreeze, &xFreeze, &yFreeze, &zFreeze,
                       &energy, &px, &py, &pz, &mass,
                       &ident, &idiag, &ib, &is, &ich,
                       &tFormation, &xFormation, &yFormation, &zFormation,
                       &iorig, &tForm);

    if (ret != 20)
      return true;
    else
    {
      StO97Track *track = mEvent->AddOscar97Track();
      track->SetId(iTrack);
      track->SetPdgId(ident);
      track->SetPx(px);
      track->SetPy(py);
      track->SetPz(pz);
      track->SetMass(mass);
      track->SetXfr(xFreeze);
      track->SetYfr(yFreeze);
      track->SetZfr(zFreeze);
      track->SetTfr(tFreeze);

      if (ich == 0)   continue;

      /*
       * Transverse sphericity calculation with following cuts:
       *   * |eta| < 0.5 && |eta| < 1.0
       *   * pT > 0.15 GeV/c
       */
      // setup particle vector
      particle.SetPxPyPzE(px, py, pz, energy);
      Float_t pt  = particle.Pt();

      if (pt > 0.15)
      {
        Float_t eta = particle.PseudoRapidity(); // it heals pt = 0 case
        
        if (TMath::Abs(eta) < 0.5)
        {
          (*matrix)(0, 0) += px*px/pt;
          (*matrix)(1, 1) += py*py/pt;
          (*matrix)(0, 1) += px*py/pt;
          (*matrix)(1, 0) += px*py/pt;

          pt_full += pt;
        }
        
        if (TMath::Abs(eta) < 1.0)
        {
          (*matrix2)(0, 0) += px*px/pt;
          (*matrix2)(1, 1) += py*py/pt;
          (*matrix2)(0, 1) += px*py/pt;
          (*matrix2)(1, 0) += px*py/pt;

          pt_full2 += pt;
        }
      }
    }
  }

  if (pt_full != 0.0)
  {
    // |eta| < 0.5 case
    *matrix  *= 1./pt_full;
    TMatrixDSymEigen death_machine(*matrix);
    TVectorD eigen = death_machine.GetEigenValues();
    mSperp = 2.*eigen.Min()/(eigen[0] + eigen[1]);

    // |eta| < 1.0 case
    *matrix2 *= 1./pt_full2;
    TMatrixDSymEigen death_machine2(*matrix2);
    TVectorD eigen2 = death_machine2.GetEigenValues();
    mSperp2 = 2.*eigen2.Min()/(eigen2[0] + eigen2[1]);
  }
  
  return false;
}