#include <RAT/PartialEnergy.hh>
#include <RAT/DB.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/Log.hh>
#include <Math/Vector3D.h>

#include <string>

using namespace std;
using namespace RAT;
using namespace RAT::Methods;

void PartialEnergy::BeginOfRun(DS::Run&)
{

  DB* db = DB::Get();
  string material = db->GetLink("GEO", "inner_av")->GetS("material_top");

  DBLinkPtr dbLink = db->GetLink("FIT_PARTIAL_ENERGY_SCALING", material);

  fSplitZ = db->GetLink( "GEO", "inner_av" )->GetD( "split_z" );
  fZConstants = dbLink->GetDArray("z_constants");
  fZPositions = dbLink->GetDArray("z_positions");
  fZPower = dbLink->GetD("z_power");
  fAbsScaling = dbLink->GetD("abs_scaling");
  fAVInnerRadius = db->GetLink( "SOLID", "acrylic_vessel_inner" )->GetD( "r_sphere" );

  if( fSplitZ < -1.0 * fAVInnerRadius )
    {
      info << "Inner AV: " << fAVInnerRadius << "\tsplit " << fSplitZ << newline;
      RAT::Log::Die("PartialEnergy::Cannot run with a split inner_av at heights lower than the negative AV radius");
    }

}

void PartialEnergy::DefaultSeed()
{

  DS::FitVertex vertex;
  vertex.SetPosition(TVector3( 0.0, 0.0, 0.0 ), false, true);
  vertex.SetPositivePositionError(ROOT::Math::XYZVectorF(fAVInnerRadius, fAVInnerRadius, fAVInnerRadius), false, true);
  vertex.SetNegativePositionError(ROOT::Math::XYZVectorF(fAVInnerRadius, fAVInnerRadius, fAVInnerRadius), false, true);
  vertex.SetEnergy(1.0, false, true);
  vertex.SetPositiveEnergyError(1.0, false, true);
  vertex.SetNegativeEnergyError(1.0, false, true);
  fSeedResult.SetVertex(0, vertex);

}

DS::FitResult PartialEnergy::GetBestFit()
{

  fFitResult.Reset();
  if( fPMTData.empty() )
    return fFitResult;
  CopySeedToResult();

  // Simply use the seed position to apply a scaling to the seed energy
  Double_t energy = fFitResult.GetVertex(0).GetEnergy();
  Double_t zPosition = fFitResult.GetVertex(0).GetPosition().Z();

  int zBin = -1;
  for(size_t iBin = 0; iBin < fZPositions.size()-1; iBin++)
    {
      if(zPosition < fZPositions[iBin+1] && zPosition >= fZPositions[iBin])
        {
          zBin = iBin;
          break;
        }
    }
  // Interpolate between bins (or extrapolate)
  bool valid = true; // Only valid if we're interpolating
  if(zBin == -1)
    {
      valid = false;
      if(zPosition < fZPositions[0])
        zBin = 0;
      else
        zBin = fZPositions.size() - 2;
    }

  double zConstant = (fZConstants[zBin+1] - fZConstants[zBin]) / (fZPositions[zBin+1] - fZPositions[zBin]) *
    (zPosition - fZPositions[zBin]) + fZConstants[zBin];

  // Scaling according to: c1 * c2 * ( zsplit + rAV ) ^ power
  // Where the constants c2 and power are dependent on the event z
  double scaling = fAbsScaling * zConstant * pow((fSplitZ + fAVInnerRadius), fZPower);

  // The scaling is the EXTRA factor of nhits (relative to full scint volume)
  // i.e. Nhit (zsplit, z) = Nhit (full) * (1 + scaling (zsplit, z))
  // so Energy(seed) = Energy(adjusted) * (1 + scaling)
  double scaledEnergy = energy / (1 + scaling);
  DS::FitVertex vertex = fFitResult.GetVertex(0);
  vertex.SetEnergy(scaledEnergy, valid);
  vertex.SetPositiveEnergyError(fFitResult.GetVertex(0).GetPositiveEnergyError());
  vertex.SetNegativeEnergyError(fFitResult.GetVertex(0).GetNegativeEnergyError());
  fFitResult.SetVertex(0, vertex);
  return fFitResult;

}