//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class TpcSpaceChargeTask
//     
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Cristoforo Simonetto  	TUM         (original author)
//	Felix Boehmer    	TUM
//
//-----------------------------------------------------------


// This Class' Header ------------------
#include "TpcSpaceChargeTask.h"

// Collaborating Class Headers --------
#include "TpcGas.h"
#include "TClonesArray.h"
#include "FairRootManager.h"
#include "FairRunAna.h"
#include "FairRuntimeDb.h"
#include "TpcPoint.h"
#include "TpcPrimaryCluster.h"
#include "TVector3.h"
#include "TpcDigiPar.h"
#include "TpcDigiMapper.h"

#include <iostream>
#include <iomanip>
#include <fstream>
#include <vector>
#include <cmath>

// Class Member definitions -----------



TpcSpaceChargeTask::TpcSpaceChargeTask()		//default constructor
  : 	FairTask("TPC Space Charge"),
	//ftpcMinR(15.5),	
	//ftpcMaxR(41.5),         //these are all default values
	//ftpcMinZ(-39.5),        //change with set-functions individually
	//ftpcMaxZ(109.5),
	frBinCount(26),
	fzBinCount(75),
	fangle(0.17453),              //TODO: parameter management!
	fionDriftVelocity(1.766e-6),  // [cm/ns]
	ftime(0),               
	ferrorCount(0),
	fprimChargeOnly(false),
	fALICEmode(false),	 
        fpointBranchName("TpcPoint"),
	fPrimClBranchName("TpcPrimaryCluster"),
	fPrimClArray(NULL),
	fEpsilon(4),
	fEvtCounter(0)

{}

TpcSpaceChargeTask::~TpcSpaceChargeTask()
{}

InitStatus
TpcSpaceChargeTask::Init()
{

  //Get ROOT Manager
  FairRootManager* ioman= FairRootManager::Instance();
  
  if(ioman==0) {
    Error("TpcSpaceChargeTask::Init","RootManager not instantiated!");
    return kERROR;
  }

  // Get input collection
  fpointArray=(TClonesArray*) ioman->GetObject(fpointBranchName);
  if(fpointArray==NULL && (!fALICEmode)) {
    Error("TpcSpaceChargeTask::Init","MC-Point-array not found!");
    return kERROR;
  }

  fPrimClArray = (TClonesArray*) ioman->GetObject(fPrimClBranchName);
  if(fPrimClArray==NULL && fALICEmode) {
    Error("TpcSpaceChargeTask::Init","PrimaryCluster-array not found!");
    return kERROR;
  }
  
  
  fsupression=fpar->getSupression();
  ftpcMaxR=fpar->getRMax();  //make sure these are the values for the ACTIVE detector region  
  ftpcMinR=fpar->getRMin();  //see geometry/tpc.geo      
  ftpcMaxZ=fpar->windowMax();
  ftpcMinZ=fpar->windowMin();
  
  //fgemGain=fpar->getGain();
  frate=fpar->getRate();              
  
  fgas=fpar->getGas();
  fWGas = fgas->W();
  
  //fgemCharge = fsupression * fgemGain; //this should be epsilon
  fdistPerTime = fionDriftVelocity / frate; 	//Ion travel distance per event

  frBinWidth = ((ftpcMaxR - ftpcMinR) / frBinCount);
  fzBinWidth = ((ftpcMaxZ - ftpcMinZ) / fzBinCount);

  std::cout<<"\n\n--------- TpcSpaceChargeTask working parameters ----------"<<std::endl;
  std::cout<<"\nEpsilon: "<<fEpsilon
	   <<";   Rate: "<<frate<<"   WGas: "<<fWGas<<std::endl;
  std::cout<<"Geometry:\n"<<"MinR: "<<ftpcMinR<<";   MaxR: "<<ftpcMaxR<<std::endl;
  std::cout<<"MinZ: "<<ftpcMinZ<<";   MaxZ: "<<ftpcMaxZ<<std::endl;
  std::cout<<"-------------------------------------------------------------\n\n"<<std::endl;
  std::cout.flush();  
  

  // Initialize the charge-map (empty at the start...)
  for (int nr = 0; nr < frBinCount; nr++)
    fchargeMap.push_back(std::vector<int>(fzBinCount, 0));  

  return kSUCCESS;
}

void 
TpcSpaceChargeTask::SetParContainers() {

  std::cout<<"TpcSpaceChargeTask::SetParContainers"<<std::endl;
  std::cout.flush();

  // Get run and runtime database
  FairRun* run = FairRun::Instance();
  if ( ! run ) Fatal("SetParContainers", "No analysis run");

  FairRuntimeDb* db = run->GetRuntimeDb();
  if ( ! db ) Fatal("SetParContainers", "No runtime database");

  // Get Tpc digitisation parameter container
  fpar= (TpcDigiPar*) db->getContainer("TpcDigiPar");
  if (! fpar ) Fatal("SetParContainers", "TpcDigiPar not found");
}

void
TpcSpaceChargeTask::Exec(Option_t* opt)
{
  std::cout<<"TpcSpaceChargeTask::Exec() on Event Number "<<fEvtCounter
	   <<std::endl;
  fEvtCounter++;
  // //only needed for ALICE mode
  // const Float_t poti = 20.77e-9; // first ionization potential for Ne/CO2
  // const Float_t w_ion = 35.97e-9; // energy for the ion-electron pair creation 
  
  if(fALICEmode) { //trivial charge conversion:
    unsigned int nCh = fPrimClArray->GetEntriesFast();
    std::cout<<"TpcSpaceChargeTask::Exec() : Converting "
	     <<nCh<<" PrimaryClusters ... "<<std::endl;
    
    
    for(unsigned int ich=0; ich<nCh; ich++) {
      TpcPrimaryCluster* icl = (TpcPrimaryCluster*)((*fPrimClArray)[ich]);
      if(icl==NULL) {
	std::cout<<"ERROR: Invalid TpcPrimaryCluster pointer, skipping"<<std::endl;
	continue;
      }
      TVector3 pos = icl->pos();
      double posZ = pos.Z();
      double r = pos.Perp();
      bool outside = false;
      if( r<ftpcMinR || r> ftpcMaxR ) 
	outside=true;      
      if( posZ<ftpcMinZ || posZ>ftpcMaxZ)
	outside = true; 
      if(outside) {
	TString warn("Hit occured outside the Tpc-Volume! (");
	char posbuf[10];
	sprintf(posbuf,"%.4f",pos.X());
	warn.Append(posbuf);
	warn+=",";
	sprintf(posbuf,"%.4f",pos.Y());
	warn.Append(posbuf);
	warn+=",";
	sprintf(posbuf,"%.4f",pos.Z());
	warn.Append(posbuf);
	warn+=")";
	Warning("TpcSpaceChargeTask::Exec", warn.Data());  
	ferrorCount++;
	continue;
      }
      
      int rBin = (int)(std::fabs(r - ftpcMinR) / frBinWidth);
      int zBin = (int)(std::fabs(posZ - ftpcMinZ) / fzBinWidth);
    
      //Now fill in the point-charge into the charge-map
      (fchargeMap[rBin])[zBin] += icl->q();
      //At the same x-y-coordinates Ions are created on the gem-plane by the 
      //impact of the drifted electrons (assuming instant electron drift) in 
      //the gems
      if(!fprimChargeOnly)
	(fchargeMap[rBin])[0] += icl->q()*fEpsilon;
    }
    return;
  }
  else {
    TpcPoint* currentPoint;
    int np = fpointArray->GetEntriesFast();	//number of MC points 
    
    for(int n=0; n<np; n++) {			//Loop over all MC points
      
      if(ftime%500==0) {
	std::cout<<". ";
	std::cout.flush();
      }
      currentPoint = (TpcPoint*) fpointArray->At(n);
      int qPoint=0;
      
      qPoint = (int)(currentPoint->GetEnergyLoss()*1e9 / fWGas); // [eV]
      
      double posX = currentPoint->GetX();
      double posY = currentPoint->GetY();
      double posZ = currentPoint->GetZ();
      double r = sqrt(posX*posX + posY*posY);
      //check for boundaries:
      bool outside = false;
      if( r<ftpcMinR || r> ftpcMaxR ) 
	outside=true;      
      if( posZ<ftpcMinZ || posZ>ftpcMaxZ)
	outside = true; 
      if(outside) {
	TString warn("Hit occured outside the Tpc-Volume! (");
	char pos[10];
	sprintf(pos,"%.4f",posX);
	warn.Append(pos);
	warn+=",";
	sprintf(pos,"%.4f",posY);
	warn.Append(pos);
	warn+=",";
	sprintf(pos,"%.4f",posZ);
	warn.Append(pos);
	warn+=")";
	Warning("TpcSpaceChargeTask::Exec", warn.Data());  		
	ferrorCount++;
	continue;
      }
      
      int rBin = (int)(std::fabs(r - ftpcMinR) / frBinWidth);
      int zBin = (int)(std::fabs(posZ - ftpcMinZ) / fzBinWidth);
      
      //Now fill in the point-charge into the charge-map
      fchargeMap.at(rBin).at(zBin) += qPoint;
      
      //At the same x-y-coordinates Ions are created on the gem-plane by the 
      //impact of the drifted electrons (assuming instant electron drift) in 
      //the gems
      if(!fprimChargeOnly)
	fchargeMap.at(rBin).at(0) += (int)(qPoint*fEpsilon);
      
    } //end Loop over all Points of this track/event
    
    ftime++;
    
    //Drift charge one bin in z-direction when enough "time" has passed
    //This takes a lot of time. Normally not used, instead done by a macro
    
    if( fmod((double)ftime * fdistPerTime, fzBinWidth) < fdistPerTime)
      {
    	std::cout<<"\n**** Drifting all the charge after  "
    		 <<ftime<<" events ****"
    		 <<std::endl;
	
    	for(int nr=0; nr<frBinCount; nr++)
    	  for(int nz=fzBinCount-1; nz>0; nz--)
    	    fchargeMap.at(nr).at(nz) = fchargeMap.at(nr).at(nz-1);
	
    	for(int nr=0; nr<frBinCount; nr++)
    	  fchargeMap.at(nr).at(0) = 0;
      }
    return;
  }
}

void TpcSpaceChargeTask::writeToFile(const char* filename)
{  
  
  std::cout<< "\n\n**** Writing Charge Density to File " << filename 
           << " ****" << std::endl;
  double innerR, outerR, xy_area, volume;
  double pi = std::acos( -1.0 );
  double e_charge = 1.602176462e-19;

  //calculate the charge density from fchargeMap and geometry
  std::vector< std::vector <double> > fchargeDens;
  for (int nr=0; nr<frBinCount; nr++)
    fchargeDens.push_back(std::vector<double>(fzBinCount, 0.));

  for (int nr=0; nr<frBinCount; nr++)
  {
    innerR = nr * frBinWidth + ftpcMinR;	//inner radius of the bin
    outerR = (nr + 1) * frBinWidth + ftpcMinR;
    xy_area = (pi - fangle) * (outerR*outerR - innerR*innerR);
    volume = xy_area*fzBinWidth;
    for (int nz=0; nz<fzBinCount; nz++) {
      double charge = (double) (fchargeMap[nr])[nz]*e_charge/volume;  //[C]
      (fchargeDens[nr])[nz] = charge;
    }
  }  
  
  //write file with 2D charge distribution density
 
  std::ofstream outfile(filename, std::fstream::out);
  outfile<<std::setprecision(6);
  outfile<<ftpcMinR<<" "<<frBinWidth<<" "<<frBinCount<<" "
  	 <<ftpcMinZ<<" "<<fzBinWidth<<" "<<fzBinCount;
  //z runs first
  for (int nr=0; nr<frBinCount; nr++)
    for (int nz=0; nz<fzBinCount; nz++)
      outfile<<"\n"<<fchargeDens.at(nr).at(nz);
  outfile<<std::endl;
  outfile.close();
  std::cout<<"\nAll done. There have been "<<ferrorCount<<" hits outside the volume."<<std::endl;
}

ClassImp(TpcSpaceChargeTask)