#include <RAT/EnergyPromptLookup.hh>
#include <RAT/DB.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/Log.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DU/LightPathCalculator.hh>
#include <RAT/DU/GroupVelocity.hh>
#include <RAT/DU/PMTInfo.hh>
#include <RAT/DU/ChanHWStatus.hh>

#include <RAT/PMTSelector.hh>
#include <RAT/PMTSelectorFactory.hh>
#include <RAT/FitterPMT.hh>
#include <RAT/SelectorMethod.hh>

#include <cmath>
#include <numeric>
#include <algorithm>
#include <iterator>
#include <functional>
#include <TVector3.h>
#include <TSpline.h>

using namespace RAT;
using namespace RAT::Methods;
using namespace RAT::DS;
using namespace std;
using std::isfinite;

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

void
EnergyPromptLookup::BeginOfRun( DS::Run& run )
{
  DB *db = DB::Get();
  string index = fIndex;
  if( fIndex.empty() )
    {
      try{ index = db->GetLink("GEO","inner_av")->GetS("material"); }
      catch( DBNotFoundError& e ) { /* Will use defaults instead */ }
    }
  // Get default parameters
  DBLinkPtr dbEPLLink = db->GetLink("FIT_ENERGY_PROMPT_LOOKUP");
  DBLinkGroup grp=db->GetLinkGroup("FIT_ENERGY_PROMPT_LOOKUP");
  DBLinkGroup::iterator it;
  // Check for material specific parameters
  // If not, keep default values
  for(it=grp.begin(); it != grp.end(); ++it){
    if(index == it->first){
      dbEPLLink = db->GetLink("FIT_ENERGY_PROMPT_LOOKUP", index);
      break;
    }
  }

  fMaterial = dbEPLLink->GetS( "index" );
  fRBins = dbEPLLink->GetD( "r_bins" );
  fCosThetaBins = dbEPLLink->GetD( "costheta_bins" );
  fRMax = dbEPLLink->GetD( "r_max" );
  fPromptWindow = dbEPLLink->GetDArray( "prompt_window" );
  fWorkingPMTs = dbEPLLink->GetD( "working_pmts" );
  fScalingFactor = dbEPLLink->GetDArray( "scaling_factor" );
  fMeVValues = dbEPLLink->GetDArray( "mev_values" );
  fNhitValues = dbEPLLink->GetDArray( "nhit_values" );

  fTimeResidualCut = PMTSelectors::PMTSelectorFactory::Get()->GetPMTSelector( "timeResidualCut" );
  fPMTCalSelector = PMTSelectors::PMTSelectorFactory::Get()->GetPMTSelector( "PMTCalSelector" );
  fTimeResidualCut->BeginOfRun( run );
}

void
EnergyPromptLookup::EndOfRun( DS::Run& )
{
  fPromptWindow.clear();
  fScalingFactor.clear();
  fMeVValues.clear();
  fNhitValues.clear();
}

void
EnergyPromptLookup::DefaultSeed()
{
  // Radius = 0.0, direction = -1z, time = 0.0
  DS::FitVertex vertex;
  vertex.SetPosition( TVector3( 0.0, 0.0, 0.0 ), false, true );
  vertex.SetDirection( TVector3( 0.0, 0.0, -1.0 ), false, true );
  vertex.SetTime( 0.0, false, true );
  fSeedResult.SetVertex( 0, vertex );
}

DS::FitResult
EnergyPromptLookup::GetBestFit()
{
  fFitResult.Reset();
  if( fPMTData.empty() )
    return fFitResult;
  CopySeedToResult();

  DS::FitVertex fitVertex = fFitResult.GetVertex( 0 );

  // Get number of working PMTs
  const DU::PMTInfo& pmtInfo = DU::Utility::Get()->GetPMTInfo();
  unsigned int numChannels = pmtInfo.GetCount();
  const DU::PMTCalStatus& PMTCalStatus = DU::Utility::Get()->GetPMTCalStatus();
  double totalActiveChannels = 0;
  for( size_t lcn = 0; lcn < numChannels; lcn++ ){
    if( pmtInfo.GetType( lcn ) != DU::PMTInfo::NORMAL &&
        pmtInfo.GetType( lcn ) != DU::PMTInfo::HQE)
      continue;
    double fractionGoodCells = PMTCalStatus.GetChannelStatus( static_cast<int>(lcn) );
    totalActiveChannels += fractionGoodCells;
  }

  // Get prompt calibrated PMTs
  fTimeResidualCut->SetD( "lowLimit", fPromptWindow[0] );
  fTimeResidualCut->SetD( "highLimit", fPromptWindow[1] );
  fTimeResidualCut->SetI( "useGroup", 1 );
  const vector<FitterPMT> promptPMTs = fTimeResidualCut->GetSelectedPMTs( fPMTData, fitVertex );
  const vector<FitterPMT> promptCalPMTs = fPMTCalSelector->GetSelectedPMTs( promptPMTs, fitVertex );
  size_t promptCalNhits = promptCalPMTs.size();

  // Get position and direction scaling value
  // First get reconstructed event information from seed
  const TVector3 eventPosition = fitVertex.GetPosition();
  const TVector3 eventDirection = fitVertex.GetDirection();

  // Get event radius and costheta from event position vector
  Double_t radius = eventPosition.Mag();
  Double_t costheta = 0;
  if( fMaterial != "lightwater_sno" )
    costheta = eventPosition.CosTheta();
  else
    costheta = (eventPosition.Dot(eventDirection)/(radius*eventDirection.Mag()));

  // Calculate radius and costheta bin centres
  vector<double> radialValues;
  for( int bin=0; bin<fRBins; bin++ ){
    radialValues.push_back( static_cast<double>(bin)*fRMax/(fRBins-1) );
  }
  vector<double> costhetaValues;
  for( int bin=0; bin<fCosThetaBins; bin++ ){
    costhetaValues.push_back( -1.0 + (static_cast<double>(bin)+0.5)*2.0/fCosThetaBins );
  }

  // Get corresponding scaling factor from database
  double posdirscalefactor = Interpolate2D( radius, costheta, radialValues, costhetaValues, fScalingFactor );

  // Scale prompt nhits with scaling factor and number of working PMTs
  Double_t promptCalNhits_scaled = ( static_cast<double>(promptCalNhits) / posdirscalefactor ) * ( fWorkingPMTs / totalActiveChannels );

  // Convert nhits to mev
  double energy = Interpolate1D( promptCalNhits_scaled, fNhitValues, fMeVValues);


  // Check if result is valid
  bool valid = true;
  if( radius > fRMax ){
    debug << "EnergyPromptLookup::GetBestFit: Radius larger than maximum in lookup table." << newline;
    valid = false;
  }
  if( fMaterial == "lightwater_sno" && energy < 0.26 ){
    debug << "EnergyPromptLookup::GetBestFit: Energy less than Cerenkov threshold." << newline;
    valid = false;
  }
  if( energy < 0 || !isfinite( energy ) ){
    debug << "EnergyPromptLookup::GetBestFit: Energy negative, infinity, or nan. Setting result invalid" << newline;
    valid = false;
  }
  // Set fit result
  fitVertex.SetEnergy( energy, valid );
  fitVertex.SetEnergyErrors( 1.0 );
  fFitResult.SetVertex( 0, fitVertex );
  return fFitResult;
}


// 1D interpolation using TSpline3 method for selected values
Double_t
EnergyPromptLookup::Interpolate1D( const Double_t inValue,
                                   const vector<double> inTable,
                                   const vector<double> outTable )
{
  if( inTable.size() != outTable.size() ){
    debug << "EnergyRSP::Interpolate1D: Tables do not match in size" << newline;
  }

  // Define vector of points to use
  vector<int> points (inTable.size(), 0);

  // Find lowest inTable value greater than inValue
  int inTableEntry = 0;
  inTableEntry = distance( inTable.begin(), find_if( inTable.begin(), inTable.end(), bind2nd(greater<double>(),inValue) ) );

  // Set points to use
  for( int i=1; i<4; i++ ){
    if( inTableEntry - i >= 0 )
      points[inTableEntry - i] = 1;
  }
  for( int i=0; i<3; i++ ){
    if( inTableEntry + i < static_cast<int>(inTable.size()) )
      points[inTableEntry + i] = 1;
  }

  // Find total number of points
  UInt_t total_points = accumulate( points.begin(), points.end(), 0 );

  // Find first point
  UInt_t firstPoint = distance( points.begin(), find( points.begin(), points.end(), 1 ) );

  // Perform cubic spline
  vector<double> x(total_points, 0);
  vector<double> y(total_points, 0);
  for( unsigned int i=0; i<total_points; i++ ){
    x[i] = inTable[firstPoint+i];
    y[i] = outTable[firstPoint+i];
  }
  TSpline3 Spline("spline",&x[0],&y[0],total_points);
  return Spline.Eval(inValue);
}


// 2D interpolation using TSpline3 method for selected values
Double_t
EnergyPromptLookup::Interpolate2D( const Double_t inValue1,
                                   const Double_t inValue2,
                                   const vector<double> inTable1,
                                   const vector<double> inTable2,
                                   const vector<double> outTable )
{
  if( inTable1.size() * inTable2.size() != outTable.size() ){
    debug << "EnergyRSP::Interpolate2D: Tables do not match in size" << newline;
  }

  // Define vector of points to use
  vector<int> points (inTable2.size(), 0);

  // Find lowest inTable value greater than inValue
  int inTable2Entry = distance( inTable2.begin(), find_if( inTable2.begin(), inTable2.end(), bind2nd(greater<double>(),inValue2) ) );

  // Set points to use
  for( int i=1; i<4; i++ ){
    if( inTable2Entry - i >= 0 )
      points[inTable2Entry - i] = 1;
  }
  for( int i=0; i<3; i++ ){
    if( inTable2Entry + i < static_cast<int>(inTable2.size()) )
      points[inTable2Entry + i] = 1;
  }

  // Find total number of points
  UInt_t total_points = accumulate( points.begin(), points.end(), 0 );

  // Find first point
  UInt_t firstPoint = distance( points.begin(), find( points.begin(), points.end(), 1 ) );

  // Get shortened inTable2
  vector<double> inTable2_shortened;
  for( unsigned int i=0; i<inTable2.size(); i++ ){
    if( points[i] == 1 )
      inTable2_shortened.push_back( inTable2[i] );
  }

  // For each inTable2 coordinate interpolate across inTable1 coordinates to create 1D interpolated outTable
  // Then interpolate across the inTable2 coordinates
  vector<double> yValues;
  for( UInt_t j=firstPoint; j<(firstPoint+total_points); j++ ){
    vector<double> xValues;
    for( unsigned int i=0; i<inTable1.size(); i++ ){
      xValues.push_back( outTable[ (j * inTable1.size()) + i ] );
    }
    yValues.push_back( Interpolate1D(inValue1, inTable1, xValues) );
  }
  return Interpolate1D(inValue2, inTable2_shortened, yValues);
}