#include <RAT/EnergyLookup.hh>
#include <RAT/DB.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/Log.hh>
using namespace RAT;
using namespace RAT::Methods;
using namespace RAT::DS;

#include <math.h>
using namespace std;

#include <TVector3.h>

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

void
EnergyLookup::BeginOfRun( DS::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 (those of labppo_scintillator)
  DBLinkPtr dbELLink = db->GetLink("FIT_ENERGY_LOOKUP");
  DBLinkGroup grp=db->GetLinkGroup("FIT_ENERGY_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){
      dbELLink = db->GetLink("FIT_ENERGY_LOOKUP", index);
      break;
    }
  }

  fNhitVEnergyVRadius = dbELLink->GetDArray( "nhit_energy_radius" );
  fRadii = dbELLink->GetDArray( "radii" );
  fEnergies = dbELLink->GetDArray( "energies" );
}

void
EnergyLookup::EndOfRun( DS::Run& )
{
  fNhitVEnergyVRadius.clear();
  fRadii.clear();
  fEnergies.clear();
}

void
EnergyLookup::DefaultSeed()
{
  //fRadius = 0.0;
  DS::FitVertex vertex;
  vertex.SetPosition( TVector3( 0.0, 0.0, 0.0 ), false, true );
  fSeedResult.SetVertex( 0, vertex );
}

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

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

  double nhit = fPMTData.size();
  double radius = fSeedResult.GetVertex( 0 ).GetPosition().Mag();

  // Find the nhit radial bins
  size_t radialBin;
  for( radialBin = 1; radialBin < fRadii.size() - 1; radialBin++ )
       if( radius < fRadii[radialBin] )
         break;

  int radialLowBin = radialBin - 1;
  int radialHighBin = radialBin;

  // Now calculate new Nhit scaling.
  // First find current scaling
  size_t energyBin;
  for( energyBin = 1; energyBin < fEnergies.size() - 1; energyBin++ )
    {
      int nhitBin = radialLowBin * fEnergies.size() + energyBin;
      if( nhit < fNhitVEnergyVRadius[nhitBin] )
        break;
    }
  // Interpolate Nhit low and high
  int energyLowBin = energyBin - 1;
  int energyHighBin = energyBin;
  double nHitLow = fNhitVEnergyVRadius[ radialLowBin  * fEnergies.size() + energyLowBin ] + ( fNhitVEnergyVRadius[ radialHighBin  * fEnergies.size() + energyLowBin ] - fNhitVEnergyVRadius[ radialLowBin  * fEnergies.size() + energyLowBin ] ) / ( fRadii[radialHighBin] - fRadii[radialLowBin] ) * ( radius - fRadii[radialLowBin] );
  double nHitHigh = fNhitVEnergyVRadius[ radialLowBin  * fEnergies.size() + energyHighBin ] + ( fNhitVEnergyVRadius[ radialHighBin  * fEnergies.size() + energyHighBin ] - fNhitVEnergyVRadius[ radialLowBin  * fEnergies.size() + energyHighBin ] ) / ( fRadii[radialHighBin] - fRadii[radialLowBin] ) * ( radius - fRadii[radialLowBin] );
  // Finally interpolate the energy
  double energy = fEnergies[energyLowBin] + ( fEnergies[energyHighBin] - fEnergies[energyLowBin] ) / ( nHitHigh - nHitLow ) * ( nhit - nHitLow );

  bool valid = true;
  if( radius >= fRadii[fRadii.size()-1] )
    {
      warn << "EnergyLookup::GetBestFit radius larger than maximum in lookup table." << newline;
      valid = false;
    }
  if( nhit >= fNhitVEnergyVRadius[(radialLowBin+1)*fEnergies.size()-1] && nhit >= fNhitVEnergyVRadius[(radialHighBin+1)*fEnergies.size()-1] )
    {
      warn << "EnergyLookup::GetBestFit nhit larger than maximum in lookup table." << newline;
      valid = false;
    }
  if( energy < 0.0 )
    valid = false;

  fitVertex.SetEnergy( energy, valid );
  fitVertex.SetEnergyErrors( 1.0 );
  fFitResult.SetVertex( 0, fitVertex );
  return fFitResult;
}