import hep.analysis.*;
import hep.physics.*;
import hep.lcd.event.*;
import hep.lcd.mc.fast.*;
import hep.lcd.util.driver.*;
import java.util.*;

public class HiggsAnalysis extends Driver
{
	public HiggsAnalysis()
	{
		add(new MCFast()); // Find tracks
		add(new HiggsSearch()); // Search for Higgs
	}
}
class HiggsSearch extends Driver
{
	// Some useful constants
	private final static double mZ0 = 91.1882; // The mass of the Z particle
	private final static double energyCM = 500; // Center of mass energy of collider
	// We will use these variables to store the candidate tracks.
	private Track z1, z2;
	
	public void process(LCDEvent event)
    {
		// We will use these to keep a list of found muons and electrons
		Vector muons = new Vector();
		Vector electrons = new Vector();
		
		// Loop over tracks
		Enumeration e = event.getTrackList().getTracks();
		while (e.hasMoreElements())
		{
			Track t = (Track) e.nextElement(); 
			// Make some cuts on the tracks. We are only interested in high-energy
			// tracks that would be well reconstructed so reject tracks that have
			//            momentum < 10 (Gev)
			//            abs(cosTheta) > 0.9
			
			double cosTheta = t.cosTheta();
			double ptot = Math.sqrt(t.getPX()*t.getPX() + t.getPY()*t.getPY() + t.getPZ()*t.getPZ());
		
			if (/* not good*/) continue;
			
			// We are only interested in elecrons or muons from a Z. Use the 
			// getTrackType method to identify the tracks. Reject any unidentified tracks. 
			// Add found electrons and muons to the Lists.			
			String trackType = getTrackType(t);
			histogram("Track Type").fill(trackType);
			// *** Add to the appropriate list (use Vector.add() method)
			
		}
		histogram("Muon N").fill(muons.size());
		histogram("Electron N").fill(electrons.size());
		
		// Now loop over the muons and electrons and find the pair which gives an invariant
		// mass closest to the Z. Since we will need to do this separately for electrons and
		// muons we have defined a method, findZee, and call it in both cases. The body
		// of the findZee method is given below.
		
		double bestMatch = -1;
		bestMatch = findZee(muons,bestMatch);
		bestMatch = findZee(electrons,bestMatch);
		if (bestMatch > 0) histogram("Best Invmass").fill(bestMatch);
		
	    // We only want to accept this event if the best mass is consistent with a Z decay
		if (/* Is Good */)
		{
			// Calculate the recoil mass for the candidate tracks stored by fillZee
			double higgsMass = recoilMass(z1,z2);
			out.println("Found a Higgs "+higgsMass);
			// *** Histogram the result
		}
	}
	/**
	 * Find the pair of tracks in v with the invariant mass closest to that of the Z
	 * Only accept the pair if they are closer than prevBest.
	 * Store the best pair in member variables z1 and z2
	 * return the best mass (or prevBest if no better pair found)
	 */
	private double findZee(Vector v, double prevBest)
	{
		int l = v.size();
		for (int i=0; i<l-1 ; i++)
		{
			Track t1 = (Track) v.get(i);
			for (int j=i+1; j<l; j++)
			{
				Track t2 = (Track) v.get(j);
				double invMass = invMass(t1,t2);
				if (prevBest < 0 || Math.abs(prevBest-mZ0) > Math.abs(invMass-mZ0))
				{
					prevBest = invMass;
					this.z1 = t1;
					this.z2 = t2;
				}
			}
		}
		return prevBest;
	}
	private static String getTrackType(Track track)
	{
		// In real life we would use calorimetry/tracking dE/dx to identify
		// electrons and muons. Here we cheat and look at the MC truth info
		int pdgid = Math.abs(track.getMCParticle().getType().getPDGID());
		if      (pdgid == 11) return "electron";
		else if (pdgid == 13) return "muon";
		else return "unknown";
	}
	/**
	 * Calculate the invariant mass of the two tracks t1 and t2
	 */
	private static double invMass(Track t1, Track t2)
	{
		double ptot1 = Math.sqrt(t1.getPX()*t1.getPX() + t1.getPY()*t1.getPY() + t1.getPZ()*t1.getPZ());
		double ptot2 = Math.sqrt(t2.getPX()*t2.getPX() + t2.getPY()*t2.getPY() + t2.getPZ()*t2.getPZ());
		double etot = ptot1 + ptot2;
		double xtot = t1.getPX() + t2.getPX();
		double ytot = t1.getPY() + t2.getPY();
		double ztot = t1.getPZ() + t2.getPZ();
		return Math.sqrt(etot*etot - xtot*xtot - ytot*ytot - ztot*ztot);
	}
	/**
	 * Caculate the mass recoiling againts the two tracks t1 and t2
	 */
	private static double recoilMass(Track t1, Track t2)
	{
	    double ptot1 = Math.sqrt(t1.getPX()*t1.getPX() + t1.getPY()*t1.getPY() + t1.getPZ()*t1.getPZ());
		double ptot2 = Math.sqrt(t2.getPX()*t2.getPX() + t2.getPY()*t2.getPY() + t2.getPZ()*t2.getPZ());
		double etot = energyCM - ptot1 - ptot2;
		double xtot = t1.getPX() + t2.getPX();
		double ytot = t1.getPY() + t2.getPY();
		double ztot = t1.getPZ() + t2.getPZ();
		return Math.sqrt(etot*etot - xtot*xtot - ytot*ytot - ztot*ztot);
	}
}