import java.util.*;
import java.io.*;

import hep.lcd.util.driver.AbstractProcessor;
import hep.lcd.util.swim.Helix;
import  hep.lcd.util.collections.*;

import hep.lcd.event.LCDEvent;
import hep.lcd.event.Track;
import hep.lcd.event.MCParticle;
import hep.lcd.event.ClusterList;
import hep.lcd.event.Cluster;
import hep.lcd.event.CalorimeterHit;
import hep.lcd.event.CalorimeterHits;

import hep.physics.BasicHepLorentzVector;
 
import hep.lcd.geometry.Detector;
import hep.lcd.geometry.PropertySet;
import hep.lcd.geometry.CalorimeterCell;
import hep.lcd.geometry.component.Calorimeter;
import hep.lcd.geometry.component.EM_Barrel;
import hep.lcd.geometry.component.HAD_Barrel;
import hep.lcd.geometry.component.Field_Strength;

import hep.analysis.partition.FixedPartition2D;
import hep.analysis.Histogram;
import hep.analysis.Histogram2D;
import hep.analysis.Partition;

/** 
* Simple example program for extrapolating a reconstructed track through the 
* calorimeters and plotting distributions of hit energy around the track. 
* This forms the basis of the energy flow algorithm which was described 
* at LCWS00 
* http://conferences.fnal.gov/lcws2000/web/norman_graf_clustering/P014.html
*
* Plan of attack is as follows:
*
* 1.) Generate sufficient samples of single charged particles to populate 
*     the geometric and momentum phase space.
*
* 2.) Propagate tracks through the calorimeters and populate the histograms as 
*     here. Refinements include correct accounting for energy loss and 
*     accounting for where the shower actually starts (i.e. number of layers
*     (interaction lengths) since start of shower, not start of calorimeter).
*
* 2.) Analyze energy deposition in calorimeters as a function of particle 
*     type, energy, momentum and polar angle. I.e. run program such as this, 
*     determine correct functional form of the distributions, fit the 2D 
*     histograms and parameterize the fits appropriately.
*
* 3.) In reconstruction, will assume tracks are pions, so one question to 
*     be answered is how much of a systematic error this introduces.
*
* 4.) Propagate found tracks through the calorimeters. At each layer (or layer 
*     since shower start) loop over hit cells in a vicinity of the track 
*     intersection at that layer of the calorimeter and calculate a 
*     chi-squared for a cell with that energy to have originated from the 
*     track in question. At the end of this loop over tracks, associate cells 
*     with the best track (or leave unassociated if above some cut on 
*     the hypothesis chi-squared). Note that some sanity checks can be 
*     introduced at this stage. For instance, we know what the intrinsic 
*     energy resolution for single charged particles is. Therefore, 
*     we can take appropriate actions (e.g. stop associating cells to a track)
*     if the E/p ratio is excessive.
*
* 5.) Associated cells are removed from the list of cells to be considered for
*     neutral hadron identification. This should make it almost trivial to identify
*     them. Recall that we do not have to do a terrific job in terms of either
*     energy or position resolution at this stage.
*
* 6.) If you insist on clustering at an earlier stage, simply go to step one and
*     replace cell with cluster. Depending on the complexity of the clustering
*     and split/merge algorithms, the number of hit calorimeter cells and the 
*     number of tracks in an event, it may reduce the number of times one loops 
*     over hit cells in the calorimeters.
*     (Note that if both cell and cluster implemented a common interface the code
*     to do all this here would be the same. OO work in progress.)
*
* 7.) This is a work in progress. Please direct questions, comments, criticisms
*     to the author.
*
*
*@author Norman A. Graf (Norman.Graf@slac.stanford.edu)
*@version 1.0
*2002-05-21
*
* Note that this example only deals with the central barrel calorimeters!
*
*/

public class TrackCellAssociationAnalyzer extends AbstractProcessor
{
	// The number of layers in the EM Calorimeter
	private int _numemlayers;
	// The number of layers in the Hadronic Calorimeter
	private  int _numhadlayers;
	// The EM sampling fraction
	private  double  _sfEMB;
	// The HAD sampling fraction
	private  double _sfHB;

	// The EM calorimeter hits, by layer
	private List[] _emLayerHits; //EM

	//The HAD calorimeter hits, by layer
	private List[] _hadLayerHits; //HAD

	// The radii of the central calorimeter layers
	private double[] _emRadii; //EM
	private double[] _hadRadii; //HAD

	// The swimmers.
	// The swimmers are crude and need refinement (a la the trf propagators)
	// but they are part of the existing LCD code distribution.
	// THIS WILL CHANGE.
	private Helix[] _emSwimmers; //EM
	private Helix[] _hadSwimmers; //HAD

	// Z extent of the central barrel calorimeters.
	private double _emZ; //EM
	private double _hadZ; //HAD

	// The central magnetic field strength
	private double _fieldStrength;
	
	// Number of events processed
	private int _nEvents;
	
	// The MC particle energy
	// Used when writing out data for ntuple analysis
	private double _mcEnergy;

	// Initialization flag (can't initialize before we have an LCDEvent and 
	// know what detector we are dealing with).
	private boolean _initialized=false; // defined for emphasis

	// Debugging will be with us for a while
	private boolean _debug=false;
	
 
	public void process(LCDEvent event)
	{
		// Initialize some things here
		if( !_initialized )
		{
			Detector det = event.getCurrentDetector();
			log("detector= "+det);
			_fieldStrength = ((Field_Strength)det.get("Field_Strength")).getField();
			EM_Barrel emb = ((EM_Barrel) det.getComponent("EM_Barrel"));
			if(_debug) log("EM Barrel= "+emb);
			_emZ = emb.getZInner();
			_numemlayers = emb.getLayers();
			int emThetaBin = emb.getThetaBins();
			int emPhiBin = emb.getPhiBins(); 

			// store CalorimeterHits by layer
			_emLayerHits = new ArrayList[_numemlayers]; 
			_emRadii = new double[_numemlayers];
			_emSwimmers = new Helix[_numemlayers];
			for (int i=0; i<_numemlayers; ++i)
			{
				_emLayerHits[i] = new ArrayList(); 
				_emRadii[i]=emb.getRadius(emThetaBin,0,i); // crude
				_emSwimmers[i] = new Helix(_fieldStrength, _emRadii[i], Math.abs(_emZ));
				// remove when we go to GEANT!
				_emSwimmers[i].setCurvatureFlip(false);
			}

			//
			HAD_Barrel hadb = ((HAD_Barrel) det.getComponent("HAD_Barrel"));
			if(_debug) log("HAD Barrel= "+hadb);
			_hadZ = hadb.getZInner();
			_numhadlayers = hadb.getLayers(); 
			int hadThetaBin = hadb.getThetaBins();
			int hadPhiBin = hadb.getPhiBins();

			_hadLayerHits = new ArrayList[_numhadlayers];
			_hadRadii = new double[_numhadlayers];
			_hadSwimmers = new Helix[_numhadlayers];
			for (int i=0; i<_numhadlayers; ++i)
			{
				_hadLayerHits[i] = new ArrayList(); 
				_hadRadii[i]=hadb.getRadius(emThetaBin,0,i); // crude
				_hadSwimmers[i] = new Helix(_fieldStrength, _hadRadii[i], Math.abs(_hadZ));
				// remove when we go to GEANT!
				_hadSwimmers[i].setCurvatureFlip(false);
			}

			// create some histograms. Fix the binning to the cell size.
			// EM
			double dphi = 2.*Math.PI/emPhiBin;
			double dtheta = Math.PI/emThetaBin;
			for(int i = 0; i<_numemlayers; ++i)
			{
				histogram("EM Layer"+i+" ThetavsPhi").setPartition
					(new FixedPartition2D(-20.*dtheta,20.*dtheta,40,-20.*dphi,20.*dphi,40));
			}

			// HAD
			dphi = 2.*Math.PI/hadPhiBin;
			dtheta = Math.PI/hadThetaBin;
			for(int i = 0; i<_numhadlayers; ++i)
			{
				histogram("HAD Layer"+i+" ThetavsPhi").setPartition
					(new FixedPartition2D(-20.*dtheta,20.*dtheta,40,-20.*dphi,20.*dphi,40));
			}

			PropertySet sf = getParameters("SamplingFractions");
			_sfEMB = sf.getDouble("sfEMB");
			_sfHB = sf.getDouble("sfHB");
			_initialized = true;
		}    

		// Create a map to hold hit cells keyed by their cell index (towerID)
		CalorimeterHitMap merge;

		// get the calorimeter hits
		CalorimeterCell calCell = event.getCalorimeterCell();

		// EM
		CalorimeterHits emHits = event.getEMCalorimeterHits();
		Enumeration emEnum = emHits.getHits();

		while (emEnum.hasMoreElements())
		{
			CalorimeterHit hit = (CalorimeterHit) emEnum.nextElement();
			calCell.setTowerID(hit.getTowerID());
			int layer = calCell.getLayer(); 
			_emLayerHits[layer].add(hit);
		}

		// HAD
		CalorimeterHits hadHits = event.getHADCalorimeterHits();
		Enumeration hadEnum = hadHits.getHits();

		while (hadEnum.hasMoreElements())
		{
			CalorimeterHit hit = (CalorimeterHit) hadEnum.nextElement();
			calCell.setTowerID(hit.getTowerID());
			int layer = calCell.getLayer();
			_hadLayerHits[layer].add(hit);
		}

		if (_debug)
		{
			for (int i=0; i<_numemlayers; ++i)
			{
				log("EM Layer "+i+" has "+_emLayerHits[i].size()+" hits");
			}		
			for (int i=0; i<_numhadlayers; ++i)
			{
				log("HAD Layer "+i+" has "+_hadLayerHits[i].size()+" hits");
			}		
		}

		// Tracks...
		for (Enumeration e = event.getTrackList().getTracks();
			e.hasMoreElements();)
		{
			Track t = (Track) e.nextElement();
			// get information about the MCParticle
			MCParticle mcpart = t.getMCParticle();
			if(mcpart!=null)
			{
			_mcEnergy += mcpart.getEnergy();
			_nEvents++;					 
			}
			//
			double[] momentum = t.getMomentum();
			double[] origin = t.getOrigin();
			int iq = t.getCharge();
			// swim through the EM calorimeter
			String calSystem = "EM_Barrel";
			for(int i=0; i<_numemlayers; ++i)
			{
				_emSwimmers[i].swim(momentum, origin, iq);
				if(_debug) log("EMSwimmer "+i+" status= "+_emSwimmers[i].getPlane());
				double[] intersection = _emSwimmers[i].getIntersect();
				if(_debug) log("r= "+_emRadii[i]);
				if(_debug) log("intersection= "+intersection[0]+" "+intersection[1]+" "+intersection[2]);
				double phi = Math.atan2(intersection[1],intersection[0]);
				if(phi<0) phi+=2*Math.PI;
				//
				double theta = Math.atan(_emRadii[i]/intersection[2]);
				if(theta<0) theta+=Math.PI;
				// Compare to the calorimeter hits at this layer
				// Meant for single particle events so just get all the hits
				// For full events will want to introduce mechanism to return hits
				// only in a specified neighborhood.
				for(int j=0; j<_emLayerHits[i].size(); ++j)
				{
					CalorimeterHit hit = (CalorimeterHit) _emLayerHits[i].get(j);
					calCell.setTowerID(hit.getTowerID());
					double cellPhi = calCell.getPhi();
					double cellTheta = calCell.getTheta();
					double cellEnergy = hit.getEnergy()/_sfEMB;
					if (_debug)
					{
						if (Math.abs(cellPhi-phi)>.1 || Math.abs(cellTheta-theta)>.1)
						{
							log("phi= "+phi+", theta= "+theta);
							log("CalorimeterHit:  phi= "+cellPhi+" theta="+cellTheta+" layer="+calCell.getLayer()+" "+cellEnergy);
						}
					}
					// Fill the histograms
					histogram("EM Layer"+calCell.getLayer()+" ThetavsPhi").fillW(cellTheta-theta,cellPhi-phi,cellEnergy);
				}
			}
			// swim through the HAD calorimeter
			// Yes, I know this is copy/pasted from above. Refactoring will fix this.
			for(int i=0; i<_numhadlayers; ++i)
			{
				_hadSwimmers[i].swim(momentum, origin, iq);
				if(_debug) log("HADSwimmer "+i+" status= "+_hadSwimmers[i].getPlane());
				double[] intersection = _hadSwimmers[i].getIntersect();
				if(_debug) log("r= "+_hadRadii[i]);
				if(_debug) log("intersection= "+intersection[0]+" "+intersection[1]+" "+intersection[2]);
				if(_debug) log("r intercept= "+Math.sqrt(intersection[0]*intersection[0]+intersection[1]*intersection[1]));
				double phi = Math.atan2(intersection[1],intersection[0]);
				if(phi<0) phi+=2*Math.PI;
				double theta = Math.atan(_hadRadii[i]/intersection[2]);
				if(theta<0) theta+=Math.PI;
				if(_debug) log("phi= "+phi+", theta= "+theta);
				// Compare to the calorimeter hits at this layer.
				// Meant for single particle events so just get all the hits
				// For full events will want to introduce mechanism to return hits
				// only in a specified neighborhood.
				for(int j=0; j<_hadLayerHits[i].size(); ++j)
				{
					CalorimeterHit hit = (CalorimeterHit) _hadLayerHits[i].get(j);
					calCell.setTowerID(hit.getTowerID());
					calCell.setTowerID(hit.getTowerID());
					double cellPhi = calCell.getPhi();
					double cellTheta = calCell.getTheta();
					double cellEnergy = hit.getEnergy()/_sfHB;
					if(_debug) log("CalorimeterHit:  phi= "+cellPhi+" theta="+cellTheta+" layer="+calCell.getLayer()+" "+cellEnergy);
					if (_debug)
					{
						if (Math.abs(cellPhi-phi)>.1 || Math.abs(cellTheta-theta)>.1)
						{
							log("phi= "+phi+", theta= "+theta);
							log("CalorimeterHit:  phi= "+cellPhi+" theta="+cellTheta+" layer="+calCell.getLayer()+" "+cellEnergy);
						}
					}
					// Fill the histograms
					histogram("HAD Layer"+calCell.getLayer()+" ThetavsPhi").fillW(cellTheta-theta,cellPhi-phi,cellEnergy);
				}
			}

			// done, so clean up
			// EM
			for(int i=0; i<_numemlayers; ++i)
			{
				_emLayerHits[i].clear();
			}
			// HAD
			for(int i=0; i<_numhadlayers; ++i)
			{
				_hadLayerHits[i].clear();
			}					
		}
		// end tracks
	}

	public void stop()
	{
		// Write to external file for use with Jas Tuple Explorer
		
      BufferedWriter bw = null;

      try {
	  	// Do we need to write the file header?
		    File file = new File("SingleParticleLayerFits.txt");
			if (!file.exists())
			{
			bw = new BufferedWriter(new FileWriter("SingleParticleLayerFits.txt", true));
	 bw.write("energy caltype layer amplitude x0 y0 sx sy rxy");
	 bw.newLine();
	 bw.flush();
	 bw.close();			
			}
     bw = new BufferedWriter(new FileWriter("SingleParticleLayerFits.txt", true));

		// Access the information in the histograms...
		for(int l=0; l<_numemlayers+_numhadlayers; ++l)
		{
			int layer=0;
			String type = "EM";
			if (l<_numemlayers)
			{
				layer=l;
			}
			else
			{
				layer=l-_numemlayers;
				type="HAD";
			}
			FixedPartition2D p = (FixedPartition2D) histogram(type+" Layer"+layer+" ThetavsPhi").getPartition();
			int nx = p.getNumberOfXBins();
			int ny = p.getNumberOfYBins();
			double xmin = p.getXMin();
			double xmax = p.getXMax();
			double ymin = p.getYMin();
			double ymax = p.getYMax();
			double dx = (xmax-xmin)/(double)nx;
			double dy = (ymax-ymin)/(double)ny;
			xmin += dx/2.;
			ymin += dy/2.;
			//log("nx= "+nx+" xmin= "+xmin+" xmax= "+xmax);
			//log("ny= "+ny+" ymin= "+ymin+" ymax= "+ymax); 
			int npoints=0;
			double[] xvals = new double[nx*ny];
			double[] yvals = new double[nx*ny];
			double[] zvals = new double[nx*ny];
			for (int i=0; i<nx; ++i)
			{
				for (int j=0; j<ny; ++j)
				{
					double val = p.getBin(i,j);
					if (val>0.)
					{
						double xval = xmin+i*dx;
						double yval = ymin+j*dy;
		                //log("i= "+i+" x= "+xval+", j= "+j+" y= "+yval+" val= "+val);
						xvals[npoints]=xval;
						yvals[npoints]=yval;
						zvals[npoints]=val;
						npoints++;
					}
				}
			}
			// now fit a 2D gaussian to the energy distribution in this layer
			// Breit-Wigner or exponential is probably better-suited, but let's
			// start somewhere.	
			double[] fitparams = LeastSquares2DGaussianFit.fit(xvals, yvals, zvals, npoints);
			log(type+" Layer= "+layer);
			log(" amplitude= "+fitparams[0]);
			log(" x0= "+fitparams[1]);
			log(" y0= "+fitparams[2]);
			log(" sx= "+fitparams[3]);
			log(" sy= "+fitparams[4]);
			log(" rxy= "+fitparams[5]);
	 bw.write( _mcEnergy/((double)_nEvents)+" "+type+" "+layer+" "+fitparams[0]+" "+fitparams[1]+" "+fitparams[2]+" "+fitparams[3]+" "+fitparams[4]+" "+fitparams[5]);
	 bw.newLine();
	 bw.flush();			
		}

      } catch (IOException ioe) {
	 ioe.printStackTrace();
      } finally {                       // always close the file
	 if (bw != null) try {
	    bw.close();
	 } catch (IOException ioe2) {
	    // just ignore it
		ioe2.printStackTrace();
	 }
      } // end try/catch/finally		
		
	}

	static void log(String s)
	{
		System.out.println(s);
	}
}