#include <map>
#include "TRandom.h"
#include "TMath.h"
#include "TList.h"
#include "TThread.h"
#include "TAtomicCount.h"
#include "TCondition.h"
#include "TStopwatch.h"
#include "TGeoManager.h"
#include "TGeoNode.h"
#include "TGeoShape.h"
#include "TGeoMedium.h"
#include "TGeoMatrix.h"

typedef std::map<Long_t, Int_t> ThreadMap_t;
ThreadMap_t thread_map;
Int_t NTHREADS;
Int_t NTRACKS;
Long_t nevents;
Double_t sums[10000];   // max 100 threads
Double_t **data = new Double_t*[10000];
TAtomicCount counter(1);
TAtomicCount finished(1);
TCondition cond;
class MyPropagator;
TList *listThreads = 0;
Int_t crtChunk;

MyPropagator *PROPAGATOR = 0;

//______________________________________________________________________________
struct GeantTrack {
   Double_t x;      // position
   Double_t y;
   Double_t z;
   Double_t px;        // momentum
   Double_t py;
   Double_t pz;
   Double_t rad;       // cumulated material budget
};

GeantTrack **TRACKS = 0;

//______________________________________________________________________________
void GenerateTracks(Int_t ntr)
{
// Generate randomly ntr tracks in the bounding box
   Double_t dbox[3] = {400., 400., 400.};
   if (TRACKS) delete [] TRACKS;
   TRACKS = new GeantTrack*[ntr];
   GeantTrack *track;
   Double_t th,phi;
   for (Int_t i=0; i<ntr; i++) {
      track = new GeantTrack();
      TRACKS[i] = track;
      track->x  = gRandom->Uniform(-dbox[0],dbox[0]);
      track->y  = gRandom->Uniform(-dbox[1],dbox[1]);
      track->z  = gRandom->Uniform(-dbox[2],dbox[2]);
      phi = 2*TMath::Pi()*gRandom->Rndm();
      th = TMath::ACos(1.-2.*gRandom->Rndm());
      track->px = TMath::Sin(th)*TMath::Cos(phi);
      track->py = TMath::Sin(th)*TMath::Sin(phi);
      track->pz = TMath::Cos(th);
      track->rad = 0.;
   }
}
   
//______________________________________________________________________________
class MyPropagator {
protected:
   struct ThreadData_t {
      Int_t indmin;
      Int_t indmax;

      ThreadData_t() : indmin(0), indmax(0) {}
      ~ThreadData_t() {}
   };
   mutable std::vector<ThreadData_t*> fThreadData; //!
   mutable Int_t                      fThreadSize; //!
public:
   MyPropagator() :fThreadSize(0) {}
   ~MyPropagator() {}
   
   MyPropagator::ThreadData_t& GetThreadData() const;
   Double_t PropagateTracks() const;
   Double_t FindRad(Double_t x, Double_t y, Double_t z, Double_t px, Double_t py, Double_t pz) const;
   
   ClassDef(MyPropagator, 0);  // Sum class   
};

ClassImp(MyPropagator)

//______________________________________________________________________________
MyPropagator::ThreadData_t& MyPropagator::GetThreadData() const
{
// Returns thread specific data.
   TThread::Lock();
   Int_t tid = TGeoManager::ThreadId();
   if (tid>=fThreadSize) {
      fThreadData.resize(tid+1);
      fThreadSize = tid+1;
   }
   ThreadData_t *tdata = fThreadData[tid];
   if (tdata == 0) {
      tdata = new ThreadData_t;
      fThreadData[tid] = tdata;
   }   
   tdata->indmin = tid*NTRACKS/NTHREADS;
   tdata->indmax = (tid==NTHREADS)?NTRACKS:(tid+1)*NTRACKS/NTHREADS;
   printf("(%ld) got chunk %d: [%d, %d]\n", TThread::SelfId(), tid, tdata->indmin, tdata->indmax);
   TThread::UnLock();
   return *tdata;
}      

//______________________________________________________________________________
Double_t MyPropagator::PropagateTracks() const
{
// Compute sum of 1/n^2 in a given range
   MyPropagator::ThreadData_t& tdata = GetThreadData();
   Double_t sum = 0;
   GeantTrack *track;
   Int_t n10 = (tdata.indmax-tdata.indmin)/10;
   for (Long_t i=tdata.indmin; i<tdata.indmax; i++) {
      Int_t crt = i-tdata.indmin;
      if (crt && (crt%n10==0)) printf ("=== thread %d: %d%%\n", TGeoManager::ThreadId(),10*crt/n10);
      track = TRACKS[i];
      sum += FindRad(track->x, track->y, track->z, track->px, track->py, track->pz);
   }   
   printf("(%ld) ComputeSum %d-%d summed up sum=%15.13f\n", TThread::SelfId(), tdata.indmin, tdata.indmax-1, sum);
   return sum; 
}

//______________________________________________________________________________
Double_t MyPropagator::FindRad(Double_t x, Double_t y, Double_t z, Double_t xp, Double_t yp, Double_t zp) const
{
   Double_t snext;
   Double_t pt[3];
   Double_t loc[3];
   Double_t epsil = 1.E-2;
   Double_t lastrad = 0.;
   Int_t ismall = 0;
   Int_t nbound = 0;
   Double_t length = 0.;
   Double_t rad    = 0.;
   TGeoMedium *med;
   TGeoShape *shape;
   TGeoNode *lastnode;
   gGeoManager->InitTrack(x,y,z,xp,yp,zp);
   TGeoNode *nextnode = gGeoManager->GetCurrentNode();
   while (nextnode) {
      med = 0;
      if (nextnode) med = nextnode->GetVolume()->GetMedium();
      else return rad;      
      shape = nextnode->GetVolume()->GetShape();
      lastnode = nextnode;
      nextnode = gGeoManager->FindNextBoundaryAndStep();
      snext  = gGeoManager->GetStep();
      if (snext<1.e-8) {
         ismall++;
         if ((ismall<3) && (lastnode != nextnode)) {
            // First try to cross a very thin layer
            length += snext;
            nextnode = gGeoManager->FindNextBoundaryAndStep();
            snext  = gGeoManager->GetStep();
            if (snext<1.E-8) continue;
            // We managed to cross the layer
            ismall = 0;
         } else {  
            // Relocate point
            if (ismall > 3) {
               return rad;
            }   
            memcpy(pt,gGeoManager->GetCurrentPoint(),3*sizeof(Double_t));
            const Double_t *dir = gGeoManager->GetCurrentDirection();
            for (Int_t i=0;i<3;i++) pt[i] += epsil*dir[i];
            snext = epsil;
            length += snext;
            rad += lastrad*snext;
            gGeoManager->CdTop();
            nextnode = gGeoManager->FindNode(pt[0],pt[1],pt[2]);
            if (gGeoManager->IsOutside()) return rad;
            TGeoMatrix *mat = gGeoManager->GetCurrentMatrix();
            mat->MasterToLocal(pt,loc);
            if (!gGeoManager->GetCurrentVolume()->Contains(loc)) {
               gGeoManager->CdUp();
               nextnode = gGeoManager->GetCurrentNode();
            }   
            continue;
         }   
      } else {
         ismall = 0;
      }      
      nbound++;
      length += snext;
      if (med) {
         Double_t radlen = med->GetMaterial()->GetRadLen();
         if (radlen>1.e-5 && radlen<1.e10) {
            lastrad = med->GetMaterial()->GetDensity()/radlen;
            rad += lastrad*snext;
         } else {
            lastrad = 0.;
         }      
      }
   }
   return rad;   
}

//______________________________________________________________________________
void *worker(void *arg)
{
// Working thread. computes a chunk of Sum(1/n^2)
   printf("(%ld) WORKER started\n", TThread::SelfId());
   Int_t tid = TGeoManager::ThreadId();
   // Create navigator for this thread
   gGeoManager->AddNavigator();
   if (TGeoManager::GetNumThreads() == NTHREADS) {
      printf("LOCKING NAVIGATORS\n");
      TGeoManager::SetNavigatorsLock(kTRUE);
   }   
   if (!PROPAGATOR) PROPAGATOR = new MyPropagator;
   sums[tid] = PROPAGATOR->PropagateTracks();
   Long_t remaining = --finished;
   if (!remaining) {
      printf("(%ld) thread %d waking everybody up\n", TThread::SelfId(), tid);
//      cond.Broadcast();
   } else {   
      printf("(%ld) thread %d going to sleep\n", TThread::SelfId(), tid);
//      cond.Wait();
      printf("(%ld) thread %d woke up\n", TThread::SelfId(), tid);
   }   
   return arg;
}

//______________________________________________________________________________
void PropagateTracks(Int_t nthrds, Long_t ntr=10000)
{
   TStopwatch timer;
   NTRACKS = ntr;
   NTHREADS = nthrds;
   // Import geometry
   TGeoManager::Import("http://root.cern.ch/files/alice2.root");
//   TGeoManager::Import("alice.root");
   gGeoManager->SetMultiThread(kTRUE);
   // Generate tracks
   GenerateTracks(ntr);
//   thread_map.clear();
   counter.Set(NTHREADS);
   finished.Set(NTHREADS);
   TThread *t;
   if (!listThreads) {
      listThreads = new TList();
      listThreads->SetOwner();
      for (Int_t j=0; j<10000; j++) {
         t = new TThread(worker);
         listThreads->Add(t);
      }   
   }   
   Int_t i;
   printf("master thread id: %ld\n", TThread::SelfId());
   timer.Start();
   for (i=0; i<NTHREADS; i++) {
      t = (TThread*)listThreads->At(i);
      t->Run();
   }
   for (i=0; i<NTHREADS; i++) {
      t = (TThread*)listThreads->At(i);
      t->Join();
   }
   Double_t total = 0;
   for (i=0; i<NTHREADS; i++) {
      total += sums[i];
      sums[i] = 0.;
   }   
   timer.Stop();
   printf("NTHREADS=%d  TOTAL=%15.13f  CPU time = %f Real time = %f\n", NTHREADS, total, timer.CpuTime(), timer.RealTime());
}

void firsttest(Int_t nthr=4)
{
   PropagateTracks(nthr);
//   ComputeSeries(4);
}