/*
    * PandoraPfoSamplingFractionAnalysislDriver
    *
    * $Id: PandoraPfoSamplingFractionAnalysislDriver.java,v 1.3 2012/02/10 15:26:26 grefe Exp $
    */
   package org.lcsim.cal.calib;
   
   import static java.lang.Math.sqrt;
   import hep.aida.ITree;
  
  import java.util.List;
  
  import org.lcsim.conditions.ConditionsManager;
  import org.lcsim.conditions.ConditionsSet;
  import org.lcsim.conditions.ConditionsManager.ConditionsSetNotFoundException;
  import org.lcsim.event.CalorimeterHit;
  import org.lcsim.event.Cluster;
  import org.lcsim.event.EventHeader;
  import org.lcsim.event.MCParticle;
  import org.lcsim.geometry.IDDecoder;
  import org.lcsim.recon.cluster.fixedcone.FixedConeClusterer;
  import org.lcsim.util.Driver;
  import org.lcsim.util.aida.AIDA;
  
  import Jama.Matrix;
  import java.util.HashMap;
  import java.util.Map;
  import org.lcsim.event.ReconstructedParticle;
  import org.lcsim.event.SimCalorimeterHit;
  
  /**
   *
   * @author Norman Graf
   */
  public class PandoraPfoSamplingFractionAnalysislDriver extends Driver
  {
  
      private ConditionsSet _cond;
      private CollectionManager _collectionmanager = CollectionManager.defaultInstance();
      // we will accumulate the raw energy values in three depths:
      // 1. Layers (0)1 through (20)21 of the EM calorimeter (note that if layer 0 is massless, SF==1.)
      // 2. last ten layers of the EM calorimeter
      // 3. the hadron calorimeter
      //
      private double[][] _acc = new double[3][3];
      private double[] _vec = new double[3];
      // let's use a clusterer to remove effects of calorimeter cells hit far, far away.
      // use the only cross-detector clusterer we have:
      private FixedConeClusterer _fcc;
      private AIDA aida = AIDA.defaultInstance();
      private ITree _tree;
      private boolean _initialized = false;
      private boolean _debug = false;
      private double[] _ecalLayering;
      boolean _useFirstLayer;
  
      boolean _isHcalDigital = false;
  
      /**
       * Creates a new instance of SamplingFractionAnalysisDriver
       */
      public PandoraPfoSamplingFractionAnalysislDriver()
      {
          _tree = aida.tree();
      }
  
      protected void process(EventHeader event)
      {
  
          if (!_initialized) {
              ConditionsManager mgr = ConditionsManager.defaultInstance();
              try {
                  _cond = mgr.getConditions("CalorimeterCalibration");
              } catch (ConditionsSetNotFoundException e) {
                  System.out.println("ConditionSet CalorimeterCalibration not found for detector " + mgr.getDetector());
                  System.out.println("Please check that this properties file exists for this detector ");
              }
  
              _ecalLayering = _cond.getDoubleArray("ECalLayering");
              _useFirstLayer = _cond.getDouble("IsFirstEmLayerSampling") == 1.;
  
              ConditionsSet hcalProperties = mgr.getConditions("SamplingFractions/HcalBarrel");
  
              _isHcalDigital = hcalProperties.getBoolean("digital");
              System.out.println("HCal is " + (_isHcalDigital == true ? "" : "not") + " read out digitally");
  
              System.out.println("initialized...");
  
              _initialized = true;
          }
  
          // organize the histogram tree by species and energy
          List<MCParticle> mcparts = event.getMCParticles();
          //MCParticle mcpart = mcparts.get(mcparts.size() - 1); // this only works if particle was generated by stdhep, does not work with GEANT gps
          MCParticle mcpart = null;
          // Look for the most energetic final state particle
          for (MCParticle myMCP : mcparts) {
              if (myMCP.getGeneratorStatus() == MCParticle.FINAL_STATE) {
                  if (mcpart == null) {
                     mcpart = myMCP;
                 } else if (mcpart.getEnergy() < myMCP.getEnergy()) {
                     mcpart = myMCP;
                 }
             }
         }
         if (mcpart != null) {
             String particleType = mcpart.getType().getName();
             double mcEnergy = mcpart.getEnergy();
             long mcIntegerEnergy = Math.round(mcEnergy);
             boolean meV = false;
             if (mcEnergy < .99) {
                 mcIntegerEnergy = Math.round(mcEnergy * 1000);
                 meV = true;
             }
 
         // to derive the sampling fractions, want access to the raw energies
             // create a map keyed on ID for the SimCalorimeterHits
             // use the ID of the CalorimeterHits in the clusters to fetch them back
             Map<Long, SimCalorimeterHit> simCalHitMap = new HashMap<Long, SimCalorimeterHit>();
 
             List<List<SimCalorimeterHit>> collections = event.get(SimCalorimeterHit.class);
             for (List<SimCalorimeterHit> collection : collections) {
                 for (SimCalorimeterHit hit : collection) {
                     simCalHitMap.put(hit.getCellID(), hit);
                 }
             }
 
             _tree.mkdirs(particleType);
             _tree.cd(particleType);
             _tree.mkdirs(mcIntegerEnergy + (meV ? "_MeV" : "_GeV"));
             _tree.cd(mcIntegerEnergy + (meV ? "_MeV" : "_GeV"));
 
             if (mcpart.getSimulatorStatus().isDecayedInCalorimeter()) {
                 List<ReconstructedParticle> rpCollection = event.get(ReconstructedParticle.class, "PandoraPFOCollection");
                 if (_debug) {
                     System.out.println("found " + rpCollection.size() + " ReconstructedParticles");
                 }
                 // System.out.println("Found "+rpCollection.size()+" ReconstructedParticles");
 //            aida.cloud1D("Number of ReconstructedParticles found").fill(rpCollection.size());
                 if (rpCollection.size() == 1) {
 
                     for (ReconstructedParticle p : rpCollection) {
                         if (_debug) {
                             System.out.println("  energy " + p.getEnergy());
                         }
                     }
                     // fetch the highest energy RP (usually the first)
                     ReconstructedParticle p = rpCollection.get(0);
                     List<Cluster> clusters = p.getClusters();
                     if (_debug) {
                         System.out.println("found " + clusters.size() + " clusters");
                     }
                     aida.histogram1D("number of found clusters", 10, -0.5, 9.5).fill(clusters.size());
                     // proceed only if we found a single cluster above threshold
                     if (clusters.size() > 0) {
                         int nHCalHits = 0;
                         for (Cluster c : clusters) {
                             aida.cloud1D("Highest energy cluster energy").fill(c.getEnergy());
                             aida.cloud1D("Highest energy cluster number of cells").fill(c.getCalorimeterHits().size());
 
                             double clusterEnergy = c.getEnergy();
                             double mcMass = mcpart.getType().getMass();
                             // subtract the mass to get kinetic energy...
                             double expectedEnergy = mcEnergy;
                             // In case of protons and neutrons we expect only the kinetic energy in the shower
                             if (mcpart.getPDGID() == 2212 || mcpart.getPDGID() == 2112) {
                                 expectedEnergy = mcEnergy - mcMass;
                             }
 //                System.out.println(mcpart.getType().getName()+" "+expectedEnergy);
                             aida.cloud1D("measured - predicted energy").fill(clusterEnergy - expectedEnergy);
 
                         // let's now break down the cluster by component.
                             // this analysis uses:
                             // 1.) first 20 EM layers
                             // 2.) next 10 EM layers
                             // 3.) Had layers
                             List<CalorimeterHit> hits = c.getCalorimeterHits();
                             double[] vals = new double[4];
                             double clusterEnergySum = 0.;
                             double rawEnergySum = 0.;
                             for (CalorimeterHit hit : hits) {
                                 long id = hit.getCellID();
                                 IDDecoder decoder = hit.getIDDecoder();
                                 decoder.setID(id);
                                 SimCalorimeterHit simHit = simCalHitMap.get(id);
                                 double rawEnergy = simHit.getRawEnergy();
                                 double correctedEnergy = hit.getCorrectedEnergy();
                                 int layer = decoder.getLayer();
                                 String detectorName = decoder.getSubdetector().getName();
 //                        System.out.println(detectorName);
                                 int caltype = 0;
                                 // FIXME Hard-coded name.
                                 if (detectorName.toUpperCase().startsWith("ECAL")) {
                                     for (int i = 1; i < _ecalLayering.length + 1; ++i) {
                                         if (layer >= _ecalLayering[i - 1]) {
                                             caltype = i - 1;
                                         }
                                     }
                                 }
                                 // FIXME Hard-coded name.
                                 if (detectorName.toUpperCase().startsWith("HCAL")) {
                                     caltype = 3;
                                     nHCalHits += 1;
                                 }
 //                        System.out.println("layer "+layer+" caltype "+caltype);
                                 clusterEnergy += hit.getCorrectedEnergy();
 //                            vals[caltype] += hit.getCorrectedEnergy();
                                 vals[caltype] += rawEnergy;
                             } // end of loop over hits in cluster
                             if (_isHcalDigital == true) {
                                 vals[3] = nHCalHits;
                             }
                             if (_debug) {
                                 System.out.println("vals: " + vals[0] + " " + vals[1] + " " + vals[2] + " " + vals[3]);
                             }
                         // set up linear least squares:
                             // expectedEnergy = a*E1 + b*E2 +c*E3
                             for (int j = 0; j < 3; ++j) {
                                 _vec[j] += expectedEnergy * vals[j + 1];
                                 for (int k = 0; k < 3; ++k) {
                                     _acc[j][k] += vals[j + 1] * vals[k + 1];
                                 }
                             }
                         }
                     } // end of cluster loop
                 } // end of one RP cut
 //
 ////            event.put("All Calorimeter Hits",allHits);
 ////            event.put("Hits To Cluster",hitsToCluster);
 //            event.put("Found Clusters", clusters);
 //            } // end of ReconstructedParticle loop
             }// end of check on decays outside tracker volume
         }
         _tree.cd("/");
     }
 
     protected void endOfData()
     {
         System.out.println("done! endOfData.");
         // calculate the sampling fractions...
         Matrix A = new Matrix(_acc, 3, 3);
         A.print(6, 4);
         Matrix b = new Matrix(3, 1);
         for (int i = 0; i < 3; ++i) {
             b.set(i, 0, _vec[i]);
         }
         b.print(6, 4);
         try {
             Matrix x = A.solve(b);
             x.print(6, 4);
             System.out.println("SamplingFractions: " + 1. / x.get(0, 0) + ", " + 1. / x.get(1, 0) + ", " + 1. / x.get(2, 0));//+ ", " + 1. / x.get(3, 0));
         } catch (Exception e) {
             e.printStackTrace();
 //            System.out.println("try reducing dimensionality...");
 //            Matrix Ap = new Matrix(_acc, 2, 2);
 //            Ap.print(6, 4);
 //            Matrix bp = new Matrix(2, 1);
 //            for (int i = 0; i < 2; ++i)
 //            {
 //                bp.set(i, 0, _vec[i]);
 //            }
 //            bp.print(6, 4);
 //            try
 //            {
 //                Matrix x = Ap.solve(bp);
 //                x.print(6, 4);
 //            } catch (Exception e2)
 //            {
 //                e2.printStackTrace();
 //            }
         }
     }
 
     public void setDebug(boolean debug)
     {
         _debug = debug;
     }
 }