#include <G4Material.hh>
#include <G4MaterialPropertyVector.hh>
#include <G4PhysicsOrderedFreeVector.hh>

#include <RAT/Log.hh>
#include <RAT/DB.hh>
#include <RAT/PositionDirectionPDF.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/DU/EffectiveVelocity.hh>
using namespace RAT;
using namespace RAT::PDFs;
using namespace RAT::DS;

using namespace std;

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

void
PositionDirectionPDF::BeginOfRun( DS::Run& )
{
  DB *db = DB::Get();
  string material;
  if( fIndex.empty() == true )
    material = db->GetLink("GEO", "inner_av")->GetS("material");
  else
    material = fIndex;

  // Set PDF to lightwater_sno as default
  const string defaultMaterial = "lightwater_sno";
  RAT::DBLinkPtr angularDistTable = RAT::DB::Get()->GetLink( "FIT_DIR", defaultMaterial );
  RAT::DBLinkPtr promptTimingTable = RAT::DB::Get()->GetLink( "GV1D", defaultMaterial );

  fConcentratorRadius = RAT::DB::Get()->GetLink( "CONCENTRATOR", "cRAT" )->GetDArray("rho_coord").back();

  // then check for material specific PDF
  DBLinkGroup angGrp=db->GetLinkGroup("FIT_DIR");
  DBLinkGroup::iterator it;
  for(it=angGrp.begin(); it != angGrp.end(); ++it){
    if(material == it->first){
      angularDistTable = RAT::DB::Get()->GetLink("FIT_DIR",material);
      break;
    }
  }
  // then check for material specific PDF
  DBLinkGroup timeGrp=db->GetLinkGroup( "GV1D" );
  for( it = timeGrp.begin(); it != timeGrp.end(); ++it )
    {
      if( material == it->first )
        {
          promptTimingTable = RAT::DB::Get()->GetLink("GV1D",material);
          break;
        }
    }

  // load in angles, time residuals & their probabilities
  vector<double> angle, angleProbability;
  vector<double> times, timeProbability;
  try{
    angle = angularDistTable->GetDArray( "angle" );
    angleProbability = angularDistTable->GetDArray( "probability" );
    times = promptTimingTable->GetDArray( "time" );
    timeProbability = promptTimingTable->GetDArray( "probability" );
  }
  catch ( RAT::DBNotFoundError &e){
    RAT::Log::Die( "PositionDirectionPDF::Initialise: Cannot find PositionDirection Parameters" );
  }

  fAngularProbability = new G4PhysicsOrderedFreeVector( &(angle[0]), &(angleProbability[0]), angle.size() );
  fTimingProbability = new G4PhysicsOrderedFreeVector( &(times[0]), &(timeProbability[0]), times.size() );

  fLightPath = DU::Utility::Get()->GetLightPathCalculator();

  return;
}

void
PositionDirectionPDF::EndOfRun( DS::Run& )
{
  delete fAngularProbability;
  delete fTimingProbability;
}

double
PositionDirectionPDF::GetProbability( const FitterPMT& pmt,
                                      const FitVertex& vertex )
{
  TVector3 position, direction;
  double eventTime;
  try {
    position = vertex.GetPosition();
    direction = vertex.GetDirection();
    eventTime = vertex.GetTime();
  }
  catch( FitVertex::NoValueError& error ) {
    warn << "PositionDirectionPDF::GetProbability: This PDF does nothing without a proposed position, direction & time" << newline;
    return 0.0;
  }

  // Calculate angle of PMT to the direction vector and look up in PDF to get probability
  const TVector3 pmtPos = DU::Utility::Get()->GetPMTInfo().GetPosition( pmt.GetID() );
  TVector3 dpmt = pmtPos - position;
  const double theta = acos(direction.Unit().Dot(dpmt.Unit()));
  double angleProb = fAngularProbability->Value(theta);

  // Solid angle correction away from centre
  double omega = 1./2 * fConcentratorRadius*fConcentratorRadius / ( dpmt.Mag()*dpmt.Mag())*((dpmt.Unit()).Dot(pmtPos.Unit()) );
  angleProb = angleProb*omega;

  // Now get the prompt timing information
  fLightPath.CalcByPosition( position, DU::Utility::Get()->GetPMTInfo().GetPosition( pmt.GetID() ) );
  double distInInnerAV = fLightPath.GetDistInInnerAV();
  double distInAV = fLightPath.GetDistInAV();
  double distInWater = fLightPath.GetDistInWater();
  const double transitTime = DU::Utility::Get()->GetGroupVelocity().CalcByDistance( distInInnerAV, distInAV, distInWater );
  const double correctedTime = pmt.GetTime() - transitTime - eventTime;
  double timeProb = 0;

  if( correctedTime > -100.0 && correctedTime < 300.0 )
	timeProb = fTimingProbability->Value( correctedTime );

  // "Punish" if way out of range
  else if( correctedTime < -100.0 )
	timeProb = exp( correctedTime + 100.0 ) * fTimingProbability->Value( correctedTime );
  else
	timeProb = exp( 300.0 - correctedTime ) * fTimingProbability->Value( correctedTime );

  double prob = angleProb * timeProb;

  return prob;
}