////////////////////////////////////////////////////////////////////////
///
/// \class
///
/// \brief
///
/// \author David Auty auty@ualberta.ca
///
/// REVISION HISTORY: Jie Hu jhu9@ualberta.ca - modified to fit J Tseng's MultiPathFunction\n
///
/// \detail
///       To fit data in a partial scint-water fill geometry
////
////////////////////////////////////////////////////////////////////////

#include <RAT/MultiPathScint.hh>
#include <CLHEP/Random/RandFlat.h>
#include <CLHEP/Units/SystemOfUnits.h>
#include <CLHEP/Units/PhysicalConstants.h>
#include <RAT/DB.hh>
#include <RAT/DU/Utility.hh>
#include "TFile.h"
#include "TH1F.h"
#include "TMath.h"
#include <TF1.h>
//#define DEBUG

using namespace RAT;
using namespace ROOT;

MultiPathScint::MultiPathScint() {
  // parameter definitions
  fPars.push_back(Parameter("x",1678,-8390.0,8390.0));
  fPars.push_back(Parameter("y",1678,-8390.0,8390.0));
  fPars.push_back(Parameter("z",1678,-8390.0,8390.0));
  fPars.push_back(Parameter("t",400,-100.0,300.0));
}

void MultiPathScint::BeginOfRun(DS::Run&) {
#ifdef MULTIPATH_FIXED_SEED
  //step = 0.0001;
  step = 0.5001;
#endif
  DB* db = DB::Get();
  DBLinkPtr dbLink = db->GetLink( "FIT_MULTIPATH" );
  fPDFDataCh = dbLink->GetDArray( "sPDF_multipathfit_shift" );
  fPDFDataScint = dbLink->GetDArray( "sPDF_scint" );
  fScintRI = dbLink->GetD( "scint_RI" ); // effective grVelocity in scint
  fWaterRI = dbLink->GetD( "water_RI_tuned" ); // effective grVelocity in water, based on scint phase MC
  double fChBinWidth = dbLink->GetD( "time_bin_width" ); // time resolution in ns, set to 0.25 ns
  double fChOffset = dbLink->GetD( "time_offset" );
  fMaxPMTtRes = dbLink->GetD( "max_PMT_tRes" );
  fBoundary = dbLink->GetD( "boundary_tolerance" );
  fNeckpathEnable = dbLink->GetZ( "neckpath_enable" );
  fPSUPRadius = dbLink->GetD( "psup_rad" );
  fPSUPRadius2 = fPSUPRadius * fPSUPRadius;

  fZoff = dbLink->GetD( "av_offset_z" ); //z offset of AV to PSUP center
  fAVRadius = db->GetLink( "SOLID", "acrylic_vessel_inner" )->GetD( "r_sphere" ); // 6005.0, AV inner radius
  fAVRadiusOuter = db->GetLink( "SOLID", "acrylic_vessel_outer" )->GetD( "r_sphere" ); // 6060, AV outer radius
  fNeckRadius = db->GetLink( "SOLID", "acrylic_vessel_inner" )->GetD( "r_neck" ); // inner radius of neck = 730 mm
  fNeckRadiusOuter = db->GetLink( "SOLID", "acrylic_vessel_outer" )->GetD( "r_neck" ); // outer radius of neck = 785 mm
  fZneckLo = sqrt( fAVRadiusOuter*fAVRadiusOuter - fNeckRadiusOuter*fNeckRadiusOuter ) + fZoff; // bottom of the neck

  double fSpeedOfLight = CLHEP::c_light;
  fSpeedOfLightWater = fSpeedOfLight / fWaterRI;
  fSpeedOfLightScint = fSpeedOfLight / fScintRI;

  fEntriesTime = fPDFDataCh.size();// 1600, same for two pdfs
  double xtemp = fPDFDataCh[0];
  int j;
  for( j = fEntriesTime - 1; j >= 0; j-- ) {
    if( fPDFDataCh[j] == fPDFDataCh[0] ) {
      fPDFDataCh[j] = xtemp;
      xtemp -= 0.001;//0.1 in FitterAW
    }
  }

  xtemp = fPDFDataCh[ fEntriesTime-1 ];

  for( j = fEntriesTime - 1; fPDFDataCh[j] == xtemp; j-- );

  for( j++; j < fEntriesTime; j++ ) {
    fPDFDataCh[j] = xtemp;
    xtemp -= 0.001;//0.1 in FitterAW
  }

  for( j = 0; j < fEntriesTime; j++ ) {
    fPDFxCh.push_back(fChBinWidth * static_cast<double>(j + 1) - fChOffset);
    fPDFCh.push_back(fPDFDataCh[j]);
  }

  for( j = 0; j < fEntriesTime - 1; j++ )
    fDerivativeCh.push_back( ( fPDFCh[j + 1] - fPDFCh[j] ) / fChBinWidth );

  /// fill scint pdf
  double xtemp1 = fPDFDataScint[0];
  for( j = fEntriesTime-1; j >= 0; j-- ) {
    if( fPDFDataScint[j] == fPDFDataScint[0] ) {
      fPDFDataScint[j] = xtemp1;
      xtemp1 -= 0.001;
    }
  }

  xtemp1 = fPDFDataScint[ fEntriesTime-1 ];

  for( j = fEntriesTime-1; fPDFDataScint[j] == xtemp1; j-- );

  for( j++; j < fEntriesTime; j++ ) {
    fPDFDataScint[j] = xtemp1;
    xtemp1 -= 0.001;
  }

  for( j = 0; j < fEntriesTime; j++ ) {
    fPDFxScint.push_back(fChBinWidth * static_cast<double>(j + 1) - fChOffset);
    fPDFScint.push_back(fPDFDataScint[j]);
  }

  for( j = 0; j < fEntriesTime - 1; j++ )
    fDerivativeScint.push_back( ( fPDFScint[j + 1] - fPDFScint[j] ) / fChBinWidth );

#ifdef DEBUG
  //Check the fitting pdfs and their derivatives for any event
  TFile f("pdfDerivatives_MultiPathScint.root","RECREATE");
  TH1F h1("h1","pdfCherenkov",fEntriesTime, -100, 300);
  TH1F h1x("h1x","pdfXCh",fEntriesTime, -100, 300);
  TH1F dh1("dh1","DerivativeCherenkov",fEntriesTime,-100, 300);

  TH1F h2("h2","pdfScint",fEntriesTime, -100, 300);
  TH1F h2x("h2x","pdfXScint",fEntriesTime, -100, 300);
  TH1F dh2("dh2","DerivativeScint",fEntriesTime, -100, 300);

  for( j = 0; j < fEntriesTime; j++ ) {
    h1.SetBinContent(j+1, fPDFCh[j]);
    h1x.SetBinContent(j+1, fPDFxCh[j]);
    dh1.SetBinContent(j+1, fDerivativeCh[j]);
  }
  for( j = 0; j < fEntriesTime; j++ ) {
    h2.SetBinContent(j+1, fPDFScint[j]);
    h2x.SetBinContent(j+1, fPDFxScint[j]);
    dh2.SetBinContent(j+1, fDerivativeScint[j]);
  }

  f.cd();
  h1.Write();h1x.Write();dh1.Write();
  h2.Write();h2x.Write();dh2.Write();
  f.Close();
#endif // DEBUG
}

void MultiPathScint::Initialise(const std::string&) {
}

void MultiPathScint::SetSeed(const DS::FitResult&) {
}

DS::FitResult MultiPathScint::MakeResult(const std::vector<double>& pars) {
  TVector3 fittedVertex(pars[0], pars[1], pars[2]);
  DS::FitVertex vertex;
  vertex.SetPosition(fittedVertex);
  vertex.SetTime(pars[3]);
  DS::FitResult result;
  result.SetVertex(0, vertex);
  return result;
}

bool MultiPathScint::ValidResult(const std::vector<double>& pars) {
  double r2 = pars[0]*pars[0] + pars[1]*pars[1] + pars[2]*pars[2];
  return r2 < fPSUPRadius2;
}

std::vector<double> MultiPathScint::GetStart() {
  // Could set up a seed here, but I'm not using it for now (from aTime and aVertex)
#ifdef MULTIPATH_FIXED_SEED
  double ran0 = step;
  double ran1 = 2.0*step - 1.0;
  double ran2Pi = step*2.0*CLHEP::pi;
  double t = 100.0 + 200.0*step;
  step += 0.0001;
  if (step > 1.0) step = 0.0001;
#else
  double ran0 = CLHEP::RandFlat::shoot( 0.0, 1.0 );
  double ran1 = CLHEP::RandFlat::shoot( -1.0, 1.0 );
  double ran2Pi = CLHEP::RandFlat::shoot( 0.0, 2.0*CLHEP::pi );
  double t = CLHEP::RandFlat::shoot( 100.0, 300.0 );
#endif

  double r = pow(ran0, 1.0/3.0) * fPSUPRadius; // mm
  double costheta = ran1;
  double sintheta = sqrt(1.0 - costheta*costheta);
  std::vector<double> v;
  v.push_back(r*sintheta*cos(ran2Pi));
  v.push_back(r*sintheta*sin(ran2Pi));
  v.push_back(r*costheta);
  // in principle, MC and data could have different PMT hit time offset,
  // but in latest incarnation of MultiPathFitter code, this doesn't seem to be the case
  v.push_back(t);
  return v;
}

/// returns Likelihood and derivative for a given time of flight t.
void MultiPathScint::LAnddLdt(double& L,  double& dLdt ) {
  if( fRes > fMaxPMTtRes ) fRes = fMaxPMTtRes; // longer than valid time (late light, crosstalk?)
  if( fRes < -fMaxPMTtRes ) fRes = -fMaxPMTtRes; // eariler than valid time

  int ibin = (int) floor( ( fRes + 100.0 ) * 4.0 );
  if(ibin > 0 && ibin < fEntriesTime - 2) {
    double dt = fRes-fPDFxCh[ibin];
    L = (fPDFCh[ibin] + fDerivativeCh[ibin] * dt);
    if (dt < 0.125) {
      double u = 0.5 + dt * 4.0;
      dLdt = -u*fDerivativeCh[ibin] - (1.0-u)*fDerivativeCh[ibin-1];
    } else {
      double u = 4.0 * dt - 0.5;
      dLdt = -(1.0-u)*fDerivativeCh[ibin] - u*fDerivativeCh[ibin+1];
    }// end of the if statement if dt<0.125
  } else if (ibin <= 0) {
    L = fPDFCh[0];
    dLdt = -fDerivativeCh[0];
  } else {
    L = fPDFCh[fEntriesTime - 2];
    dLdt = -fDerivativeCh[fEntriesTime - 2];
  }
}

void MultiPathScint::LAnddLdtScint( double& L, double& dLdt ) {
  if ( fRes > fMaxPMTtRes ) fRes = fMaxPMTtRes; //max_PMT_tRes
  else if ( fRes < -fMaxPMTtRes ) fRes = -fMaxPMTtRes;
  int ibin = (int) floor( ( fRes + 100.0 ) * 4.0 );
  if(ibin > 0 && ibin < fEntriesTime - 2) {
    double dt = fRes - fPDFxScint[ibin];
    L = (fPDFScint[ibin] + fDerivativeScint[ibin] * dt);
    if (dt < 0.125) {
      double u = 0.5 + dt * 4.0;
      dLdt = -u*fDerivativeScint[ibin] - (1.0-u)*fDerivativeScint[ibin-1];
    } else {
      double u = 4.0 * dt - 0.5;
      dLdt = -(1.0-u)*fDerivativeScint[ibin] - u*fDerivativeScint[ibin+1];
    }// end of the if statement if dt<0.125
  } else if (ibin <= 0) {
    L = fPDFScint[0];
    dLdt = -fDerivativeScint[0];
  } else {
    L = fPDFScint[fEntriesTime - 2];
    dLdt = -fDerivativeScint[fEntriesTime - 2];
  }
}

double MultiPathScint::Calculate(const std::vector<double>& pmt,
                                 const std::vector<double>& par,
                                 std::vector<double>& diff) {

  double dx = pmt[0] - par[0];
  double dy = pmt[1] - par[1];
  double dz = pmt[2] - par[2];
  double dr = sqrt(dx*dx + dy*dy + dz*dz); // | pmtpos - X0|

  TVector3 fVertex( par[0], par[1], par[2] ); // X0 = (x0,y0,z0), trial vertex
  TVector3 fPMTpos( pmt[0], pmt[1], pmt[2] );
  /// following is to calculate time residual by: PMTtime - tof - trial time;
  // fRes = pmt[3] - tof - par[3];
  // in partial fitter, tof is calculated separately in different ways in different situations;
  // instead of fRes, (tDiff, tof) is passed to the likelihood calculation
  double tof = 0.0;
  double L = 0.0;
  double dLdt = 0.0, dLdx = 0.0, dLdy = 0.0, dLdz = 0.0;
  double fx = par[0]; // trial x0
  double fy = par[1]; // trial y0
  double fz = par[2]; // trial z0

  double distInScint = 0.0;
  /// evaluate line-AV sphere intersection first
  double scintpathInAV = 0.0;
  double udDotVtx = fx*dx/dr + fy*dy/dr + (fz-fZoff)*dz/dr; // (P-X0)/|P-X0| * (X0 - oAV)
  double vtx2 = fx*fx + fy*fy + (fz-fZoff)*(fz-fZoff); // (X0 - oAV)*(X0 - oAV)
  double rVertex = sqrt( fx*fx + fy*fy + (fz-fZoff)*(fz-fZoff) );
  double sqrVal = udDotVtx*udDotVtx - vtx2 + fAVRadiusOuter*fAVRadiusOuter;
  double a1 = 0.0, a2 = 0.0, aplus = 0.0, abig = 0.0, asmall = 0.0;

  if( sqrVal>0 ) { // line passes AV sphere
    /// find the line-sphere intersect points; a1, a2 are the path lengths btw vertex and intersection points
    a1 = -udDotVtx + sqrt(sqrVal);
    a2 = -udDotVtx - sqrt(sqrVal); // always a2<a1

    if( rVertex<fAVRadiusOuter ) { // vertex inside the AV, equivalent to a1*a2<0
      if ( a1*a2<0 ) {
        aplus = a1; // a1>0>a2
        scintpathInAV = aplus;
      }
    } // vertex in AV

    else { // rVertex>=fAVRadiusOuter, vertex in external
      if( a1>0 && a2>0) {
        abig = a1; // far intersection point
        asmall = a2; // near intersection point
        scintpathInAV = abig - asmall;
      } // ensure a1 and a2 are positive
    } // vertex in external
  } // pass through AV

  #ifdef DEBUG
  if(scintpathInNeck>0)
  {
    debug <<"scint path: "<<scintpathInAV<<" mm \n";
    debug <<"(x,y,z,t ) = ("<<fx<<", "<<fy<<", "<<fz<<") "<<par[3]<<"\n";
    debug <<"(pmtx,y,z,t) = pmt= (" << pmt[0] << "," << pmt[1] << "," << pmt[2]<<") "<<pmt[3]<<"\n";
  }
  #endif

  if( scintpathInAV<0 ) scintpathInAV = 0;

  distInScint = scintpathInAV;
  tof = ( dr - distInScint)/fSpeedOfLightWater + distInScint/fSpeedOfLightScint;
  fRes = pmt[3] - tof - par[3];
  bool checkInScint = ( rVertex<fAVRadiusOuter || ( fz>fZneckLo && fz<fPSUPRadius && fVertex.Perp()<fNeckRadiusOuter) );
  if(checkInScint) LAnddLdtScint(L, dLdt);
  else LAnddLdt(L, dLdt);
  double dLdt0 = dLdt/L;
  /// derivatives are simplified as the vertex moves a little bit in the opposite direction to the light path in water or scint
  double fSpeedOfLightEff = checkInScint? fSpeedOfLightScint:fSpeedOfLightWater;
  double factor = -dLdt0 / ( dr * fSpeedOfLightEff );
  dLdx = factor * dx;
  dLdy = factor * dy;
  dLdz = factor * dz;

  diff.assign(4, 0.0);
  diff[0] = dLdx;
  diff[1] = dLdy;
  diff[2] = dLdz;
  diff[3] = dLdt0;

  L = log(L);

#ifdef DEBUG
  debug << "AWP::CL pmt= (" << pmt[0] << "," << pmt[1] << "," << pmt[2] << ")\n";
  debug << "        fRes=" << fRes << " dL=" << dLdtCh << "\n";
  debug << "        L=" << L << " diff=(" << diff[0] << "," << diff[1] << "," << diff[2] << "," << diff[3] << ")\n";
#endif

  return L;
}