import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import hep.physics.jet.JadeEJetFinder;
import hep.physics.jet.JetFinder;
import hep.physics.vec.HepLorentzVector;

import org.lcsim.event.EventHeader;
import org.lcsim.event.MCParticle;
import org.lcsim.util.Driver;

public class MyJetFindingDriver extends Driver {
    private double yCut = 0.001;

    public MyJetFindingDriver() {
    }

    public void process(EventHeader event) {
        List<MCParticle> particles = event.getMCParticles();
        Map<HepLorentzVector, MCParticle> map = new HashMap<HepLorentzVector, MCParticle>();
        for (MCParticle particle : particles) {
            if (particle.getCharge() != 0) {
                map.put(particle.asFourVector(), particle);
            }
        }
        JetFinder jetFind = new JadeEJetFinder(yCut);
        List<List<MCParticle>> jetList = new ArrayList<List<MCParticle>>();
        jetFind.setEvent(map.keySet());
        for (int i = 0; i < jetFind.njets(); ++i) {
            List<HepLorentzVector> jetVecList = jetFind.particlesInJet(i); 
            List<MCParticle> jetPartList = new ArrayList<MCParticle>();
            for (HepLorentzVector iVec : jetVecList) {
                jetPartList.add(map.get(iVec));
            }
            jetList.add(jetPartList);
        }
        event.put("jetList", jetList);
    }
}