#include "Riostream.h"
#include "TPad.h"
#include "TView.h"
#include "TROOT.h"
#include "TGeoNode.h"
#include "TGeoVolume.h"
#include "TGeoManager.h"
#include "TVirtualGeoPainter.h"

//_______________________________________________________________________________
class iterplugin : public TGeoIteratorPlugin
{
public:
   iterplugin() : TGeoIteratorPlugin(), fColor(kGreen), fReplicaX(0),fReplicaY(0),fReplicaZ(0) {}
   virtual ~iterplugin() {}
   // Process current node
   virtual void ProcessNode();
   void         Select(Int_t replicaX, Int_t replicaY, Int_t replicaZ, Int_t color) 
                   {fReplicaX=replicaX; fReplicaY=replicaY; fReplicaZ=replicaZ; fColor=color;}
   Bool_t       MatchReplicasFromPath(const char *path);
   Int_t fColor;             // Current color
   Int_t fReplicaX;          // X slice
   Int_t fReplicaY;          // Y slice
   Int_t fReplicaZ;          // Z slice
  

   ClassDef(iterplugin, 0)   // A simple user iterator plugin that changes volume color    
};

ClassImp(iterplugin)

void iterplugin::ProcessNode()
{
   if (!fIterator) return;
   TString path;
   fIterator->GetPath(path);
   if (!MatchReplicasFromPath(path)) return;
   Int_t level = fIterator->GetLevel();
   TGeoVolume *vol = fIterator->GetNode(level)->GetVolume();
   vol->SetLineColor(fColor);
   vol->SetTransparency(0);
}

Bool_t iterplugin::MatchReplicasFromPath(const char *path)
{
//   printf("%s\n", path);
   TString s(path);
   TString snum;
   Int_t ind1, ind2;
   ind1 = s.Index("SLICEX_");
   if (ind1<0) return kFALSE;
   ind2 = s.Index("/",ind1);
   if (ind2<0) ind2 = s.Length();
   snum = s(ind1+7, ind2-ind1-7);   
   if (fReplicaX != snum.Atoi()) return kFALSE;
   ind1 = s.Index("SLICEY_");
   if (ind1<0) return kFALSE;
   ind2 = s.Index("/",ind1);
   if (ind2<0) ind2 = s.Length();
   snum = s(ind1+7,ind2-ind1-7);
   if (fReplicaY != snum.Atoi()) return kFALSE;
   ind1 = s.Index("SLICEZ_");
   if (ind1<0) return kFALSE;
   ind2 = s.Index("/",ind1);
   if (ind2<0) ind2 = s.Length();
   snum = s(ind1+7,ind2-ind1-7);
   if (fReplicaZ != snum.Atoi()) return kFALSE;
   return kTRUE;
}   

iterplugin *plugin = 0;

void select(Int_t rowx, Int_t rowy, Int_t rowz, Int_t color)
{
  plugin->Select(rowx,rowy,rowz,color); 
  gGeoManager->GetGeomPainter()->ModifiedPad(); 
}
  
void test()
{
  new TGeoManager("top","");
  TGeoMaterial *matVacuum = new TGeoMaterial("Vacuum", 0,0,0);
  TGeoMedium *Vacuum = new TGeoMedium("Vacuum",1, matVacuum);
  TGeoVolume *vbox = gGeoManager->MakeBox("vbox", Vacuum, 20,30,40);
  gGeoManager->SetTopVolume(vbox);

  cout << gGeoManager->AddVolume(vbox) << endl;
 
  TGeoVolume *divx = vbox->Divide("SLICEX",1, 5, -10, 2);
  TGeoVolume *divy = divx->Divide("SLICEY",2, 20, -10, 2);
  TGeoVolume *divz = divy->Divide("SLICEZ",3, 20, -10, 2);
  divz->SetLineColor(kBlue);
  divz->SetTransparency(50);
  gGeoManager->CloseGeometry();
  vbox->Draw("ogl");
  TView *view = gPad->GetView(); 
  view->ShowAxis();
  gGeoManager->GetGeomPainter()->ModifiedPad(); 

  plugin = new iterplugin();
  cout << "Total nodes: " << divy->GetNtotal() << endl;

  TString name("/vbox_1/SLICEX_5/SLICEY_20/SLICEZ_20");
  gGeoManager->cd(name.Data());
  TGeoIterator * next = new TGeoIterator(gGeoManager->GetCurrentNode()->GetVolume());
  gGeoManager->GetGeomPainter()->SetIteratorPlugin(plugin);
  plugin->SetIterator(next);
  select(5,20,20,kRed);
}