#include <RAT/FastZ.hh>
#include <RAT/DB.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DU/PMTInfo.hh>

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

void FastZ::BeginOfRun(DS::Run&) {
  DB* const db = DB::Get();
  DBLinkPtr dblink = db->GetLink("FIT_FAST_Z");
  // constants needed for calculating weighted average
  fTop_cos = dblink->GetD("top_weight");
  fBottom_cos = dblink->GetD("bottom_weight");
  fTop_Water = dblink->GetD("top_water");
  fTop_Scint = dblink->GetD("top_scint");
  fCs = dblink->GetD("c_scint");
  fCw = dblink->GetD("c_water");
  fWindow = dblink->GetD("time_window");
  fTop_Range = dblink->GetD("top_range");
  fBottom_Range = dblink->GetD("bottom_range");

}

DS::FitResult FastZ::GetBestFit() {
  const DU::PMTInfo& pmtInfo = DU::Utility::Get()->GetPMTInfo();

  // default (or error) fit result
  DS::FitVertex dummyVertex;
  dummyVertex.SetPosition(TVector3(0.0, 0.0, 0.0), false);
  dummyVertex.SetTime(0.0, false);
  SelectPMTData(dummyVertex);
  fFitResult.SetVertex(0, dummyVertex);

  // check whether there are enough PMT's
  if (fSelectedPMTData.size() < 2) return fFitResult;

  // calculate fit result
  double fitTime = 0.0;

  std::vector<double> vZ;
  std::vector<double> vTime;

  //storing pmt information of interest
  for (size_t ipmt = 0; ipmt < fSelectedPMTData.size(); ipmt++) {
    const RAT::FitterPMT& pmt = fSelectedPMTData[ipmt];
    const TVector3 pmtPos = pmtInfo.GetPosition(pmt.GetID());
    double pmtZ = pmtPos.Z();


    double pmtTime = pmt.GetTime();

    if (pmtZ > fTop_Range) {
      top_pmt.Fill(pmtTime);
      vZ.push_back(pmtZ);
      vTime.push_back(pmtTime);
    }

    if (pmtZ < fBottom_Range) {
      bot_pmt.Fill(pmtTime);
      vZ.push_back(pmtZ);
      vTime.push_back(pmtTime);
    }
  }

  double top_mode = top_pmt.GetMaximumBin() * 10 + 5;
  double bot_mode = bot_pmt.GetMaximumBin() * 10 + 5;

  double top_time = 10000, bot_time = 10000;
  double top_z = 0, bot_z = 0;


  //finding earliest pmt hit within search range and time window
  for (size_t i = 0; i < vZ.size(); i++) {
    if (vZ[i] > fTop_Range && vTime[i] < top_time && (top_mode - vTime[i]) < fWindow && (top_mode - vTime[i]) > 0) {
      top_time = vTime[i];
      top_z = vZ[i];
    }

    if (vZ[i] < fBottom_Range && vTime[i] < bot_time && (bot_mode - vTime[i]) < fWindow && (bot_mode - vTime[i]) > 0) {
      bot_time = vTime[i];
      bot_z = vZ[i];
    }
  }

  if (top_z == 0 || bot_z == 0) return fFitResult;

  double fitZ = (bot_time - top_time) + bot_z/(fBottom_cos*fCw) + top_z*(fTop_Scint/(fTop_cos*fCs) + fTop_Water/(fTop_cos*fCw));
  fitZ /= (1/(fBottom_cos*fCw) + fTop_Scint/(fTop_cos*fCs) + fTop_Water/(fTop_cos*fCw));

  TVector3 fitPos(0.0, 0.0, fitZ);
  DS::FitVertex fitVertex;
  fitVertex.SetPosition(fitPos);
  fitVertex.SetTime(fitTime);
  fFitResult.SetVertex(0, fitVertex);
  return fFitResult;
}