////////////////////////////////////////////////////////////////////////
///// \class RAT::Methods::QPartialEnergy
/////
///// \brief Energy estimator based on charge for the partial scint fill
/////
///// \author Yang Zhang <yzhang23@ualberta.ca>
////          Rohith Sekar rohithve@ualberta.ca
////
///// REVISION HISTORY:
/////  - 2019-09-07 : Yang Zhang - first instance
/////
///// \details This method estimates the energy based on charge for the
////   partial fill. Basically, it is just an estimatation function and
////   the lool-up table. The current fitter is valid up to 30 MeV.
//////////////////////////////////////////////////////////////////////////

#include <RAT/QPartialEnergy.hh>
#include <RAT/DB.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/Log.hh>
#include <RAT/FitterPMT.hh>
#include <Math/Vector3D.h>
#include <RAT/PMTSelector.hh>
#include <RAT/PMTSelectorFactory.hh>
#include <RAT/ChannelEfficiency.hh>


#include <string>

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

void
QPartialEnergy::Initialise( const std::string& param)
{
   //fIndex is material
   fIndex = param;
}

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

  DB* db = DB::Get();
  string index;
  if( fIndex.empty() )
  {
    try{
      index = db->GetLink("GEO","inner_av")->GetS("material");
    }
    catch(DBNotFoundError & ){
      try{
        // If we are in a partial-fill geometry, use the upper
        // material
        index = db->GetLink("GEO","inner_av")->GetS("material_top");
      }
      catch(DBNotFoundError & ) {
        Log::Die("QEnergy: inner_av material and material_top "
        "undefined. Try something like:\n"
        "/rat/db/set GEO[inner_av] material "
        "\"labppo_scintillator\" or "
        "/rat/db/set GEO[inner_av] material_top "
        "\"labppo_scintillator\"");
      }
    }
  }
  else
  {
    index = fIndex;
  }

  DBLinkPtr dbLink;
  try{
    (dbLink = db->GetLink("QPartialEnergy",index))->GetS("index");
  }
  catch(DBNotFoundError & ){
    warn << "Warning : QPartialEnergy::BeginOfRun: No full (QPartialEnergy) definition for " + index + " using defaults instead.\n";
    dbLink = db->GetLink("QPartialEnergy");
  }

  fQratio = dbLink->GetDArray( "Qratio_scint" );
  fzCorrectP = dbLink->GetDArray( "zCorrectP" );
  fEnergyP = dbLink->GetDArray( "EnergyP" );
  fEnergyErorrP = dbLink->GetDArray( "EnergyErorrP" );
  fSplitZ = db->GetLink( "GEO", "inner_av" )->GetD( "split_z" );
  fZoff = db->GetLink( "GEO", "av" )->GetDArray( "position" );
  fAVInnerRadius = db->GetLink( "SOLID", "acrylic_vessel_inner" )->GetD( "r_sphere" );
  fEnergy = dbLink->GetD( "EnergyCut" );
  fCoorChanEff = dbLink->GetD("chan_eff");
  fCoorActiveChannels = dbLink->GetD("active_channels");
  fQDummy = dbLink->GetD("QDummy");

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

  ApplyDetectorStateCorrection();
}

void QPartialEnergy::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);

}

void QPartialEnergy::ApplyDetectorStateCorrection(){
     const DU::PMTCalStatus& PMTCalStatus = DU::Utility::Get()->GetPMTCalStatus();
     const DU::PMTInfo& pmtInfo = DU::Utility::Get()->GetPMTInfo();
     const DU::ChanHWStatus& channelHardwareStatus = DU::Utility::Get()->GetChanHWStatus();

     //Get the fraction of active channels to correct for detector state
     int numChannels = DB::Get()->GetLink( "PMT_DQXX" )->GetIArray( "dqch" ).size();
     unsigned int totalActiveChannels = 0;
     for( size_t lcn = 0; lcn < numChannels; lcn++ ){
         if (channelHardwareStatus.IsEnabled()){
             if (!channelHardwareStatus.IsDAQEnabled(lcn))
                  continue;
             if( pmtInfo.GetType( lcn ) != DU::PMTInfo::NORMAL &&
                 pmtInfo.GetType( lcn ) != DU::PMTInfo::HQE)
             continue;
             double fractionGoodCells = PMTCalStatus.GetChannelStatus( static_cast<int>(lcn) );
             if( fractionGoodCells )
               totalActiveChannels++;
         }
       }
     double chanNumCorr = static_cast<double>(totalActiveChannels)/fCoorActiveChannels;

     //Get the channel efficiency correction
     double chanEff = ChannelEfficiency::Get()->GetAverageEfficiency();
     double chanEffCorr = chanEff/fCoorChanEff;
     //Apply both corrections
     fCorrt_NCh_Eff = chanNumCorr*chanEffCorr;
}

DS::FitResult QPartialEnergy::GetBestFit()
{

  fFitResult.Reset();
  if( fPMTData.empty() )
    return fFitResult;
  CopySeedToResult(); //use seed results

  // Use the seed position to estimate the energy
  DS::FitVertex fitVertex = fFitResult.GetVertex( 0 );
  fEventPosition = fitVertex.GetPosition();

  Double_t xPosition = fEventPosition.X();
  Double_t yPosition = fEventPosition.Y();
  Double_t zPosition = fFitResult.GetVertex(0).GetPosition().Z(); //same as above line
  double rVertex = sqrt( xPosition*xPosition + yPosition*yPosition + (zPosition-fZoff[2])*(zPosition-fZoff[2]) );
  double RPosition = sqrt(xPosition*xPosition + yPosition*yPosition);
  //set Q-correction
  int ZMatrix = int((zPosition+6000)/500);
  int RMatrix = int (RPosition/500);
  int CMatrix = ZMatrix + RMatrix*24;
  if (CMatrix>287) CMatrix=287 ; //bin set, 12*24=288
  bool valid = false; // Only valid if
  if (rVertex>fAVInnerRadius) {
     fQcorrection = fQDummy;//dummy value for outside ofAV;
     valid = false;
  }
  else {
     fQcorrection = fQratio[CMatrix];
     valid = true;
  }

  double totalcharge = 0;
  //PMTCalSelector could be improved for future use
  for( unsigned int pmt = 0; pmt < fPMTData.size(); pmt++ ){
     const RAT::FitterPMT& Gpmt = fPMTData[pmt];
     double charge = Gpmt.GetQHS();
     if(charge>0) totalcharge += charge;
  }

  //set the z-correction
  double zCorrect = fzCorrectP[0]+fzCorrectP[1]*(fSplitZ)+fzCorrectP[2]*fSplitZ*fSplitZ+fzCorrectP[3]*fSplitZ*fSplitZ*fSplitZ;
  fQcorrection = fQcorrection*zCorrect;
  totalcharge = totalcharge/fQcorrection; //corrected to reference point; detector states considered

  double Energy = fEnergyP[0] + fEnergyP[1]*totalcharge/fCorrt_NCh_Eff + fEnergyP[2]*totalcharge*totalcharge/fCorrt_NCh_Eff/fCorrt_NCh_Eff/fCorrt_NCh_Eff + fEnergyP[3]*totalcharge*totalcharge*totalcharge;
  double error = sqrt(fEnergyErorrP[0]*fEnergyErorrP[0]+fEnergyErorrP[1]*fEnergyErorrP[1]*totalcharge*totalcharge
         +fEnergyErorrP[2]*fEnergyErorrP[2]*totalcharge*totalcharge*totalcharge*totalcharge + fEnergyErorrP[3]*fEnergyErorrP[3]*totalcharge*totalcharge*totalcharge*totalcharge*totalcharge*totalcharge);

  // store results
  if (Energy<fEnergy) valid = false;
  fitVertex.SetEnergy(Energy, valid);
  fitVertex.SetPositiveEnergyError(error);
  fitVertex.SetNegativeEnergyError(error);
  fFitResult.SetVertex(0, fitVertex);
  return fFitResult;

}