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

#include <vector>
//_______________________________________________________________________________
class iterplugin : public TGeoIteratorPlugin
{
public:
  iterplugin() : TGeoIteratorPlugin() {}
   virtual ~iterplugin() {}
   // Process current node
   virtual void ProcessNode();
  void         Select(vector<Int_t> replicaX, vector<Int_t> replicaY, vector<Int_t> replicaZ,
		      vector<Int_t> color) 
                      {fReplicaX=replicaX; fReplicaY=replicaY; fReplicaZ=replicaZ; fColor=color;}
  Bool_t       MatchReplicasFromPath(int index, const char *path);
   vector<Int_t> fColor;             // Current color
   vector<Int_t> fReplicaX;          // X slice
   vector<Int_t> fReplicaY;          // Y slice
   vector<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);
  
  //printf("path: %s\n", path.Data());
  if (!path.Contains("SLICEZ") ) return;
  Int_t level = fIterator->GetLevel();
  TGeoVolume *vol = fIterator->GetNode(level)->GetVolume();
  
  if (fColor.size()==0) {
    vol->SetTransparency(100);
  }

  for (int i=0; i<fColor.size(); i++) {
    if (!MatchReplicasFromPath(i, path)) {
      vol->SetTransparency(100);
      continue;
    }
    else {
      vol->SetLineColor(fColor[i]);
      vol->SetTransparency(0);
      fColor.erase(fColor.begin()+i);
      fReplicaX.erase(fReplicaX.begin()+i);
      fReplicaY.erase(fReplicaY.begin()+i);
      fReplicaZ.erase(fReplicaZ.begin()+i);
      printf("Ncolor: %i\n", fColor.size());
      return;
    }
  }
}

Bool_t iterplugin::MatchReplicasFromPath(int index, 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[index] != 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[index] != 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[index] != snum.Atoi()) return kFALSE;
   return kTRUE;
}   

iterplugin *plugin = 0;

void select(vector<Int_t> rowx, vector<Int_t> rowy, vector<Int_t> rowz, vector<Int_t> color)
{
  plugin->Select(rowx,rowy,rowz,color); 
  gGeoManager->GetGeomPainter()->ModifiedPad(); 
}
  
void test_many()
{
  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);
  
  gGeoManager->CloseGeometry();
  divz->SetLineColor(kBlue);
  divz->SetTransparency(0);
  //divz->SetVisibility(kFALSE);
  vbox->Draw("ogl");
  //vbox->SetVisContainers();
  TView *view = gPad->GetView(); 
  view->ShowAxis();
  gGeoManager->GetGeomPainter()->ModifiedPad(); 

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

  gGeoManager->GetGeomPainter()->SetIteratorPlugin(plugin);

  vector<int> x, y, z, color;

  x.push_back(1);
  y.push_back(1);
  z.push_back(1);
  color.push_back(3);

  x.push_back(2);
  y.push_back(20);
  z.push_back(20);
  color.push_back(4);

  x.push_back(4);
  y.push_back(20);
  z.push_back(19);
  color.push_back(5);

  x.push_back(4);
  y.push_back(20);
  z.push_back(20);
  color.push_back(2);

  x.push_back(4);
  y.push_back(3);
  z.push_back(9);
  color.push_back(7);

  x.push_back(3);
  y.push_back(20);
  z.push_back(20);
  color.push_back(6);

  select(x,y,z,color);
}