/*
  Some utility functions to convolute spectra and fragmentation functions
  used by FF_RAA_brick.C
  some of these functions are also contained in plot_spec.C

  Author: M. van Leeuwen, Utrecht University
*/

Float_t interpolate(TGraph *gr, Float_t x) {
  Int_t ip = 0;
  Double_t x2, y2;
  gr->GetPoint(ip,x2,y2);
  while (ip < gr->GetN() && x > x2) {
    ip++;
    gr->GetPoint(ip,x2,y2);
  }
  if (ip == 0 || ip == gr->GetN())
    return 0;
  Double_t x1, y1;
  gr->GetPoint(ip-1,x1,y1);
  return ((x-x1)*y2 + (x2-x)*y1) / (x2-x1);
}

TGraph *convolute_spectrum_inv(TGraph *parton_gr, TF1 *frag_f, Float_t min_pt_had = 2, Float_t max_pt_had = 100) {
  // parton spectrum is dN/pt/dpt !
  TGraph *pt_had_gr = new TGraph();

  const Int_t n_pt = 100;
  const Float_t dpt_had = (max_pt_had-min_pt_had)/n_pt;
  const Float_t dz = 0.025;
  const Float_t zmin = 0.05;
  const Float_t Qmax = 100; // maximum Q for KKP
  for (Int_t i_pt = 0; i_pt < n_pt; i_pt++) {
    Float_t pt_had = min_pt_had + i_pt*dpt_had;
    Float_t z = zmin;
    Double_t dn_tot = 0;
    while ( z <= 1 ) {
      Float_t pt_part = pt_had/z;
      Float_t Q = TMath::Min(pt_part,Qmax);
      frag_f->SetParameter(0,Q);

      // need 1/z^2 (PHENIX nucl-ex/0605039) for dN/ptdpt 
      //cout << "pt_had " << pt_had << " z  " << z << " pt_part " << pt_part 
      //   << " dN/pt/dpt_p " << interpolate(parton_gr,pt_part) << " ff " << frag_f->Eval(z) << endl;
      dn_tot += interpolate(parton_gr,pt_part) * frag_f->Eval(z) / z / z; 
      z += dz;
    }
    pt_had_gr->SetPoint(i_pt,pt_had,dn_tot);
  }
  return pt_had_gr;
}

TGraph *convolute_spectrum_inv(TGraph *parton_gr, TGraph *frag_gr, Float_t min_pt_had=2, Float_t max_pt_had=100) {
  // Assume parton spectrum is dN/dpt/pt !
  TGraph *pt_had_gr = new TGraph();
  const Int_t n_pt = 100;
  const Float_t dpt_had = (max_pt_had-min_pt_had)/n_pt;
  const Float_t dz = 0.05;
  const Float_t Qmax = 100; // maximum Q for KKP
  for (Int_t i_pt = 0; i_pt < n_pt; i_pt++) {
    Float_t pt_had = min_pt_had + i_pt*dpt_had;
    Float_t z = dz;
    Double_t dn_tot = 0;
    while ( z <= 1 ) {
      //cout << " z " << z << endl;
      Float_t pt_part = pt_had/z;
      Float_t Q = TMath::Min(pt_part,Qmax);
      frag_f->SetParameter(0,Q);
      
      // dpt,p = 1/z * dpt,h
      // need 1/z^2 (PHENIX nucl-ex/0605039) for dN/ptdpt 
      // cout << "pt_part " << pt_part << " ff " << frag_f->Eval(z) << endl;
      // cout << interpolate(pt_part_gr[i_type],pt_part) << endl;
      dn_tot += interpolate(parton_gr,pt_part) * interpolate(frag_gr,z) / z / z; 
      z += dz;
    }
    pt_had_gr->SetPoint(i_pt,pt_had,dn_tot);
  }
  return pt_had_gr;
}

TGraph *convolute_ff(Float_t E, TGraph *rad_gr, TGraph *edge_gr, TF1 *frag_f, Int_t np = 51, Float_t min_z = 0, Float_t max_z = 1) {
  // convolute over x = DeltaE/E

  Float_t dz = (max_z-min_z)/(np-1);
  
  Int_t np_x = 100;
  Float_t dx = 1./np_x;

  Double_t dum, p0, p1 = 0;
  edge_gr->GetPoint(0,dum,p0);
  if (edge_gr->GetN() > 1)
    edge_gr->GetPoint(1,dum,p1);

  TGraph * gr = new TGraph(np);
  for (Int_t i_pt = 0; i_pt < np; i_pt++ ) {
    Float_t pt = (min_z + i_pt*dz) * E;
    Float_t dndpt = 0;

    dndpt += p0 * frag_f->Eval(pt/E) * 1./E; // 1/E/(1-x) is dz/dpt and x = 0 (P(0))

    for (Int_t i_x = 0; i_x < np_x; i_x++) {
      Float_t x = (i_x + 0.5) * dx;

      Float_t z = pt/E/(1.-x);

      if (z <= 1)
	dndpt += interpolate(rad_gr, x) * frag_f->Eval(z) * 1./E/(1-x) * dx;
      // 1/E/(1-x) is dz/dpt
    }
    gr->SetPoint(i_pt, pt/E, dndpt*E);
  }
  return gr;
}