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

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

#include <string>
using namespace std;

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

void
DirectionPDF::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
  string defaultMaterial = "lightwater_sno";
  RAT::DBLinkPtr angularDistTable = RAT::DB::Get()->GetLink( "FIT_DIR", defaultMaterial );

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

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

  // load in angles & probabilities
  vector<double> angle, probability;
  try{
    angle = angularDistTable->GetDArray( "angle" );
    probability = angularDistTable->GetDArray( "probability" );
  }
  catch ( RAT::DBNotFoundError &e){
    RAT::Log::Die( "DirectionPDF::Initialise: Cannot find Direction Parameters" );
  }

  fProbability = new G4PhysicsOrderedFreeVector( &(angle[0]), &(probability[0]), angle.size() );
}

void
DirectionPDF::EndOfRun( DS::Run& )
{
  delete fProbability;
}

double
DirectionPDF::GetProbability( const FitterPMT& pmt,
                              const FitVertex& vertex )
{
  TVector3 position, direction;
  try {
    position = vertex.GetPosition();
    direction = vertex.GetDirection();
  }
  catch( FitVertex::NoValueError& error ) {
    warn << "DirectionPDF::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 prob = fProbability->Value(theta);

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

  prob = prob*omega;

  return prob;
}