#include "MipFinder.hh"
#include <vector>
#include <algorithm>
#include <EVENT/LCIO.h>
#include <EVENT/LCCollection.h>
#include <EVENT/CalorimeterHit.h>
#include <IMPL/CalorimeterHitImpl.h>
#include <IMPL/ClusterImpl.h>
#include <IMPL/LCCollectionVec.h>
#include <NVector_t.hh>
#include <values.h>
#include <LinearRegression.hh>
#include <marlin/ConditionsProcessor.h>
#include <cassert>
#include <UTIL/CellIDDecoder.h>

#include <cmath>

namespace CALICE {
  // create instance to make this Processor known to Marlin
  MipFinder a_MipFinder_instance;
  
  /** Mips finder .
   * Mip tracks are search within the calorimeter hits. 
   * Developped for cosmics.
   * Derived from the MipSelect processor inside calice_reco
   */
  MipFinder::MipFinder() : Processor("MipFinder")

  {
    _description = "Finds mips between two layers (i.e. straight clusters) and put them into a cluster collection." ;
    
    _nHits.push_back(3);
    _nHits.push_back(50);
    registerProcessorParameter( "NHits" , 
				"The minimum and the maximum number of hits per track for accepted tracks." ,
				_nHits ,
				_nHits);
    
    
    _maxDist=15.;
    registerProcessorParameter( "MaxDist" , 
				"The maximum allowed distance in x, y and z (mm) of hits of grouped to the same cluster." ,
				_maxDist ,
				_maxDist);
    
    _clusterEnergy.push_back(1.);
    _clusterEnergy.push_back(30*150);
    registerProcessorParameter( "ClusterEnergy" , 
				"Minimum and maximum total energy of accepted clusters." ,
				_clusterEnergy ,
				_clusterEnergy);
    
    _maxChi2=1000.;
    registerProcessorParameter( "MaxChi2" , 
				"Maximum chi^2/n.d.f. of line fits to the x and y coordinates." ,
				_maxChi2 ,
				_maxChi2);
    
    
    _w0=4.8;
    registerProcessorParameter( "WeightCutOff" , 
				"Cut parameter which defines the energy threshold of hits considered in the centre-of-gravity determination." ,
				_w0 ,
			       _w0);
    
    _thetaCut=0.1;
    registerProcessorParameter( "ThetaCut" , 
				"Cut parameter which defines the angle of the cone to add a linear box to another one." ,
				_thetaCut ,
				_thetaCut);
    
    
    
    _posError.clear();
    _posError.push_back(1.5);
    _posError.push_back(1.5);
    _posError.push_back(.5);
    registerProcessorParameter( "ClusterPositionError" , 
				"The error on the cluster position." ,
				_posError ,
				_posError,
				_posError.size());
    

    
    registerInputCollection( LCIO::CALORIMETERHIT , 
			     "CalorimeterHitCollection",
			     "Name of the Calorimeter hit collection"  ,
			     _colName ,
			     std::string("ConvCalorimeterHits") ) ;
    
    registerOutputCollection( LCIO::CLUSTER , 
			      "ClusterCollectionName" , 
			      "Name of the Cluster collection"  ,
			      _clusterColName ,
			      std::string("EcalClusters") ) ;
    
    registerProcessorParameter( "InferiorLayer" , 
				"Inferior layer to take hits" ,
				_layerMin ,
				_layerMin);
    
    registerProcessorParameter( "UpperLayer" , 
				"Upper layer to take hits" ,
				_layerMax ,
				_layerMax);

  }
  
  void MipFinder::init() {
    printParameters();

  }

  void MipFinder::processEvent( LCEvent * evtP ) {
    if (evtP) {
      try {
	
	LCCollection* col_hits = evtP->getCollection( _colName ) ;
	CellIDDecoder<CalorimeterHit> cd(col_hits);
	Bool_t hitLayers[30] ;
	Float_t posLayers[30] ;
	for ( int i = 0 ; i < 30 ; i++ ) { hitLayers[i] = false ; posLayers[i] = 0 ;}
	_cluster.clear() ;
	
	for ( UInt_t hit_i = 0 ; hit_i < (UInt_t) col_hits->getNumberOfElements() ; hit_i++ )
	  // just a countable loop to be sure it will stop
	  {
	    // create a collection of hits to be used by the algorithm
	    IMPL::LCCollectionVec *col_toBeDepleted = new IMPL::LCCollectionVec ( EVENT::LCIO::CALORIMETERHIT ) ;
	    col_toBeDepleted->parameters().setValue("CellIDEncoding",col_hits->getParameters().getStringVal("CellIDEncoding")) ;
	    
	    for ( UInt_t hit_j = 0 ; hit_j < (UInt_t) col_hits->getNumberOfElements() ; hit_j++ )
	      // loop over all hits and add to the collection those which are not in _cluster
	      // _cluster contains Boxes identified as Mips
	      {
		CalorimeterHit *a_hit = dynamic_cast<CalorimeterHit*>( col_hits->getElementAt( hit_j ));
		hitLayers[ cd(a_hit)["K-1"] ] = true ;
		posLayers[ cd(a_hit)["K-1"] ] = a_hit->getPosition()[2] ;

		bool isInCluster = false ;
		for ( std::vector<Box>::iterator cluster_iter=_cluster.begin(); cluster_iter!=_cluster.end(); cluster_iter++ )
		  {
		    for ( std::vector<const EVENT::CalorimeterHit*>::iterator hit_iter=cluster_iter->_hits.begin() ; hit_iter!=cluster_iter->_hits.end() ; hit_iter++ )
		      {
			if ( (*hit_iter)->getPosition()[0] == a_hit->getPosition()[0] && (*hit_iter)->getPosition()[1] == a_hit->getPosition()[1] && (*hit_iter)->getPosition()[2] == a_hit->getPosition()[2] && (*hit_iter)->getEnergy() == a_hit->getEnergy() ) { isInCluster = true ;}
		      }
		  }
		if ( !isInCluster )
		  {
		    //IMPL::CalorimeterHitImpl *new_hit = new IMPL::CalorimeterHitImpl() ;
		    //new_hit->setPosition(a_hit->getPosition()) ;
		    //new_hit->setEnergy(a_hit->getEnergy()) ;
		    //new_hit->setCellID0(a_hit->getCellID0()) ;
		    col_toBeDepleted->addElement(a_hit) ;
		  }
		
	      }// col_toBeDepleted filled
	    
	    
	    // apply the MipFinding procedure to col_toBeDepleted's hits
	    groupHits( col_toBeDepleted );
	    // _temp is filled

	    // check the maximum number of hits in the boxes of _temp
	    int nbHitsMaxInBox = 0 ;
	    for ( std::vector<Box>::iterator temp_iter=_temp.begin(); temp_iter!=_temp.end(); temp_iter++ )
	      {
		int nbHitsInBox = temp_iter->_hits.size() ;
		if ( nbHitsInBox > nbHitsMaxInBox )
		  { nbHitsMaxInBox = nbHitsInBox ; }
	      }
	    // if > 2 then take the largest box and implement it in _cluster
	    if ( nbHitsMaxInBox > 2 )
	      {
		for ( std::vector<Box>::iterator temp_iter=_temp.begin(); temp_iter!=_temp.end(); temp_iter++ )
		  {
		    bool done = false ;
		    if ( temp_iter->_hits.size() == (UInt_t) nbHitsMaxInBox && !done )
		      {
			_cluster.push_back(Box()) ;
			_cluster.back().add(*temp_iter) ;
			done = true ;
		      }
		  }
	      }
	    // if not then just add the other boxes to _cluster and stop the loop
	    else
	      {
		for ( std::vector<Box>::iterator temp_iter=_temp.begin(); temp_iter!=_temp.end(); temp_iter++ )
		  {
		    _cluster.push_back(Box()) ;
		    _cluster.back().add(*temp_iter) ;
		  }
		hit_i = (UInt_t) col_hits->getNumberOfElements() ;
	      }
	  }

	
	// --- final step: create Cluster collection
	
        // first : eliminate the hits after _layerMax by copying _cluster to _temp but not these hits
        _temp.clear() ;
	
	int mark = 0 ;
	bool isFirstLayer = false ;
	int firstHitLayer = _layerMin ;
	int maxLayer = _layerMax ;
	for ( int layer = 0 ; layer < 30 ; layer++ )
	  {
	    if ( hitLayers[layer] )
	      {
		mark++ ; 
		if ( !isFirstLayer )
		  {
		    isFirstLayer = true ;
		    firstHitLayer = layer ;
		  }
	      }
	    if ( mark == _layerMax - _layerMin + 1 ) { maxLayer = layer ;  }
	  }
        if ( mark < _layerMax - _layerMin + 1 )
          {
            maxLayer = 29 ;
            isFirstLayer = false ;
            for ( int layer = _layerMin ; layer > 0 ; layer-- )
              {
                if ( hitLayers[layer] )
                  {
                    mark++ ;
                    if ( mark == _layerMax - _layerMin + 1 )
                      {
                        isFirstLayer = true ;
                        firstHitLayer = layer ;                                                                                                                   
                      }                                                                                                                                                           
                  }
              }
            if ( !isFirstLayer ) { firstHitLayer = 0 ; }
          }
	Float_t layerMinPos = posLayers[ firstHitLayer ] ;
	Float_t layerMaxPos = posLayers[ maxLayer ] ;
	
        for(std::vector<Box>::const_iterator cluster_iter=_cluster.begin(); cluster_iter!=_cluster.end(); cluster_iter++ )
          {
            _temp.push_back(Box()) ;
            for (std::vector<const EVENT::CalorimeterHit*>::const_iterator hit_iter=cluster_iter->_hits.begin(); hit_iter!=cluster_iter->_hits.end() ; hit_iter++ )
	      {
		if ( (*hit_iter)->getPosition()[2] >= layerMinPos && (*hit_iter)->getPosition()[2] <= layerMaxPos )
		  { _temp.back().add(*hit_iter) ; }
	      }
	  }
	
	
	// second : copy all the information in _temp in the new cluster                    
	IMPL::LCCollectionVec *col_cluster = new IMPL::LCCollectionVec( EVENT::LCIO::CLUSTER ) ;
	FloatVec hengneChi2;
	for(std::vector<Box>::const_iterator cluster_iter=_temp.begin(); cluster_iter!=_temp.end(); cluster_iter++) {
	  if (!cluster_iter->isMerged()) {
	    if (   cluster_iter->_hits.size() >= static_cast<UInt_t>(_nHits[0]) 
		   && cluster_iter->_hits.size() <  static_cast<UInt_t>(_nHits[1]) ) {
	      IMPL::ClusterImpl *a_cluster = new IMPL::ClusterImpl() ;
	      Double_t sum[3]={0.,0.,0.};
	      Double_t sum2[3]={0.,0.,0.};
	      Double_t sum_of_weights=0.;
	      Double_t total_energy=0.;
	      UInt_t   n_weights=0;
	      
	      LinearRegression linreg[2];
	      hengneChi2.push_back( (float) 0 );
	      
	      for (std::vector<const EVENT::CalorimeterHit *>::const_iterator hit_iter=cluster_iter->_hits.begin();
		   hit_iter != cluster_iter->_hits.end();
		   hit_iter++) {
		if ((*hit_iter)->getEnergy()>0.) {
		  total_energy+=(*hit_iter)->getEnergy();
		}
	      }
	      
	      for (std::vector<const EVENT::CalorimeterHit *>::const_iterator hit_iter=cluster_iter->_hits.begin();
		   hit_iter != cluster_iter->_hits.end();
		   hit_iter++) {

		if ((*hit_iter)->getEnergy()>0.) {
		  Double_t energy=(*hit_iter)->getEnergy();
		  Double_t weight=_w0+std::log(energy/total_energy);
		  if (weight>0) {
		    for (UInt_t coord_i=0; coord_i<3; coord_i++) {
		      sum[coord_i]+=(*hit_iter)->getPosition()[coord_i] * weight;
		      sum2[coord_i]+= sqr((*hit_iter)->getPosition()[coord_i]) * weight;
		    }
		    linreg[0].add((*hit_iter)->getPosition()[2],(*hit_iter)->getPosition()[0],weight);
		    linreg[1].add((*hit_iter)->getPosition()[2],(*hit_iter)->getPosition()[1],weight);
		    sum_of_weights+=weight;
		    n_weights++;
		  }
		}
		
		a_cluster->addHit(const_cast<EVENT::CalorimeterHit *>(*hit_iter),1.);
	      }
	      
	      if (   n_weights >= static_cast<UInt_t>(_nHits[0]) 
		     && n_weights <  static_cast<UInt_t>(_nHits[1]) ) {
		
		if (total_energy >= _clusterEnergy[0] && total_energy < _clusterEnergy[1]) {
		  Float_t pos[3];
		  EVENT::FloatVec pos_err;
		  pos_err.resize(6,0.);
		  linreg[0].calc();
		  linreg[1].calc();
		  
		  Double_t chi2=0;
		  Double_t sum_of_weights=0;
		  UInt_t   n_weights=0; // number of correct calorimeter hits in the cluster
		  for (std::vector<const EVENT::CalorimeterHit *>::const_iterator hit_iter=cluster_iter->_hits.begin();
		       hit_iter != cluster_iter->_hits.end();
		       hit_iter++) {
		    
		    Double_t energy=(*hit_iter)->getEnergy();
		    Double_t weight=_w0+std::log(energy/total_energy);
		    
		    if (energy>0 && weight>0) {
		      Double_t res_x=(*hit_iter)->getPosition()[0]-linreg[0].eval((*hit_iter)->getPosition()[2]);
		      Double_t res_y=(*hit_iter)->getPosition()[1]-linreg[1].eval((*hit_iter)->getPosition()[2]);
		      
		      chi2+=sqr(res_x/weight);
		      chi2+=sqr(res_y/weight);
		      sum_of_weights+=weight;
		      n_weights++;

		      hengneChi2.back() += sqr(res_x/10) + sqr(res_y/10) ; // 10 is the size of a Si cell
		    }
		  }
		  
		  if (n_weights>0) {
		    chi2 /= (sum_of_weights-sum_of_weights*2/n_weights);
		    hengneChi2.back() /= n_weights ;
		  }
		  
		  if (sum_of_weights>0. && chi2<_maxChi2) {
		    
		    for (UInt_t coord_i=0; coord_i<3; coord_i++) {
		      Double_t mean=sum[coord_i]/sum_of_weights;
		      sum2[coord_i]=sqrt((sum2[coord_i]-mean*sum[coord_i])/(sum_of_weights));
		      pos[coord_i]=mean;
		    }
		    
		    pos_err[1]=pos_err[2]=pos_err[4]=0.;
		    pos_err[0]=_posError[0];
		    pos_err[3]=_posError[1];
		    pos_err[5]=_posError[2];
		    a_cluster->setPosition(pos);
		    a_cluster->setPositionError(pos_err);
		    a_cluster->setEnergy(total_energy);
		  }
		}
	      }
	      col_cluster->addElement(a_cluster);
	      
	    }
	  }
	  
	}
	
	col_cluster->parameters().setValues("HengneChi2", hengneChi2 ) ;
	col_cluster->parameters().setValue("firstHitLayer", firstHitLayer );
	col_cluster->parameters().setValue("layerMax", maxLayer );
	evtP->addCollection(col_cluster, _clusterColName);
	
      }
  
      
      catch (lcio::DataNotAvailableException err) {
      }
      
    }
  }
  
  void MipFinder::groupHits(LCCollection * col_toBeDepleted ) {
    
    // first loop to check how many layers were not hit and add further one to keep the same number of hit layers
    int maxLayer = _layerMax ;
    bool hitlayers[30] ;
    for ( int j = 0 ; j < 30 ; j++ )
      { hitlayers[j] = false ; }
    for(UInt_t element_i=0; element_i < (UInt_t) col_toBeDepleted->getNumberOfElements(); element_i++) {
      CalorimeterHit *a_hit=dynamic_cast<CalorimeterHit *>( col_toBeDepleted->getElementAt(element_i) );
      
      CellIDDecoder<CalorimeterHit> cd(col_toBeDepleted);
      //    *<table>
      //    *<tr><td>M:3,<td> wafer row</tr>
      //    *<tr><td>S-1:3,<td> wafer column</tr>
      //    *<tr><td>I:9,<td>pad coloumn</tr>
      //    *<tr><td>J:9,<td>pad row</tr>
      //    *<tr><td>K-1:6<td>layer</tr>
      //    *</table>	  
      int layer = cd(a_hit)["K-1"];
      hitlayers[layer] = true ;
    }
    int mark = 0 ;
    int firstHitLayer = _layerMin ;
    bool isFirstLayer = false ;
    for ( int layer = _layerMin ; layer < 30 ; layer++ )
      {
	if ( hitlayers[layer] )
	  { 
	    mark++ ;
	    if ( !isFirstLayer )
	      {
		isFirstLayer = true ;
		firstHitLayer = layer ;
	      }
	  }
	if ( mark == _layerMax - _layerMin + 1 ) { maxLayer = layer ;  }
      }
    if ( mark < _layerMax - _layerMin + 1 )
      {
	maxLayer = 29 ;
	isFirstLayer = false ;
	for ( int layer = _layerMin ; layer > 0 ; layer-- )
	  {
	    if ( hitlayers[layer] )
	      {
		mark++ ;
		if ( mark == _layerMax - _layerMin + 1 )
		  {
		    isFirstLayer = true ;
		    firstHitLayer = layer ;
		  }
	      }
	  }
	if ( !isFirstLayer ) { firstHitLayer = 0 ; }
      }
    // redefinition of the new minimum and maximum layer done
    
    
    // --- first step group nearby hits
    _temp.clear() ;
    for(UInt_t element_i=0; element_i < (UInt_t) col_toBeDepleted->getNumberOfElements(); element_i++) {
      CalorimeterHit *a_hit=dynamic_cast<CalorimeterHit *>( col_toBeDepleted->getElementAt(element_i) );
      
      CellIDDecoder<CalorimeterHit> cd(col_toBeDepleted);
      int layer = cd(a_hit)["K-1"];
      if ( layer >= firstHitLayer && layer <= maxLayer ) {
	
#ifdef CAN_CAST_CALORIMETERHIT_GETPOSITION	
	const ThreeVector_t *pointP=reinterpret_cast<const ThreeVector_t *>(a_hit->getPosition());
#else
	ThreeVector_t temp;
	temp.set(a_hit->getPosition());
	const ThreeVector_t *pointP=&temp;
#endif
	
	Double_t min_distance=FLT_MAX;
	UInt_t min_cluster_index=_temp.size();
	UInt_t cluster_index=0;
	for(std::vector<Box>::const_iterator cluster_iter=_temp.begin(); cluster_iter!=_temp.end(); cluster_iter++, cluster_index++) {
	  if (cluster_iter->closerToEnvelopThan(_maxDist,pointP)) {
	    Double_t a_min_distance=cluster_iter->minDistanceToPoints(pointP);
	    if (a_min_distance<_maxDist) {
	      if (a_min_distance<min_distance) {
		// if the new hit *a_hit is close enough to the box
		// check if the new hit can fit in the linear regression given by the other hits in the Box and closer to the _maxDist of this line
		if ( cluster_iter->_hits.size() > 1 ) {
		  LinearRegression boxRegression[2];
		  for(std::vector<const EVENT::CalorimeterHit *>::const_iterator hits_iter=cluster_iter->_hits.begin(); hits_iter!=cluster_iter->_hits.end(); hits_iter++) {
		    boxRegression[0].add((*hits_iter)->getPosition()[2],(*hits_iter)->getPosition()[0]);
		    boxRegression[1].add((*hits_iter)->getPosition()[2],(*hits_iter)->getPosition()[1]);
		  }
		  boxRegression[0].calc();
		  boxRegression[1].calc();
		  Double_t res_x=(a_hit)->getPosition()[0] - boxRegression[0].eval((a_hit)->getPosition()[2]);
		  Double_t res_y=(a_hit)->getPosition()[1] - boxRegression[1].eval((a_hit)->getPosition()[2]);
		  Double_t dist=sqrt(sqr(res_x)+sqr(res_y));
		  
		  if ( dist < _maxDist ) { min_cluster_index=cluster_index; }
		}
		else if ( cluster_iter->_hits.size() == 1 ) { min_cluster_index=cluster_index; }
	      }
	    }
	  }
	}
	
	if (min_cluster_index<_temp.size()) {
	  _temp[min_cluster_index].add(a_hit);
	}
	else {
	  _temp.push_back(Box());
	  _temp.back().add(a_hit);
	}
      }
    }
    
    // --- second step merge nearby groups 
    //std::vector<Box> merged_cluster;
    for(std::vector<Box>::iterator cluster_iter=_temp.begin(); cluster_iter!=_temp.end(); cluster_iter++) {
      
      bool isInCone = false ;
      
      LinearRegression boxRegression[2];
      for ( Int_t i1 = 0 ; i1 < (Int_t) cluster_iter->_hits.size() ; i1++ )
	{
	  boxRegression[0].add( (cluster_iter->_hits[i1])->getPosition()[2] , (cluster_iter->_hits[i1])->getPosition()[0] );
	  boxRegression[1].add( (cluster_iter->_hits[i1])->getPosition()[2] , (cluster_iter->_hits[i1])->getPosition()[1] );
	}
      boxRegression[0].calc();
      boxRegression[1].calc();		
      
      std::vector<Box>::iterator iter=cluster_iter;
      iter++;
      for(; iter!=_temp.end(); iter++)
	{
	  if ( !cluster_iter->isMerged() && iter->closerToEnvelopThan(_maxDist,*cluster_iter))
	    {
	      LinearRegression box2Regression[2];
	      for ( Int_t i2 = 0 ; i2 < (Int_t) iter->_hits.size() ; i2++ )
		{
		  box2Regression[0].add( (iter->_hits[i2])->getPosition()[2] , (iter->_hits[i2])->getPosition()[0] );
		  box2Regression[1].add( (iter->_hits[i2])->getPosition()[2] , (iter->_hits[i2])->getPosition()[1] );
		}
	      box2Regression[0].calc();
	      box2Regression[1].calc();
	      
	      if ( sqr(atan(boxRegression[0].ascent())-atan(box2Regression[0].ascent())) < sqr(_thetaCut)
		   && sqr(atan(boxRegression[1].ascent())-atan(box2Regression[1].ascent())) < sqr(_thetaCut) )
		{ isInCone = true ; }
	      
	      // treat the special case in which in box is made of two hits in separate layers with different x or y (just 1 pixel difference)
	      if ( cluster_iter->_hits.size() == 2 )
		{
		  if ( (cluster_iter->_hits[0])->getPosition()[2]-(cluster_iter->_hits[1])->getPosition()[2] != 0
		       && ( ( (cluster_iter->_hits[0])->getPosition()[1]-(cluster_iter->_hits[1])->getPosition()[1] != 0 ) ||
			    ( sqr((cluster_iter->_hits[0])->getPosition()[0]-(cluster_iter->_hits[1])->getPosition()[0]) != 0 ) )
		       )
		    { isInCone = true ;}
		}
	      
	      if ( iter->_hits.size() == 2 )
		{
		  if ( (iter->_hits[0])->getPosition()[2]-(iter->_hits[1])->getPosition()[2] != 0
		       && ( ( (iter->_hits[0])->getPosition()[1]-(iter->_hits[1])->getPosition()[1] != 0 ) ||
			    ( sqr((iter->_hits[0])->getPosition()[0]-(iter->_hits[1])->getPosition()[0]) != 0 ) )
		       )
		    { isInCone = true ;}
		}
	      
	      
	      if ( isInCone && !cluster_iter->isMerged() )
		{
		  iter->add(*cluster_iter);
		  cluster_iter->setMerged();
		  //break;
		}
	    }
	}
    }
    
  }
  
  
  void MipFinder::end() {
  }
  

}