#include <RAT/MultiPathWaterDirection.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"

//#define DEBUG

using namespace RAT;
using namespace ROOT;

MultiPathWaterDirection::MultiPathWaterDirection() {
  // parameter definitions
  fPars.push_back(Parameter("theta",628,-CLHEP::pi,CLHEP::pi));
  fPars.push_back(Parameter("phi",942,-CLHEP::pi,2.0*CLHEP::pi));
}

void MultiPathWaterDirection::BeginOfRun(DS::Run&) {
#ifdef MULTIPATH_FIXED_SEED
  stepd = 0.0001;
#endif
  DB* db = DB::Get();
  DBLinkPtr dbLink = db->GetLink( "FIT_MULTIPATH_ANGLE" );
  std::vector<double> fPDFDataChAngle = dbLink->GetDArray( "sPDFAngle_ChFit" );
  double fAngleBinWidth = dbLink->GetD( "angle_bin_width" );
  double fAngleOffset = dbLink->GetD( "angle_offset" );

  //Drive correction parameters
  DBLinkPtr dbLinkDrive = db->GetLink( "FIT_MULTIPATH" );
  fDriveCorEnable = dbLinkDrive->GetZ( "driveCor_enable" );
  fDrivePositionScale = dbLinkDrive->GetD( "driveCor_p0" );
  fDriveDirectionScale = dbLinkDrive->GetD( "driveCor_p1" );

  fEntriesAngle = fPDFDataChAngle.size();
  for (int j = 0; j < fEntriesAngle; j++) {
    fPDFxChAngle.push_back(fAngleBinWidth * static_cast<double>(j+1) - fAngleOffset);
    fPDFChAngle.push_back(fPDFDataChAngle[j]);
  }
  for (int j = 0; j < fEntriesAngle - 1; j++) {
    fDerivativeChAngle.push_back((fPDFChAngle[j+1] - fPDFChAngle[j])/fAngleBinWidth);
  }

#ifdef DEBUG
  //Check the fitting pdfs and their derivatives for any event
  TFile f("pdfDerivatives_MultiPathWaterDirection.root","RECREATE");
  TH1F h3("h3","pdfAngleCh",200,1.0,200.0);
  TH1F h3x("h3x","pdfXAngleCh",200,1.0,200.0);
  TH1F dh3("dh3","DerivativeAngleCh",200,1.0,200.0);

  for(int j = 0; j < fEntriesAngle - 1 ; j++) {
    h3.SetBinContent(j+1, fPDFChAngle[j]);
    h3x.SetBinContent(j+1, fPDFxChAngle[j]);
    dh3.SetBinContent(j+1, fDerivativeChAngle[j]);
  }

  f.cd();
  h3.Write();
  h3x.Write();
  dh3.Write();
  f.Close();
#endif // DEBUG
}

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

std::vector<double> MultiPathWaterDirection::GetStart() {
  // Could set up a seed here, but I'm not using it for now (from aTime and aVertex)
#ifdef MULTIPATH_FIXED_SEED
  double ranPi = stepd;
  double ran2Pi = 2.0*stepd;
  stepd += 0.0001;
  if (stepd > CLHEP::pi) stepd = 0.0001;
#else
  double ranPi = CLHEP::RandFlat::shoot(0.0, CLHEP::pi);
  double ran2Pi = CLHEP::RandFlat::shoot(0.0, 2.0*CLHEP::pi);
#endif
  std::vector<double> v;
  v.push_back(ranPi); // theta
  v.push_back(ran2Pi); // phi
  return v;
}

void MultiPathWaterDirection::LAnddLdCosTheta_Ch(double cosTheta, double& L_Ch, double& dLdCosTheta_Ch) {
  int ibin = (int)floor( ( cosTheta + 1.0 ) * 100.0 );
  if (ibin > 0 && ibin < fEntriesAngle - 2) {
    double dAngle = cosTheta - fPDFxChAngle[ibin];
    L_Ch = (fPDFChAngle[ibin] + fDerivativeChAngle[ibin] * dAngle);
    if (dAngle < 0.005) {
      double u = 100.0 * dAngle + 0.02;
      dLdCosTheta_Ch = u*fDerivativeChAngle[ibin] + (1.0-u)*fDerivativeChAngle[ibin-1];
    } else {
      double u = 100.0 * dAngle - 0.02;
      dLdCosTheta_Ch = (1.0-u)*fDerivativeChAngle[ibin] + u*fDerivativeChAngle[ibin+1];
    }// end of the if statement if dAngle<0.005
  } else if (ibin <= 0) {
    L_Ch = fPDFChAngle[0];
    dLdCosTheta_Ch = fDerivativeChAngle[0];
  } else {
    L_Ch = fPDFChAngle[fEntriesAngle - 2];
    dLdCosTheta_Ch = fDerivativeChAngle[fEntriesAngle - 2];
  }
}

void MultiPathWaterDirection::SetSeed(const DS::FitResult& result) {
  DS::FitVertex vertex = result.GetVertex(0);
  TVector3 pos = vertex.GetPosition();
  fVertex.clear();
  fVertex.push_back(pos.X());
  fVertex.push_back(pos.Y());
  fVertex.push_back(pos.Z());
  fVertex.push_back(vertex.GetTime());
}

DS::FitResult MultiPathWaterDirection::MakeResult(const std::vector<double>& pars) {
  TVector3 v(fVertex[0], fVertex[1], fVertex[2]);
  TVector3 d;
  d.SetMagThetaPhi(1.0, pars[0], pars[1]);
  DS::FitVertex vertex;
  vertex.SetPosition(v);
  vertex.SetDirection(d);
  if( fDriveCorEnable ) {
    TVector3 vCor;
    vCor = fDrivePositionScale*v + fDriveDirectionScale*d;
    vertex.SetPosition(vCor);
  }
  vertex.SetTime(fVertex[3]);
  DS::FitResult result;
  result.SetVertex(0, vertex);
  return result;
}

bool MultiPathWaterDirection::ValidResult(const std::vector<double>&) {
  return true;
}

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

  double coszenith = cos(par[0]);
  double sinzenith = sin(par[0]);
  double cosazimuth = cos(par[1]);
  double sinazimuth = sin(par[1]);

  double dx = pmt[0] - fVertex[0];
  double dy = pmt[1] - fVertex[1];
  double dz = pmt[2] - fVertex[2];

  double dmag = sqrt(dx*dx + dy*dy + dz*dz);
  dx /= dmag;
  dy /= dmag;
  dz /= dmag;

  double dirx = sinzenith * cosazimuth;
  double diry = sinzenith * sinazimuth;
  double dirz = coszenith;

  double ddthetax = coszenith * cosazimuth;
  double ddthetay = coszenith * sinazimuth;
  double ddthetaz = - sinzenith;

  double ddphix = - sinzenith * sinazimuth;
  double ddphiy =   sinzenith * cosazimuth;
  double ddphiz = 0.0;

  double cosTheta = dirx*dx + diry*dy + dirz*dz;

  double L = 0.0;
  double dLdCosThetaCh = 0.0;
  LAnddLdCosTheta_Ch(cosTheta, L, dLdCosThetaCh);

  double factor = dLdCosThetaCh / L;
  diff.assign(2, 0.0);
  diff[0] = factor * (ddthetax*dx + ddthetay*dy + ddthetaz*dz);
  diff[1] = factor * (ddphix*dx + ddphiy*dy + ddphiz*dz);

  L = log(L);

#ifdef DEBUG
  debug << "AWD::CL pmt=(" << pmt[0] << "," << pmt[1] << "," << pmt[2] << ")\n";
  debug << "        costheta=" << cosTheta << " dL=" << dLdCosThetaCh << "\n";
  debug << "        L=" << L << " diff=(" << diff[0] << "," << diff[1] << ")\n";
#endif

  return L;
}