/*
  Full convolution of parton spectra, energy loss and fragmentation functions
  
  * Parton energy loss taken from AliQuenchingWeights (aliroot or standalone)
  * Fragmentation functions from KKP library
  * Input parton spectra from a textfile (generated using calc_xsec4.C)

  Uses utils.C and convolute.C

  Author: M. van Leeuwen, Utrecht University
*/


{
  gROOT->SetStyle("Plain");
  gSystem->Load("kkp/libkkp.so");
  gROOT->LoadMacro("utils.C");
  gROOT->LoadMacro("convolute.C");

  gStyle->SetTitleOffset(1.2,"Y");

  TCanvas *c1 = new TCanvas("c1","c1: FF",800,600);
  c1->Divide(2,1);
  TCanvas *c2 = new TCanvas("c2","c2: P(DE)",600,600);
  TCanvas *c3 = new TCanvas("c3","c3: pt-spectra",800,600);
  c3->Divide(2,1);

  const Float_t qhat = 3; // [GeV^2/fm]
  const Float_t L = 5; //2; // [fm]
  const Float_t alpha_s = 0.3;
  const Float_t hbarc = 0.197; //[GeVfm]

  const Float_t Eparton = 20;
  //const Int_t gluon_frag = 1;  // flag to include fragments from radiated gluon

  Char_t *type_label[2]={"quark","gluon"};
  Float_t kkp_type[2]={1,0}; // u, g   

  //Char_t *parton_fname="calc_xsec/cteq61_LO_spec_200GeV_4.txt";
  Char_t *parton_fname="calc_xsec/cteq61_LO_spec_2TeV76_4.txt";

  TF1 *frag_f = new TF1("frag_f",kkp_func,0,1,2);
  frag_f->SetParameter(0,Eparton); // Q for fragmentation

  AliQuenchingWeights qw;
  qw->InitMult();

  // Read parton spectra
  
  TGraph *pt_part_gr[2] = {0,0};
  
  TGraph *q_pt_gr = new TGraph;
  TGraph *g_pt_gr = new TGraph;
  pt_part_gr[1] = g_pt_gr;
  pt_part_gr[0] = q_pt_gr;
  ifstream fin(parton_fname);
  while (fin.good()) {
    Char_t dum[25];
    Double_t pt, qq, gg;
    fin >> pt >> qq >> gg;
    q_pt_gr->SetPoint(q_pt_gr->GetN(),pt,qq);
    g_pt_gr->SetPoint(g_pt_gr->GetN(),pt,gg);
  }

  c3->cd(1);
  gPad->SetLogy();
  h1 = gPad->DrawFrame(0,1e-10,20,1e-3);
  h1->SetXTitle("p_{T} (GeV)");
  h1->SetYTitle("dN/dp_{T}");
  c3->cd(2);
  h1 = gPad->DrawFrame(0,0,20,1);
  h1->SetXTitle("p_{T} (GeV)");
  h1->SetYTitle("R_{AA}");


  c1->cd(1);
  gPad->SetLogy();
  h1 = gPad->DrawFrame(0,1e-3,1,10);
  h1->SetXTitle("z");
  h1->SetYTitle("D(z)");

  c1->cd(2);
  h1 = gPad->DrawFrame(0,0,1,1.2);
  h1->SetXTitle("z");
  h1->SetYTitle("D_{AA}(z)/D_{pp}(z)");

  c2->cd();
  h1 = gPad->DrawFrame(0,0,1,2);
  h1->SetXTitle("x=#omega/E");
  h1->SetYTitle("dN/dx");

  TLatex *ltx = new TLatex();
  ltx->SetNDC();
  Char_t tmp_str[256];
  sprintf(tmp_str,"L = %.0f fm, E = %0.f GeV",L,Eparton);
  ltx->DrawLatex(0.4,0.83,tmp_str);

  Float_t omega_c = 0.5*qhat*L*L/hbarc;
  Float_t R = omega_c * L / hbarc;

  Double_t domega = 0.1;

  for (Int_t iparton = 0; iparton < 2; iparton++) {    
    frag_f->SetParameter(1,kkp_type[iparton]); // parton flavour 0=g 1=u    
    Int_t lstyle = iparton+1;
    Int_t lcol = 2*(iparton+1);
    frag_f->SetLineWidth(1);
    frag_f->SetLineStyle(lstyle);
    c1->cd(1);
    frag_f->DrawCopy("same");
  
    pt_had_gr = convolute_spectrum_inv(pt_part_gr[iparton], frag_f);
    pt_had_gr->SetLineStyle(lstyle);

    c3->cd(1);
    pt_had_gr->Draw("l");

    TGraph *bdmps_cont_gr = new TGraph();
    TGraph *bdmps_disc_gr = new TGraph();

    for (Int_t i=0; i < 600; i++) {
      Double_t cont, disc;
      Float_t omega = domega*(i+1);
      qw->CalcMult(iparton+1, R, omega/omega_c, cont, disc);
      if (i==0) 
	bdmps_disc_gr->SetPoint(0,0,disc);
      Float_t dpdomega = cont/omega_c;
      bdmps_cont_gr->SetPoint(i,omega/Eparton,dpdomega*Eparton);
    }

    c2->cd();
    bdmps_cont_gr->SetLineColor(lcol);
    bdmps_cont_gr->SetLineStyle(lstyle);
    bdmps_cont_gr->Draw("l");
    bdmps_disc_gr->SetMarkerColor(lcol);
    bdmps_disc_gr->SetMarkerStyle(20+iparton);
    bdmps_disc_gr->Draw("p");

    asw_ff_gr = convolute_ff(Eparton, bdmps_cont_gr, bdmps_disc_gr, frag_f,49,0.02,0.98);

    c1->cd(1);
    asw_ff_gr->SetLineColor(lcol);
    asw_ff_gr->Draw("l");

    c1->cd(2);
    asw_rat_gr = divide_graph(asw_ff_gr,frag_f);
    asw_rat_gr->SetLineColor(lcol);
    asw_rat_gr->SetLineStyle(lstyle);
    asw_rat_gr->SetLineWidth(2);
    asw_rat_gr->Draw("l");

    pt_had_asw_gr = convolute_spectrum_inv(pt_part_gr[iparton], asw_ff_gr);
    c3->cd(1);
    pt_had_asw_gr->SetLineWidth(lstyle);
    pt_had_asw_gr->SetLineColor(lcol);
    pt_had_asw_gr->Draw("l");
    
    c3->cd(2);
    raa_asw_gr = divide_graphs_int(pt_had_asw_gr,pt_had_gr);
    raa_asw_gr->SetLineWidth(2);
    raa_asw_gr->SetLineStyle(lstyle);
    raa_asw_gr->SetLineColor(lcol);
    raa_asw_gr->Draw("L");
  }
}