/*
 *  Delphes: a framework for fast simulation of a generic collider experiment
 *  Copyright (C) 2012-2014  Universite catholique de Louvain (UCL), Belgium
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "display/DelphesPlotSummary.h"
#include "TRootEmbeddedCanvas.h"
#include <algorithm>

bool vecsorter (TLorentzVector i,TLorentzVector j) { return (i.Pt()>j.Pt()); }

DelphesPlotSummary::DelphesPlotSummary(TEveWindowTab* tab):tab_(tab) {}

DelphesPlotSummary::~DelphesPlotSummary() {}

void DelphesPlotSummary::Progress(Int_t p)
{
  Emit("Progress(Int_t)",p);
}

void DelphesPlotSummary::Init(std::vector<DelphesBranchBase*>& elements) {
  elements_ = &elements;
  // loop on the elements, and create tabs
  for(std::vector<DelphesBranchBase*>::iterator data=elements.begin();data<elements.end();++data) {
    // the canvas
    TEveWindowSlot* slot = tab_->NewSlot();
    TRootEmbeddedCanvas* trec = new TRootEmbeddedCanvas();
    TCanvas* canvas = trec->GetCanvas();
    TEveWindowFrame * wf = slot->MakeFrame(trec);
    wf->SetElementName((*data)->GetName());
    canvas->Divide(3,3);
    canvases_[(*data)->GetName()] = canvas;
    // the histograms
    TH1F* h;
    std::vector<TH1F*> histograms;
    histograms.reserve(9);
    h = new TH1F(Form("%sPt",(*data)->GetName()),Form("%s Pt",(*data)->GetName()),100,0,-1);
    histograms.push_back(h);
    h = new TH1F(Form("%sEta",(*data)->GetName()),Form("%s Eta",(*data)->GetName()),100,0,-1);
    histograms.push_back(h);
    h = new TH1F(Form("%sPhi",(*data)->GetName()),Form("%s Phi",(*data)->GetName()),100,0,-1);
    histograms.push_back(h);
    h = new TH1F(Form("l%sPt",(*data)->GetName()),Form("leading %s Pt",(*data)->GetName()),100,0,-1);
    histograms.push_back(h);
    h = new TH1F(Form("l%sEta",(*data)->GetName()),Form("leading %s Eta",(*data)->GetName()),100,0,-1);
    histograms.push_back(h);
    h = new TH1F(Form("l%sPhi",(*data)->GetName()),Form("leading %s Phi",(*data)->GetName()),100,0,-1);
    histograms.push_back(h);
    h = new TH1F(Form("sl%sPt",(*data)->GetName()),Form("subleading %s Pt",(*data)->GetName()),100,0,-1);
    histograms.push_back(h);
    h = new TH1F(Form("sl%sEta",(*data)->GetName()),Form("subleading %s Eta",(*data)->GetName()),100,0,-1);
    histograms.push_back(h);
    h = new TH1F(Form("sl%sPhi",(*data)->GetName()),Form("subleading %s Phi",(*data)->GetName()),100,0,-1);
    histograms.push_back(h);
    histograms_[(*data)->GetName()] = histograms;
    // the event histograms
    TH1F* h1 = (TH1F*)histograms[0]->Clone(); h1->Reset(); h1->SetLineColor(kBlue);
    TH1F* h2 = (TH1F*)histograms[1]->Clone(); h2->Reset(); h2->SetLineColor(kBlue);
    TH1F* h3 = (TH1F*)histograms[2]->Clone(); h3->Reset(); h3->SetLineColor(kBlue);
    std::vector<TH1F*> hv;
    hv.push_back(h1);
    hv.push_back(h2);
    hv.push_back(h3);
    eventProfiles_[(*data)->GetName()] = hv;
    // the event markers
    TMarker *m;
    std::vector<TMarker*> mv;
    m = new TMarker(0,0,29); m->SetMarkerColor(kBlue); m->SetMarkerSize(3);
    mv.push_back(m);
    m = new TMarker(0,0,29); m->SetMarkerColor(kBlue); m->SetMarkerSize(3);
    mv.push_back(m);
    m = new TMarker(0,0,29); m->SetMarkerColor(kBlue); m->SetMarkerSize(3);
    mv.push_back(m);
    m = new TMarker(0,0,29); m->SetMarkerColor(kBlue); m->SetMarkerSize(3);
    mv.push_back(m);
    m = new TMarker(0,0,29); m->SetMarkerColor(kBlue); m->SetMarkerSize(3);
    mv.push_back(m);
    m = new TMarker(0,0,29); m->SetMarkerColor(kBlue); m->SetMarkerSize(3);
    mv.push_back(m);
    eventMarkers_[(*data)->GetName()] = mv;
  }
}

void DelphesPlotSummary::FillSample(ExRootTreeReader* treeReader, Int_t event_id) {
  Int_t entries = treeReader->GetEntries();
  for(Int_t i=0;i<entries;++i) {
    treeReader->ReadEntry(i);
    for(std::vector<DelphesBranchBase*>::iterator element = elements_->begin();element<elements_->end();++element) {
      std::vector<TLorentzVector> vectors = (*element)->GetVectors();
      std::sort(vectors.begin(), vectors.end(), vecsorter); 
      std::vector<TH1F*> histograms = histograms_[(*element)->GetName()];
      for(std::vector<TLorentzVector>::iterator it=vectors.begin(); it<vectors.end();++it) {
        histograms[0]->Fill(it->Pt());
        histograms[1]->Fill(it->Eta());
        histograms[2]->Fill(it->Phi());
        if(it==vectors.begin()) {
          histograms[3]->Fill(it->Pt());
          histograms[4]->Fill(it->Eta());
          histograms[5]->Fill(it->Phi());
        }
        if(it==vectors.begin()+1) {
          histograms[6]->Fill(it->Pt());
          histograms[7]->Fill(it->Eta());
          histograms[8]->Fill(it->Phi());
        }
      }
    }
    Progress(int(100*i/entries));
  }
  treeReader->ReadEntry(event_id);
  Progress(100);
}

void DelphesPlotSummary::Draw() {
  for(std::map< TString, TCanvas* >::iterator it=canvases_.begin(); it!=canvases_.end(); ++it) {
    TCanvas* c = it->second;
    std::vector<TH1F*> histograms = histograms_[it->first];
    std::vector<TH1F*> eventProfiles = eventProfiles_[it->first];
    std::vector<TMarker*> eventMarkers = eventMarkers_[it->first];
    for(Int_t i=0;i<9;++i) {
      c->cd(i+1);
      if(histograms[i]->GetEntries()==0) continue;
      histograms[i]->Draw();
      if(i<3) {
        eventProfiles[i]->Draw("same");
      } else {
        eventMarkers[i-3]->Draw("same");
      }
    }
    c->Update();
  } 
}

void DelphesPlotSummary::FillEvent() {
  // clear event histograms and markers
  for(std::map< TString, std::vector<TH1F*> >::iterator hv = eventProfiles_.begin(); hv!=eventProfiles_.end();++hv) {
    for(std::vector<TH1F*>::iterator h = hv->second.begin(); h<hv->second.end();++h) {
      (*h)->Reset();
    }
  }
  for(std::map< TString, std::vector<TMarker*> >::iterator mv = eventMarkers_.begin(); mv!=eventMarkers_.end();++mv) {
    for(std::vector<TMarker*>::iterator m = mv->second.begin(); m<mv->second.end();++m) {
      (*m)->SetMarkerSize(0);
    }
  }
  // loop over the elements and fill markers with event data
  for(std::vector<DelphesBranchBase*>::iterator element = elements_->begin();element<elements_->end();++element) {
    std::vector<TLorentzVector> vectors = (*element)->GetVectors();
    std::sort(vectors.begin(), vectors.end(), vecsorter); 
    std::vector<TH1F*> hv = eventProfiles_[(*element)->GetName()];
    TH1F* h1 = hv[0]; h1->Reset();
    TH1F* h2 = hv[1]; h1->Reset();
    TH1F* h3 = hv[2]; h1->Reset();
    std::vector<TMarker*> mv = eventMarkers_[(*element)->GetName()];
    for(std::vector<TLorentzVector>::iterator it=vectors.begin(); it<vectors.end();++it) {
      h1->Fill(it->Pt());
      h2->Fill(it->Eta());
      h3->Fill(it->Phi());
      if(it==vectors.begin()) {
        mv[0]->SetX(it->Pt()); mv[0]->SetMarkerSize(3);
        mv[1]->SetX(it->Eta()); mv[1]->SetMarkerSize(3);
        mv[2]->SetX(it->Phi()); mv[2]->SetMarkerSize(3);
      }
      if(it==vectors.begin()+1) {
        mv[3]->SetX(it->Pt()); mv[3]->SetMarkerSize(3);
        mv[4]->SetX(it->Eta()); mv[4]->SetMarkerSize(3);
        mv[5]->SetX(it->Phi()); mv[5]->SetMarkerSize(3);
      }
    }
  }
}