#include "tantra_muon_veto.hh"
#include "TFile.h"
#include "TH3.h"
#include <iostream>
using namespace std;

// the following was copied from frederico versani

std::vector<AntDSTHit *> GetCoincidentHits(std::vector<AntDSTHit>& _hits, float _gate /*= 20*/)
{
  std::vector<AntDSTHit *> coin;

  std::map<int, std::vector<AntDSTHit *>> lcm_map;
  for (unsigned i(0); i < _hits.size(); ++i)
  {
    unsigned lcm_id = _hits.at(i).GetFloorId() + 100 * _hits.at(i).GetLineId();
    lcm_map[lcm_id].push_back(&(_hits.at(i)));
  }
  std::vector<AntDSTHit *>::iterator vec_it1, vec_it2;
  std::map<int, std::vector<AntDSTHit *>>::iterator map_it;
  for (map_it = lcm_map.begin(); map_it != lcm_map.end(); ++map_it)
  {
    for (vec_it1 = map_it->second.begin(); vec_it1 != map_it->second.end(); ++vec_it1)
      for (vec_it2 = map_it->second.begin(); vec_it2 != map_it->second.end(); ++vec_it2)
      {
        if (vec_it1 == vec_it2)
          continue;
        if (fabs((*vec_it1)->GetTime() - (*vec_it2)->GetTime()) < _gate)
        {
          coin.push_back(*vec_it1);
          break;
        }
      }
  }
  return coin;
}

double tantra_muon_veto(AntDST &antdst_event )
{

  // copy from frederico's code, just to be use it's exactly the same
  static const double wc_light = 0.299792458;
  static const double wkCorrectionForGroupVelocity = 0.0298;
  static const double wkWaterIndex = 1.3499;
  static const double wv_light = wc_light / (wkWaterIndex + wkCorrectionForGroupVelocity);

  static TH3D *pdf3show = 0;
  static TH3D *pdf3mup = 0;

  if (!pdf3show)
  {
    //========================================================
    // open Tino's pdfs
    //========================================================

    cout << " loading pdfs for muon veto " << endl;

    TFile *pdf_show = new TFile("/sps/km3net/users/heijboer/aanet/antares/outfile2_sig.root");
    TFile *pdf_mupa = new TFile("/sps/km3net/users/heijboer/aanet/antares/outfile2_mu.root");

    if (!pdf_show || !pdf_mupa ) {
      cout << " failed find root files with tantra-veto likelihoods. " << endl;
    }

    pdf3show = (TH3D *)pdf_show->Get("h43dsig"),
    pdf3mup = (TH3D *)pdf_mupa->Get("h43dmu");
    pdf3show->Scale(1. / pdf3show->Integral());
    pdf3mup->Scale(1. / pdf3mup->Integral());
  }

  auto re = antdst_event.GetStrategy( eShowerTantraFit, eFinalPDFFit ); // step seems to be 3.

  if ( re.fRecParticles.size() != 1 ) {
    cout << "strange number of recparticles for tantra" << endl;
    return 0;
  }

  RecParticle& tantra_track = re.fRecParticles[0];

  std::vector<AntDSTHit> allHits = antdst_event.GetHitsVector();
  auto coiHits = GetCoincidentHits(allHits);

  double Like = 0;
  int NOntime = 0;

  for (unsigned i(0); i < allHits.size(); ++i)
  {
    double TRes(allHits.at(i).GetTime() - tantra_track.GetTime() - (allHits.at(i).GetOMPosition() - tantra_track.GetPosition()).Mag() / wv_light);
    if (TRes > -20 && TRes < 60)
      ++NOntime;
  }

  for (std::vector<AntDSTHit *>::iterator hit_it = coiHits.begin(); hit_it != coiHits.end(); ++hit_it)
  {
    double dist = ((*hit_it)->GetOMPosition() - tantra_track.GetPosition()).Mag();
    double TRes((*hit_it)->GetTime() - tantra_track.GetTime() - dist / wv_light);

    if (fabs(TRes) > 1000) continue;
    double pmu = pdf3mup->GetBinContent(pdf3mup->FindBin(NOntime, TRes, dist));
    double psig = pdf3show->GetBinContent(pdf3show->FindBin(NOntime, TRes, dist));
    static double cap_p(1e-6);
    pmu = std::max(pmu, cap_p);
    psig = std::max(psig, cap_p);
    Like -= (log(pmu) - pmu);
    Like += (log(psig) - psig);
  }


Like = std::min(Like, 5e3);
Like = std::max(Like, -1e3);
return Like;
}