#include <set>
#include <map>
#include <utility>
#include <vector>
#include <algorithm>
#include <sstream>
#include <exception>
#include "TGeoManager.h"
#include "TGeoVolume.h"
#include "TGeoMedium.h"
#include "TGeoNode.h"
#include "TGeoCompositeShape.h"
#include "TGeoMatrix.h"
#include "TGeoTube.h"
#include "TGeoCone.h"
#include "TGeoArb8.h"
#include "external/ExRootAnalysis/ExRootConfReader.h"
#include "external/ExRootAnalysis/ExRootTreeReader.h"
#include "display/DelphesCaloData.h"
#include "display/DelphesDisplay.h"
#include "display/DelphesBranchElement.h"
#include "display/Delphes3DGeometry.h"
#include "classes/DelphesClasses.h"
#include "TF2.h"
#include "TH1F.h"
#include "TChain.h"
#include "TEveElement.h"
#include "TEveJetCone.h"
#include "TEveTrack.h"
#include "TEveTrackPropagator.h"
#include "TEveCalo.h"
#include "TMath.h"
#include "TSystem.h"
#include "TEveManager.h"
#include "TEveGeoNode.h"
#include "TEveTrans.h"
#include "TEveViewer.h"
#include "TEveBrowser.h"
#include "TRootBrowser.h"
#include "TGLViewer.h"
#include "TGButton.h"
#include "TCollection.h"
#include "TClonesArray.h"
#include "TGLClip.h"
#include "TEveArrow.h"

/*
 * assembly.C: sauvegarde as shape-extract -> implement in the geometry class (read/write)
 * histobrowser.C: intégration d'histogrammes dans le display (on pourrait avoir Pt, eta, phi pour les principales collections)
 * also from alice_esd: summary html table
 * 
 */
using namespace std;

// Forward declarations.
class Delphes3DGeometry;
class ExRootTreeReader;
class DelphesCaloData;
class DelphesDisplay;
class DelphesBranchBase;
template<typename EveContainer> class DelphesBranchElement;
void make_gui();
void load_event();
void delphes_read();
void delphes_read_towers(TClonesArray* data, DelphesBranchBase* element);
void delphes_read_tracks(TClonesArray* data, DelphesBranchBase* element);
void delphes_read_jets(TClonesArray* data, DelphesBranchBase* element);
void delphes_read_vectors(TClonesArray* data, DelphesBranchBase* element);
void readConfig(const char *configFile, const Delphes3DGeometry& det3D, std::vector<DelphesBranchBase*>& elements, std::vector<TClonesArray*>& arrays);

// Configuration and global variables.
Int_t event_id           = 0; // Current event id.
Double_t gRadius = 1.29;
Double_t gTotRadius = 2.0;
Double_t gHalfLength = 3.0;
Double_t gBz = 3.8;

TChain gChain("Delphes");

ExRootTreeReader *gTreeReader = 0;

std::vector<DelphesBranchBase*> gElements;
std::vector<TClonesArray*> gArrays;

DelphesDisplay *gDelphesDisplay = 0;

/******************************************************************************/
// Initialization and steering functions
/******************************************************************************/

// function that parses the config to extract the branches of interest and prepare containers
void readConfig(const char *configFile, Delphes3DGeometry& det3D, std::vector<DelphesBranchBase*>& elements, std::vector<TClonesArray*>& arrays) {
   ExRootConfReader *confReader = new ExRootConfReader;
   confReader->ReadFile(configFile);
   Double_t tk_radius = det3D.getTrackerRadius();
   Double_t tk_length = det3D.getTrackerHalfLength();
   Double_t tk_Bz     = det3D.getBField();
   Double_t mu_radius = det3D.getDetectorRadius();
   Double_t mu_length = det3D.getDetectorHalfLength();
   TAxis*   etaAxis   = det3D.getCaloAxes().first;
   TAxis*   phiAxis   = det3D.getCaloAxes().second;
   ExRootConfParam branches = confReader->GetParam("TreeWriter::Branch");
   Int_t nBranches = branches.GetSize()/3;
   DelphesBranchElement<TEveTrackList>* tlist;
   DelphesBranchElement<DelphesCaloData>* clist;
   DelphesBranchElement<TEveElementList>* elist;
   for(Int_t b = 0; b<nBranches; ++b) {
     TString input = branches[b*3].GetString();
     TString name = branches[b*3+1].GetString();
     TString className = branches[b*3+2].GetString();
     if(className=="Track") {
       if(input.Contains("eflow",TString::kIgnoreCase) || name.Contains("eflow",TString::kIgnoreCase)) continue; //no eflow
       tlist = new DelphesBranchElement<TEveTrackList>(name,"track",kBlue);
       elements.push_back(tlist);
       TEveTrackPropagator *trkProp = tlist->GetContainer()->GetPropagator();
       trkProp->SetMagField(0., 0., -tk_Bz);
       trkProp->SetMaxR(tk_radius);
       trkProp->SetMaxZ(tk_length);
     } else if(className=="Tower") {
       if(input.Contains("eflow",TString::kIgnoreCase) || name.Contains("eflow",TString::kIgnoreCase)) continue; //no eflow
       clist = new DelphesBranchElement<DelphesCaloData>(name,"tower",kBlack);
       clist->GetContainer()->SetEtaBins(etaAxis);
       clist->GetContainer()->SetPhiBins(phiAxis);
       elements.push_back(clist);
     } else if(className=="Jet") {
       if(input.Contains("GenJetFinder")) {
         elist = new DelphesBranchElement<TEveElementList>(name,"jet",kCyan);
         elist->GetContainer()->SetRnrSelf(false);
         elist->GetContainer()->SetRnrChildren(false);
         elements.push_back(elist);
       } else {
         elements.push_back(new DelphesBranchElement<TEveElementList>(name,"jet",kYellow));
       }
     } else if(className=="Electron") {
       tlist = new DelphesBranchElement<TEveTrackList>(name,"track",kRed);
       elements.push_back(tlist);
       TEveTrackPropagator *trkProp = tlist->GetContainer()->GetPropagator();
       trkProp->SetMagField(0., 0., -tk_Bz);
       trkProp->SetMaxR(tk_radius);
       trkProp->SetMaxZ(tk_length);
     } else if(className=="Photon") {
       tlist = new DelphesBranchElement<TEveTrackList>(name,"photon",kYellow);
       elements.push_back(tlist);
       TEveTrackPropagator *trkProp = tlist->GetContainer()->GetPropagator();
       trkProp->SetMagField(0., 0., 0.);
       trkProp->SetMaxR(tk_radius);
       trkProp->SetMaxZ(tk_length);
     } else if(className=="Muon") {
       tlist = new DelphesBranchElement<TEveTrackList>(name,"track",kGreen);
       elements.push_back(tlist);
       TEveTrackPropagator *trkProp = tlist->GetContainer()->GetPropagator();
       trkProp->SetMagField(0., 0., -tk_Bz);
       trkProp->SetMaxR(mu_radius);
       trkProp->SetMaxZ(mu_length);
     } else if(className=="MissingET") {
       elements.push_back(new DelphesBranchElement<TEveElementList>(name,"vector",kViolet));
     } else if(className=="GenParticle") {
       tlist = new DelphesBranchElement<TEveTrackList>(name,"track",kCyan);
       elements.push_back(tlist);
       tlist->GetContainer()->SetRnrSelf(false);
       tlist->GetContainer()->SetRnrChildren(false);
       TEveTrackPropagator *trkProp = tlist->GetContainer()->GetPropagator();
       trkProp->SetMagField(0., 0., -tk_Bz);
       trkProp->SetMaxR(tk_radius);
       trkProp->SetMaxZ(tk_length);
     }
//TODO one possible simplification could be to add the array to the element class.
     arrays.push_back(gTreeReader->UseBranch(name));
   }
}

void delphes_event_display(const char *configFile, const char *inputFile, Delphes3DGeometry& det3D)
{

   // initialize the application
   TEveManager::Create(kTRUE, "IV");
   TGeoManager* geom = gGeoManager;

   // build the detector
   gRadius = det3D.getTrackerRadius();
   gTotRadius = det3D.getDetectorRadius();
   gHalfLength = det3D.getTrackerHalfLength();
   gBz = det3D.getBField();

   //TODO specific to some classical detector... could use better the det3D
   TGeoVolume* top = det3D.getDetector(false);
   geom->SetTopVolume(top);
   TEveElementList *geometry = new TEveElementList("Geometry");
   TEveGeoTopNode* trk = new TEveGeoTopNode(gGeoManager, top->FindNode("tracker_1"));
   trk->SetVisLevel(6);
   geometry->AddElement(trk);
   TEveGeoTopNode* calo = new TEveGeoTopNode(gGeoManager, top->FindNode("Calorimeter_barrel_1"));
   calo->SetVisLevel(3);
   geometry->AddElement(calo);
   calo = new TEveGeoTopNode(gGeoManager, top->FindNode("Calorimeter_endcap_1"));
   calo->SetVisLevel(3);
   calo->UseNodeTrans();
   geometry->AddElement(calo);
   calo = new TEveGeoTopNode(gGeoManager, top->FindNode("Calorimeter_endcap_2"));
   calo->SetVisLevel(3);
   calo->UseNodeTrans();
   geometry->AddElement(calo);
   TEveGeoTopNode* muon = new TEveGeoTopNode(gGeoManager, top->FindNode("muons_barrel_1"));
   muon->SetVisLevel(4);
   geometry->AddElement(muon);
   muon = new TEveGeoTopNode(gGeoManager, top->FindNode("muons_endcap_1"));
   muon->SetVisLevel(4);
   muon->UseNodeTrans();
   geometry->AddElement(muon);
   muon = new TEveGeoTopNode(gGeoManager, top->FindNode("muons_endcap_2"));
   muon->SetVisLevel(4);
   muon->UseNodeTrans();
   geometry->AddElement(muon);
   //gGeoManager->DefaultColors();

   // Create chain of root trees
   gChain.Add(inputFile);

   // Create object of class ExRootTreeReader
   printf("*** Opening Delphes data file ***\n");
   gTreeReader = new ExRootTreeReader(&gChain);

   // prepare data collections
   readConfig(configFile, det3D, gElements, gArrays);

   // viewers and scenes
   gDelphesDisplay = new DelphesDisplay;
   gEve->AddGlobalElement(geometry);
   gDelphesDisplay->ImportGeomRPhi(geometry);
   gDelphesDisplay->ImportGeomRhoZ(geometry);
   // find the first calo data and use that to initialize the calo display
   for(std::vector<DelphesBranchBase*>::iterator data=gElements.begin();data<gElements.end();++data) {
     if(TString((*data)->GetType())=="tower") {
//TODO: why do I have to split this in two lines??? seems a cint bug?
       DelphesBranchElement<DelphesCaloData>* data_tmp = dynamic_cast<DelphesBranchElement<DelphesCaloData>*>(*data);
       DelphesCaloData* container =  data_tmp->GetContainer();
//       DelphesCaloData* container = dynamic_cast<DelphesBranchElement<DelphesCaloData>*>((*data))->GetContainer();
       assert(container);
       TEveCalo3D *calo3d = new TEveCalo3D(container);
       calo3d->SetBarrelRadius(gRadius);
       calo3d->SetEndCapPos(gHalfLength);
       gEve->AddGlobalElement(calo3d);
       gDelphesDisplay->ImportCaloRPhi(calo3d);
       gDelphesDisplay->ImportCaloRhoZ(calo3d);
       TEveCaloLego *lego = new TEveCaloLego(container);
       lego->InitMainTrans();
       lego->RefMainTrans().SetScale(TMath::TwoPi(), TMath::TwoPi(), TMath::Pi());
       lego->SetAutoRebin(kFALSE);
       lego->Set2DMode(TEveCaloLego::kValSizeOutline);
       gDelphesDisplay->ImportCaloLego(lego);
       break;
     }
   }
   gEve->Redraw3D(kTRUE);
}

//______________________________________________________________________________
void load_event()
{
   // Load event specified in global event_id.
   // The contents of previous event are removed.

   //TODO move this to the status bar ???
   printf("Loading event %d.\n", event_id);

   // clear the previous event
   gEve->GetViewers()->DeleteAnnotations();
   for(std::vector<DelphesBranchBase*>::iterator data=gElements.begin();data<gElements.end();++data) {
     (*data)->Reset();
   }

   // read the new event
   delphes_read();

   // update display
   TEveElement* top = (TEveElement*)gEve->GetCurrentEvent();
   gDelphesDisplay->DestroyEventRPhi();
   gDelphesDisplay->ImportEventRPhi(top);
   gDelphesDisplay->DestroyEventRhoZ();
   gDelphesDisplay->ImportEventRhoZ(top);
   //update_html_summary();
   gEve->Redraw3D(kFALSE, kTRUE);
}

void delphes_read()
{

  // safety
  if(event_id >= gTreeReader->GetEntries() || event_id<0 ) return;

  // Load selected branches with data from specified event
  gTreeReader->ReadEntry(event_id);

  // loop over selected branches, and apply the proper recipe to fill the collections.
  // this is basically to loop on gArrays to fill gElements.

//TODO: one option would be to have templated methods in the element classes. We could simply call "element.fill()"
  std::vector<TClonesArray*>::iterator data = gArrays.begin();
  std::vector<DelphesBranchBase*>::iterator element = gElements.begin();
  std::vector<TClonesArray*>::iterator data_tracks = gArrays.begin();
  std::vector<DelphesBranchBase*>::iterator element_tracks = gElements.begin();
  Int_t nTracks = 0;
  for(; data<gArrays.end() && element<gElements.end(); ++data, ++element) {
    TString type = (*element)->GetType();
    // keep the most generic track collection for the end
    if(type=="track" && TString((*element)->GetClassName())=="Track" && nTracks==0) {
      data_tracks = data;
      element_tracks = element;
      nTracks = (*data_tracks)->GetEntries();
      continue;
    }
    // branch on the element type
    if(type=="tower") delphes_read_towers(*data,*element);
    else if(type=="track" || type=="photon") delphes_read_tracks(*data,*element);
    else if(type=="jet") delphes_read_jets(*data,*element);
    else if(type=="vector") delphes_read_vectors(*data,*element);
  }
  // finish whith what we consider to be the main track collection
  if(nTracks>0) delphes_read_tracks(*data,*element);
}

void delphes_read_towers(TClonesArray* data, DelphesBranchBase* element) {
  DelphesCaloData* container = dynamic_cast<DelphesBranchElement<DelphesCaloData>*>(element)->GetContainer();
  assert(container);
  // Loop over all towers
  TIter itTower(data);
  Tower *tower;
  while((tower = (Tower *) itTower.Next()))
  {
    container->AddTower(tower->Edges[0], tower->Edges[1], tower->Edges[2], tower->Edges[3]);
    container->FillSlice(0, tower->Eem);
    container->FillSlice(1, tower->Ehad);
  }
  container->DataChanged();
}

void delphes_read_tracks(TClonesArray* data, DelphesBranchBase* element) {
  TEveTrackList* container = dynamic_cast<DelphesBranchElement<TEveTrackList>*>(element)->GetContainer();
  assert(container);
  TString className = element->GetClassName();
  TIter itTrack(data);
  Int_t counter = 0;
  TEveTrack *eveTrack;
  TEveTrackPropagator *trkProp = container->GetPropagator();
  if(className=="Track") {
    // Loop over all tracks
    Track *track;
    while((track = (Track *) itTrack.Next())) {
      TParticle pb(track->PID, 1, 0, 0, 0, 0,
                   track->P4().Px(), track->P4().Py(),
                   track->P4().Pz(), track->P4().E(),
                   track->X, track->Y, track->Z, 0.0);

      eveTrack = new TEveTrack(&pb, counter, trkProp);
      eveTrack->SetName(Form("%s [%d]", pb.GetName(), counter++));
      eveTrack->SetStdTitle();
      eveTrack->SetAttLineAttMarker(container);
      container->AddElement(eveTrack);
      eveTrack->SetLineColor(element->GetColor());
      eveTrack->MakeTrack();
    }
  } else if(className=="Electron") {
    // Loop over all electrons
    Electron *electron;
    while((electron = (Electron *) itTrack.Next())) {
      TParticle pb(electron->Charge<0?11:-11, 1, 0, 0, 0, 0,
                   electron->P4().Px(), electron->P4().Py(),
                   electron->P4().Pz(), electron->P4().E(),
                   0., 0., 0., 0.);

      eveTrack = new TEveTrack(&pb, counter, trkProp);
      eveTrack->SetName(Form("%s [%d]", pb.GetName(), counter++));
      eveTrack->SetStdTitle();
      eveTrack->SetAttLineAttMarker(container);
      container->AddElement(eveTrack);
      eveTrack->SetLineColor(element->GetColor());
      eveTrack->MakeTrack();
  }
  } else if(className=="Muon") {
    // Loop over all muons
    Muon *muon;
    while((muon = (Muon *) itTrack.Next())) {
      TParticle pb(muon->Charge<0?13:-13, 1, 0, 0, 0, 0,
                   muon->P4().Px(), muon->P4().Py(),
                   muon->P4().Pz(), muon->P4().E(),
                   0., 0., 0., 0.);

      eveTrack = new TEveTrack(&pb, counter, trkProp);
      eveTrack->SetName(Form("%s [%d]", pb.GetName(), counter++));
      eveTrack->SetStdTitle();
      eveTrack->SetAttLineAttMarker(container);
      container->AddElement(eveTrack);
      eveTrack->SetLineColor(element->GetColor());
      eveTrack->MakeTrack();
    }
  } else if(className=="Photon") {
    // Loop over all photons
    Photon *photon;
    while((photon = (Photon *) itTrack.Next())) {
      TParticle pb(22, 1, 0, 0, 0, 0,
                   photon->P4().Px(), photon->P4().Py(),
                   photon->P4().Pz(), photon->P4().E(),
                   0., 0., 0., 0.);

      eveTrack = new TEveTrack(&pb, counter, trkProp);
      eveTrack->SetName(Form("%s [%d]", pb.GetName(), counter++));
      eveTrack->SetStdTitle();
      eveTrack->SetAttLineAttMarker(container);
      container->AddElement(eveTrack);
      eveTrack->SetLineColor(element->GetColor());
      eveTrack->MakeTrack();
    }
  }
}

void delphes_read_jets(TClonesArray* data, DelphesBranchBase* element) {
  TEveElementList* container = dynamic_cast<DelphesBranchElement<TEveElementList>*>(element)->GetContainer();
  assert(container);
  TIter itJet(data);
  Jet *jet;
  TEveJetCone *eveJetCone;
  // Loop over all jets
  Int_t counter = 0;
  while((jet = (Jet *) itJet.Next()))
  {
    eveJetCone = new TEveJetCone();
    eveJetCone->SetTitle(Form("jet [%d]: Pt=%f, Eta=%f, \nPhi=%f, M=%f",counter,jet->PT, jet->Eta, jet->Phi, jet->Mass));
    eveJetCone->SetName(Form("jet [%d]", counter++));
    eveJetCone->SetMainTransparency(60);
    eveJetCone->SetLineColor(element->GetColor());
    eveJetCone->SetFillColor(element->GetColor());
    eveJetCone->SetCylinder(gRadius - 10, gHalfLength - 10);
    eveJetCone->SetPickable(kTRUE);
    eveJetCone->AddEllipticCone(jet->Eta, jet->Phi, jet->DeltaEta, jet->DeltaPhi);
    container->AddElement(eveJetCone);
  }
}

void delphes_read_vectors(TClonesArray* data, DelphesBranchBase* element) {
  TEveElementList* container = dynamic_cast<DelphesBranchElement<TEveElementList>*>(element)->GetContainer();
  assert(container);
  TIter itMet(data);
  MissingET *MET;
  TEveArrow *eveMet;
  // Missing Et
  Double_t maxPt = 50.;
  // TODO to be changed as we don't have access to maxPt anymore. MET scale could be a general parameter set in GUI
  while((MET = (MissingET*) itMet.Next())) {
    eveMet = new TEveArrow((gRadius * MET->MET/maxPt)*cos(MET->Phi), (gRadius * MET->MET/maxPt)*sin(MET->Phi), 0., 0., 0., 0.);
    eveMet->SetMainColor(element->GetColor());
    eveMet->SetTubeR(0.04);
    eveMet->SetConeR(0.08);
    eveMet->SetConeL(0.10);
    eveMet->SetPickable(kTRUE);
    eveMet->SetName("Missing Et");
    eveMet->SetTitle(Form("Missing Et (%.1f GeV)",MET->MET));
    container->AddElement(eveMet);
  }
}

/******************************************************************************/
// GUI
/******************************************************************************/

//______________________________________________________________________________
// 
// EvNavHandler class is needed to connect GUI signals.

class EvNavHandler
{     
public:
   void Fwd()
   {  
      if (event_id < gTreeReader->GetEntries() - 1) {
         ++event_id;
         load_event();
      } else {
         printf("Already at last event.\n");
      }
   }
   void Bck()
   {
      if (event_id > 0) {
         --event_id;
         load_event();
      } else {
         printf("Already at first event.\n");
      }
   }
};

//______________________________________________________________________________
void make_gui()
{
   // Create minimal GUI for event navigation.
   // TODO: better GUI could be made based on the ch15 of the manual (Writing a GUI)

   // add a tab on the left
   TEveBrowser* browser = gEve->GetBrowser();
   browser->StartEmbedding(TRootBrowser::kLeft);

   // set the main title
   TGMainFrame* frmMain = new TGMainFrame(gClient->GetRoot(), 1000, 600);
   frmMain->SetWindowName("Delphes Event Display");
   frmMain->SetCleanup(kDeepCleanup);

   // build the navigation menu
   TGHorizontalFrame* hf = new TGHorizontalFrame(frmMain);
   {
      TString icondir;
      if(gSystem->Getenv("ROOTSYS"))
        icondir = Form("%s/icons/", gSystem->Getenv("ROOTSYS"));
      if(!gSystem->OpenDirectory(icondir)) 
        icondir = Form("%s/icons/", (const char*)gSystem->GetFromPipe("root-config --etcdir") );
      TGPictureButton* b = 0;
      EvNavHandler    *fh = new EvNavHandler;

      b = new TGPictureButton(hf, gClient->GetPicture(icondir+"GoBack.gif"));
      hf->AddFrame(b);
      b->Connect("Clicked()", "EvNavHandler", fh, "Bck()");

      b = new TGPictureButton(hf, gClient->GetPicture(icondir+"GoForward.gif"));
      hf->AddFrame(b);
      b->Connect("Clicked()", "EvNavHandler", fh, "Fwd()");
   }
   frmMain->AddFrame(hf);
   frmMain->MapSubwindows();
   frmMain->Resize();
   frmMain->MapWindow();
   browser->StopEmbedding();
   browser->SetTabTitle("Event Control", 0);
}

/******************************************************************************/
// MAIN 
/******************************************************************************/

void geometry(const char* filename = "delphes_card_CMS.tcl", const char* ParticlePropagator="ParticlePropagator",
                                                             const char* TrackingEfficiency="ChargedHadronTrackingEfficiency",
                                                             const char* MuonEfficiency="MuonEfficiency",
                                                             const char* Calorimeters="Calorimeter") 
{

   // load the libraries
   gSystem->Load("libGeom");
   //gSystem->Load("../libDelphes");
   gSystem->Load("../libDelphesDisplay");

   // create the detector representation
   Delphes3DGeometry det3D(new TGeoManager("delphes", "Delphes geometry"));
   det3D.readFile(filename, ParticlePropagator, TrackingEfficiency, MuonEfficiency, Calorimeters);

   // create the application items
   delphes_event_display("delphes_card_CMS.tcl", "../delphes_output.root", det3D); //TODO root file as input cfg
   make_gui();
   load_event();
   gEve->Redraw3D(kTRUE); // Reset camera after the first event has been shown.

   // EClipType not exported to CINT (see TGLUtil.h):
   // 0 - no clip, 1 - clip plane, 2 - clip box
   TGLViewer *v = gEve->GetDefaultGLViewer();
   //Double_t plane[4] = { 0., 1., 0., 0. };
   //v->GetClipSet()->SetClipState(1,plane);
   //v->GetClipSet()->SetClipType(1);
   //v->ColorSet().Background().SetColor(kMagenta+4);
   //v->SetGuideState(TGLUtil::kAxesEdge, kTRUE, kFALSE, 0);
   v->RefreshPadEditor(v);
   v->CurrentCamera().RotateRad(-1.2, 0.5);
   v->DoDraw();

}