// CHLEP stuff
#include <G4SystemOfUnits.hh>
#include <G4PhysicsOrderedFreeVector.hh>
#include <G4PhysicalConstants.hh>

#include <vector>
#include <string>
#include <sstream>

#include <RAT/Log.hh>
#include <RAT/DB.hh>
#include <RAT/Processor.hh>
#include <RAT/CherenkovSourcePDF.hh>
#include <RAT/FitterPMT.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/DS/Run.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DU/PMTInfo.hh>
#include <RAT/DU/LightPathCalculator.hh>
#include <RAT/PhysicsUtil.hh>

using namespace RAT;
using namespace RAT::PDFs;
using namespace RAT::DS;

using std::stringstream;
using std::vector;
using CLHEP::nm;
using CLHEP::c_light;

void CherenkovSourcePDF::Initialise( const std::string& param)
{
    fIndex=param;
}

void CherenkovSourcePDF::BeginOfRun( DS::Run&)
{
    DB* db = DB::Get();
    RAT::DBLinkPtr timingTable =  db->GetLink("CHERENKOV_SOURCE_PDF",fIndex);

    vector<double> times, probability;
    // Get table of hit probability vs. time residual from RATDB inputs
    try {
        times = timingTable->GetDArray( "time" );
        probability = timingTable->GetDArray( "probability" );
    } catch ( RAT::DBNotFoundError& e) {
        RAT::Log::Die(
                "CherenkovSourcePDF::Initialise: Cannot find database table" );
    }

    if(times.size() != probability.size()) {
        stringstream msg;
        msg << "\nCherenkovSourcePDF::Initialise: fatal error\n"
            << "Time and probability arrays have different sizes:  "
            << times.size() <<" "<< probability.size() << "\n\n";
        RAT::Log::Die( msg.str());
    }

    // Create fProbability table from times and probability vector arrays
    fProbability = new G4PhysicsOrderedFreeVector(
                            &(times[0]), &(probability[0]), times.size() );

    fLightPath = DU::Utility::Get()->GetLightPathCalculator();
    // Typical wavelength for photons from Cherenkov source with UVA acrylic:
    const double energy = RAT::util::WavelengthToEnergy( 455.0 * nm ); // 455 nm in MeV
    fIdxScint = fLightPath.GetInnerAVRI(energy);
    fIdxAV    = fLightPath.GetAVRI(energy);
    fIdxWater = fLightPath.GetWaterRI(energy);

    return;
}

double CherenkovSourcePDF::GetProbability( const FitterPMT& pmt, const FitVertex& vertex )
{
    TVector3 startPos;
    double eventTime;
    try {
        startPos = vertex.GetPosition();
        eventTime = vertex.GetTime();
    } catch( FitVertex::NoValueError& error ) {
        warn << "\nCherenkovSourcePDF::GetProbability: " <<
                "Vertex position or time not available -- " <<
                error.what() << newline;
        warn << "    Returning hit probability of zero"<< newline;
        return 0.0;
    }

    const DU::PMTInfo& pmtInfo = DU::Utility::Get()->GetPMTInfo();
    TVector3 pmtPos = pmtInfo.GetPosition( pmt.GetID() );
    double transitTime=-1.;

    const double energy = RAT::util::WavelengthToEnergy( 455.0 * nm ); //default wavelength value 455 [nm] in MeV

    // Calculate distInInnerAV, distInAV, and distInWater given
    // startPos and the PMT position, taking refractions into account
    fLightPath.CalcByPosition( startPos, pmtInfo.GetPosition( pmt.GetID() ),
                              energy, 200.0 );

    double distInInnerAV = fLightPath.GetDistInInnerAV();
    double distInAV = fLightPath.GetDistInAV();
    double distInWater = fLightPath.GetDistInWater();

    if( fLightPath.GetTIR() ) {  // total internal reflection encountered
        transitTime = 100.;
    }else{
        transitTime = (fIdxScint*distInInnerAV + fIdxAV*distInAV +
                                        fIdxWater*distInWater)/c_light;
    }

    double tResid = pmt.GetTime() - transitTime - eventTime;
    double prob = fProbability->Value( tResid);

    // "Punish" if way out of range
    if( tResid < -100.0 ){
        prob = exp( 0.01*(tResid + 100.0) )*prob;
    }else if( tResid > 300.0 ) {
        prob = exp( 0.01*(300.0 - tResid) )*prob;
    }

    return prob;
}