PointIdAlg.cxx

Go to the documentation of this file.

// Class:       PointIdAlg
// Authors:     D.Stefan (Dorota.Stefan@ncbj.gov.pl),      from DUNE,   CERN/NCBJ, since May 2016
//              R.Sulej (Robert.Sulej@cern.ch),            from DUNE,   FNAL/NCBJ, since May 2016
//              P.Plonski,                                 from DUNE,   WUT,       since May 2016
//              D.Smith,                                   from LArIAT, BU, 2017: real data dump

#include "larreco/RecoAlg/ImagePatternAlgs/Tensorflow/PointIdAlg/PointIdAlg.h"
#include "tensorflow/core/public/session.h"

#include "art/Framework/Principal/Event.h"
#include "art/Framework/Principal/Handle.h"
#include "art/Framework/Services/Registry/ServiceHandle.h"

#include "larcorealg/Geometry/ChannelMapAlg.h" // geo::InvalidWireIDError
#include "larcore/CoreUtils/ServiceUtil.h" // lar::providerFrom<>()

#include "lardataobj/Simulation/SimChannel.h"
#include "larsim/Simulation/LArG4Parameters.h"

#include "larevt/CalibrationDBI/Interface/ChannelStatusService.h"
#include "larevt/CalibrationDBI/Interface/ChannelStatusProvider.h"

#include "messagefacility/MessageLogger/MessageLogger.h"

#include "CLHEP/Random/RandGauss.h"

#include <sys/stat.h>

// ------------------------------------------------------
// -------------------ModelInterface---------------------
// ------------------------------------------------------

std::vector< std::vector<float> > nnet::ModelInterface::Run(std::vector< std::vector< std::vector<float> > > const & inps, int samples)
{
    if ((samples == 0) || inps.empty() || inps.front().empty() || inps.front().front().empty())
        return std::vector< std::vector<float> >();

    if ((samples == -1) || (samples > (int)inps.size())) { samples = inps.size(); }

    std::vector< std::vector<float> > results;
    for (int i = 0; i < samples; ++i)
    {
        results.push_back(Run(inps[i]));
    }
    return results;
}


std::string nnet::ModelInterface::findFile(const char* fileName) const
{
    std::string fname_out;
    cet::search_path sp("FW_SEARCH_PATH");
    if (!sp.find_file(fileName, fname_out))
    {
        struct stat buffer;
        if (stat(fileName, &buffer) == 0) { fname_out = fileName; }
        else
        {
            throw art::Exception(art::errors::NotFound)
                << "Could not find the model file " << fileName;
        }
    }
    return fname_out;
}


// ------------------------------------------------------
// ----------------KerasModelInterface-------------------
// ------------------------------------------------------

nnet::KerasModelInterface::KerasModelInterface(const char* modelFileName) :
        m(nnet::ModelInterface::findFile(modelFileName).c_str())
{
        mf::LogInfo("KerasModelInterface") << "Keras model loaded.";
}
// ------------------------------------------------------

std::vector<float> nnet::KerasModelInterface::Run(std::vector< std::vector<float> > const & inp2d)
{
        std::vector< std::vector< std::vector<float> > > inp3d;
        inp3d.push_back(inp2d); // lots of copy, should add 2D to keras...

        keras::DataChunk *sample = new keras::DataChunk2D();
        sample->set_data(inp3d); // and more copy...
        std::vector<float> out = m.compute_output(sample);
        delete sample;
        return out;
}


// ------------------------------------------------------
// -----------------TfModelInterface---------------------
// ------------------------------------------------------

nnet::TfModelInterface::TfModelInterface(const char* modelFileName)
{
        g = tf::Graph::create(nnet::ModelInterface::findFile(modelFileName).c_str(), {"cnn_output", "_netout"});
        if (!g) { throw art::Exception(art::errors::Unknown) << "TF model failed."; }

        mf::LogInfo("TfModelInterface") << "TF model loaded.";
}
// ------------------------------------------------------

std::vector< std::vector<float> > nnet::TfModelInterface::Run(std::vector< std::vector< std::vector<float> > > const & inps, int samples)
{
    if ((samples == 0) || inps.empty() || inps.front().empty() || inps.front().front().empty())
        return std::vector< std::vector<float> >();

    if ((samples == -1) || (samples > (long long int)inps.size())) { samples = inps.size(); }

    long long int rows = inps.front().size(), cols = inps.front().front().size();

    tensorflow::Tensor _x(tensorflow::DT_FLOAT, tensorflow::TensorShape({ samples, rows, cols, 1 }));
    auto input_map = _x.tensor<float, 4>();
    for (long long int s = 0; s < samples; ++s) {
        const auto & sample = inps[s]; 
        for (long long int r = 0; r < rows; ++r) {
            const auto & row = sample[r]; 
            for (long long int c = 0; c < cols; ++c) {
                input_map(s, r, c, 0) = row[c];
            }
        }
    }

    return g->run(_x);
}
// ------------------------------------------------------

std::vector<float> nnet::TfModelInterface::Run(std::vector< std::vector<float> > const & inp2d)
{
    long long int rows = inp2d.size(), cols = inp2d.front().size();

    tensorflow::Tensor _x(tensorflow::DT_FLOAT, tensorflow::TensorShape({ 1, rows, cols, 1 }));
    auto input_map = _x.tensor<float, 4>();
    for (long long int r = 0; r < rows; ++r) {
        const auto & row = inp2d[r]; 
        for (long long int c = 0; c < cols; ++c) {
            input_map(0, r, c, 0) = row[c];
        }
    }

    auto out = g->run(_x);
    if (!out.empty()) return out.front();
    else return std::vector<float>();
}


// ------------------------------------------------------
// --------------------PointIdAlg------------------------
// ------------------------------------------------------

nnet::PointIdAlg::PointIdAlg(const Config& config) : img::DataProviderAlg(config),
        fNNet(0),
        fPatchSizeW(config.PatchSizeW()), fPatchSizeD(config.PatchSizeD()),
        fCurrentWireIdx(99999), fCurrentScaledDrift(99999)
{
        fNNetModelFilePath = config.NNetModelFile();
        fNNetOutputs = config.NNetOutputs();

        deleteNNet();

        if ((fNNetModelFilePath.length() > 5) &&
            (fNNetModelFilePath.compare(fNNetModelFilePath.length() - 5, 5, ".nnet") == 0))
        {
                fNNet = new nnet::KerasModelInterface(fNNetModelFilePath.c_str());
        }
        else if ((fNNetModelFilePath.length() > 3) &&
            (fNNetModelFilePath.compare(fNNetModelFilePath.length() - 3, 3, ".pb") == 0))
        {
                fNNet = new nnet::TfModelInterface(fNNetModelFilePath.c_str());
        }
        else
        {
                mf::LogError("PointIdAlg") << "File name extension not supported.";
        }

    if (!fNNet) { throw cet::exception("nnet::PointIdAlg") << "Loading model from file failed."; }

    resizePatch();
}
// ------------------------------------------------------

nnet::PointIdAlg::~PointIdAlg(void)
{
        deleteNNet();
}
// ------------------------------------------------------

void nnet::PointIdAlg::resizePatch(void)
{
        fWireDriftPatch.resize(fPatchSizeW);
        for (auto & r : fWireDriftPatch) r.resize(fPatchSizeD);
}
// ------------------------------------------------------

float nnet::PointIdAlg::predictIdValue(unsigned int wire, float drift, size_t outIdx) const
{
        float result = 0.;

        if (!bufferPatch(wire, drift))
        {
                mf::LogError("PointIdAlg") << "Patch buffering failed.";
                return result;
        }

        if (fNNet)
        {
                auto out = fNNet->Run(fWireDriftPatch);
                if (!out.empty()) { result = out[outIdx]; }
                else
                {
                        mf::LogError("PointIdAlg") << "Problem with applying model to input.";
                }
        }

        return result;
}
// ------------------------------------------------------

std::vector<float> nnet::PointIdAlg::predictIdVector(unsigned int wire, float drift) const
{
        std::vector<float> result;

        if (!bufferPatch(wire, drift))
        {
                mf::LogError("PointIdAlg") << "Patch buffering failed.";
                return result;
        }

        if (fNNet)
        {
        result = fNNet->Run(fWireDriftPatch);
        if (result.empty())
        {
            mf::LogError("PointIdAlg") << "Problem with applying model to input.";
        }
        }

        return result;
}
// ------------------------------------------------------

std::vector< std::vector<float> > nnet::PointIdAlg::predictIdVectors(std::vector< std::pair<unsigned int, float> > points) const
{
        if (points.empty() || !fNNet) { return std::vector< std::vector<float> >(); }

        std::vector< std::vector< std::vector<float> > > inps(
                points.size(), std::vector< std::vector<float> >(
                        fPatchSizeW, std::vector<float>(fPatchSizeD)));
        for (size_t i = 0; i < points.size(); ++i)
        {
                unsigned int wire = points[i].first;
                float drift = points[i].second;
                if (!bufferPatch(wire, drift, inps[i]))
                {
                        throw cet::exception("PointIdAlg") << "Patch buffering failed" << std::endl;
                }

        }

        return fNNet->Run(inps);
}
// ------------------------------------------------------

bool nnet::PointIdAlg::isSamePatch(unsigned int wire1, float drift1, unsigned int wire2, float drift2) const
{
    if (fDownscaleFullView)
    {
        size_t sd1 = (size_t)(drift1 / fDriftWindow);
        size_t sd2 = (size_t)(drift2 / fDriftWindow);
        if ((wire1 == wire2) && (sd1 == sd2))
                return true; // the same position
    }
    else
    {
        if ((wire1 == wire2) && ((size_t)drift1 == (size_t)drift2))
                return true; // the same position
    }

    return false; // not the same position
}

bool nnet::PointIdAlg::isCurrentPatch(unsigned int wire, float drift) const
{
    if (fDownscaleFullView)
    {
        size_t sd = (size_t)(drift / fDriftWindow);
        if ((fCurrentWireIdx == wire) && (fCurrentScaledDrift == sd))
                return true; // still within the current position
    }
    else
    {
        if ((fCurrentWireIdx == wire) && (fCurrentScaledDrift == drift))
                return true; // still within the current position
    }

    return false; // not a current position
}
// ------------------------------------------------------

std::vector<float> nnet::PointIdAlg::flattenData2D(std::vector< std::vector<float> > const & patch)
{
        std::vector<float> flat;
        if (patch.empty() || patch.front().empty())
        {
                mf::LogError("DataProviderAlg") << "Patch is empty.";
                return flat;
        }

        flat.resize(patch.size() * patch.front().size());

        for (size_t w = 0, i = 0; w < patch.size(); ++w)
        {
                auto const & wire = patch[w];
                for (size_t d = 0; d < wire.size(); ++d, ++i)
                {
                        flat[i] = wire[d];
                }
        }

        return flat;
}
// ------------------------------------------------------

bool nnet::PointIdAlg::isInsideFiducialRegion(unsigned int wire, float drift) const
{
        size_t marginW = fPatchSizeW / 8; // fPatchSizeX/2 will make patch always completely filled
        size_t marginD = fPatchSizeD / 8;

        size_t scaledDrift = (size_t)(drift / fDriftWindow);
        if ((wire >= marginW) && (wire < fNWires - marginW) &&
            (scaledDrift >= marginD) && (scaledDrift < fNScaledDrifts - marginD)) return true;
        else return false;
}
// ------------------------------------------------------

// ------------------------------------------------------
// ------------------TrainingDataAlg---------------------
// ------------------------------------------------------

nnet::TrainingDataAlg::TrainingDataAlg(const Config& config) : img::DataProviderAlg(config),
        fEdepTot(0),
        fWireProducerLabel(config.WireLabel()),
        fHitProducerLabel(config.HitLabel()),
        fTrackModuleLabel(config.TrackLabel()),
        fSimulationProducerLabel(config.SimulationLabel()),
        fSaveVtxFlags(config.SaveVtxFlags()),
    fAdcDelay(config.AdcDelayTicks()),
    fEventsPerBin(100, 0)
{
    fSaveSimInfo = !fSimulationProducerLabel.label().empty();
}
// ------------------------------------------------------

nnet::TrainingDataAlg::~TrainingDataAlg(void)
{
}
// ------------------------------------------------------

void nnet::TrainingDataAlg::resizeView(size_t wires, size_t drifts)
{
        img::DataProviderAlg::resizeView(wires, drifts);

        fWireDriftEdep.resize(wires);
        for (auto & w : fWireDriftEdep)
        {
                w.resize(fNCachedDrifts);
                std::fill(w.begin(), w.end(), 0.0F);
        }

        fWireDriftPdg.resize(wires);
        for (auto & w : fWireDriftPdg)
        {
                w.resize(fNCachedDrifts);
                std::fill(w.begin(), w.end(), 0);
        }
}
// ------------------------------------------------------

bool nnet::TrainingDataAlg::setWireEdepsAndLabels(
        std::vector<float> const & edeps, std::vector<int> const & pdgs, size_t wireIdx)
{
        if ((wireIdx >= fWireDriftEdep.size()) || (edeps.size() != pdgs.size())) { return false; }

        size_t dstep = 1;
        if (fDownscaleFullView) { dstep = fDriftWindow; }

        if (edeps.size() / dstep > fNCachedDrifts) { return false; }

        auto & wEdep = fWireDriftEdep[wireIdx];
        auto & wPdg = fWireDriftPdg[wireIdx];

        for (size_t i = 0; i < fNCachedDrifts; ++i)
        {
                size_t i0 = i * dstep;
                size_t i1 = (i + 1) * dstep;

                int best_pdg = pdgs[i0] & nnet::TrainingDataAlg::kPdgMask;
                int vtx_flags = pdgs[i0] & nnet::TrainingDataAlg::kVtxMask;
                int type_flags = pdgs[i0] & nnet::TrainingDataAlg::kTypeMask;
                float max_edep = edeps[i0];
                for (size_t k = i0 + 1; k < i1; ++k)
                {
                        float ek = edeps[k];
                        if (ek > max_edep)
                        {
                                max_edep = ek;
                                best_pdg = pdgs[k] & nnet::TrainingDataAlg::kPdgMask; // remember best matching pdg
                        }
                        type_flags |= pdgs[k] & nnet::TrainingDataAlg::kTypeMask; // accumulate track type flags
                        vtx_flags |= pdgs[k] & nnet::TrainingDataAlg::kVtxMask;   // accumulate all vtx flags
                }

                wEdep[i] = max_edep;

        best_pdg |= type_flags;
                if (fSaveVtxFlags) best_pdg |= vtx_flags;
                wPdg[i] = best_pdg;
        }

        return true;
}
// ------------------------------------------------------

nnet::TrainingDataAlg::WireDrift nnet::TrainingDataAlg::getProjection(const TLorentzVector& tvec, unsigned int plane) const
{
        auto const* detprop = lar::providerFrom<detinfo::DetectorPropertiesService>();
        nnet::TrainingDataAlg::WireDrift wd;
        wd.Wire = 0; wd.Drift = 0; wd.TPC = -1; wd.Cryo = -1;

    try
    {
        double vtx[3] = {tvec.X(), tvec.Y(), tvec.Z()};
            if (fGeometry->FindTPCAtPosition(vtx).isValid)
            {
                geo::TPCID tpcid = fGeometry->FindTPCAtPosition(vtx);
                unsigned int tpc = tpcid.TPC, cryo = tpcid.Cryostat;

                // correct for the time offset
                float dx = tvec.T() * 1.e-3 * detprop->DriftVelocity();
                int driftDir = fGeometry->TPC(tpcid).DetectDriftDirection();
                        if (driftDir == 1) { dx *= -1; }
                        else if (driftDir != -1)
                        {
                            throw cet::exception("nnet::TrainingDataAlg") << "drift direction is not X." << std::endl;
                        }
                        vtx[0] = tvec.X() + dx;
                                
                wd.Wire = fGeometry->NearestWire(vtx, plane, tpc, cryo);
                wd.Drift = fDetProp->ConvertXToTicks(vtx[0], plane, tpc, cryo);
                wd.TPC = tpc; wd.Cryo = cryo;
            }
        }
        catch (const geo::InvalidWireIDError & e)
        {
            mf::LogWarning("TrainingDataAlg") << "Vertex projection out of wire planes, just skipping this vertex.";
        }
        catch (...)
        {
            mf::LogWarning("TrainingDataAlg") << "Vertex projection out of wire planes, skip MC vertex.";
        }
        return wd;
}
// ------------------------------------------------------

bool nnet::TrainingDataAlg::isElectronEnd(const simb::MCParticle & particle,
    const std::unordered_map< int, const simb::MCParticle* > & particleMap) const
{
    const float minElectronLength2 = 2.5*2.5;
    const float maxDeltaLength2 = 0.15*0.15;

    int pdg = abs(particle.PdgCode());
        if (pdg != 11) return false; // should be applied only to electrons

        size_t nSec = particle.NumberDaughters();
        for (size_t d = 0; d < nSec; ++d)
        {
                auto d_search = particleMap.find(particle.Daughter(d));
                if (d_search != particleMap.end())
                {
                        auto const & daughter = *((*d_search).second);
                        int d_pdg = abs(daughter.PdgCode());
                        if (d_pdg != 22) { return false; } // not the end of the shower
                }
        }

    float trkLength2 = 0;
        auto const * p = &particle;
        bool branching = false;
        while (!branching)
        {
        trkLength2 += particleRange2(*p);
        auto m_search = particleMap.find(p->Mother());
                if (m_search != particleMap.end())
        {
                        p = (*m_search).second;
                        int m_pdg = abs(p->PdgCode());
                        if (m_pdg == 11)
                        {
                            nSec = p->NumberDaughters();
                            size_t ne = 0;
                            for (size_t d = 0; d < nSec; ++d)
                            {
                                auto d_search = particleMap.find(p->Daughter(d));
                                if (d_search != particleMap.end())
                                {
                                    auto const & daughter = *((*d_search).second);
                                    int d_pdg = abs(daughter.PdgCode());
                                    if (d_pdg == 11)
                                    {
                                if (particleRange2(daughter) > maxDeltaLength2) { ne++; }
                                    }
                                }
                            }
                            if (ne > 1) { branching = true; }
                        }
                        else break;
        }
        else break;
    }

    return (trkLength2 > minElectronLength2);
}

bool nnet::TrainingDataAlg::isMuonDecaying(const simb::MCParticle & particle,
    const std::unordered_map< int, const simb::MCParticle* > & particleMap) const
{
    bool hasElectron = false, hasNuMu = false, hasNuE = false;

    int pdg = abs(particle.PdgCode());
        if ((pdg == 13) && (particle.EndProcess() == "FastScintillation")) // potential muon decay at rest
        {
                unsigned int nSec = particle.NumberDaughters();
                for (size_t d = 0; d < nSec; ++d)
                {
                        auto d_search = particleMap.find(particle.Daughter(d));
                        if (d_search != particleMap.end())
                        {
                                auto const & daughter = *((*d_search).second);
                                int d_pdg = abs(daughter.PdgCode());
                                if (d_pdg == 11) hasElectron = true;
                                else if (d_pdg == 14) hasNuMu = true;
                                else if (d_pdg == 12) hasNuE = true;
                        }
                }
        }

        return (hasElectron && hasNuMu && hasNuE);
}

void nnet::TrainingDataAlg::collectVtxFlags(
        std::unordered_map< size_t, std::unordered_map< int, int > > & wireToDriftToVtxFlags,
        const std::unordered_map< int, const simb::MCParticle* > & particleMap,
        unsigned int plane) const
{
        for (auto const & p : particleMap)
        {
                auto const & particle = *p.second;

                double ekStart = 1000. * (particle.E() - particle.Mass());
                double ekEnd = 1000. * (particle.EndE() - particle.Mass());

                int pdg = abs(particle.PdgCode());
                int flagsStart = nnet::TrainingDataAlg::kNone;
                int flagsEnd = nnet::TrainingDataAlg::kNone;
        
                switch (pdg)
                {
                        case 22:   // gamma
                                if ((particle.EndProcess() == "conv") &&
                                    (ekStart > 40.0)) // conversion, gamma > 40MeV
                                {
                                        //std::cout << "---> gamma conversion at " << ekStart << std::endl;
                                        flagsEnd = nnet::TrainingDataAlg::kConv;
                                }
                                break;

                        case 11:   // e+/-
                            if (isElectronEnd(particle, particleMap))
                            {
                                flagsEnd = nnet::TrainingDataAlg::kElectronEnd;
                            }
                                break;


                        case 13:   // mu+/-
                            if (isMuonDecaying(particle, particleMap))
                            {
                                //std::cout << "---> mu decay to electron" << std::endl;
                                flagsEnd = nnet::TrainingDataAlg::kDecay;
                            }
                                break;

                        case 111:  // pi0
                                //std::cout << "---> pi0" << std::endl;
                                flagsStart = nnet::TrainingDataAlg::kPi0;
                                break;

                        case 321:  // K+/-
                        case 211:  // pi+/-
                        case 2212: // proton
                                if (ekStart > 50.0)
                                {
                                        if (particle.Mother() != 0)
                                        {
                                                auto search = particleMap.find(particle.Mother());
                                                if (search != particleMap.end())
                                                {
                                                        auto const & mother = *((*search).second);
                                                        int m_pdg = abs(mother.PdgCode());
                                                        unsigned int nSec = mother.NumberDaughters();
                                                        unsigned int nVisible = 0;
                                                        if (nSec > 1)
                                                        {
                                                                for (size_t d = 0; d < nSec; ++d)
                                                                {
                                                                        auto d_search = particleMap.find(mother.Daughter(d));
                                                                        if (d_search != particleMap.end())
                                                                        {
                                                                                auto const & daughter = *((*d_search).second);
                                                                                int d_pdg = abs(daughter.PdgCode());
                                                                                if (((d_pdg == 2212) || (d_pdg == 211) || (d_pdg == 321)) &&
                                                                                    (1000. * (daughter.E() - daughter.Mass()) > 50.0))
                                                                                {
                                                                                        ++nVisible;
                                                                                }
                                                                        }
                                                                }
                                                        }
                                                        // hadron with Ek > 50MeV (so well visible) and
                                                        // produced by another hadron (but not neutron, so not single track from nothing) or
                                                        // at least secondary hadrons with Ek > 50MeV (so this is a good kink or V-like)
                                                        if (((m_pdg != pdg) && (m_pdg != 2112)) || ((m_pdg != 2112) && (nVisible > 0)) || ((m_pdg == 2112) && (nVisible > 1)))
                                                        {
                                                                // std::cout << "---> hadron at " << ekStart
                                                                //      << ", pdg: " << pdg << ", mother pdg: " << m_pdg
                                                                //      << ", vis.daughters: " << nVisible << std::endl;
                                                                flagsStart = nnet::TrainingDataAlg::kHadr;
                                                        } 
                                                }
                                                // else std::cout << "---> mother not found for tid: " << particle.Mother() << std::endl;
                                        }

                                        if (particle.EndProcess() == "FastScintillation") // potential decay at rest
                                        {
                                                unsigned int nSec = particle.NumberDaughters();
                                                for (size_t d = 0; d < nSec; ++d)
                                                {
                                                        auto d_search = particleMap.find(particle.Daughter(d));
                                                        if (d_search != particleMap.end())
                                                        {
                                                                auto const & daughter = *((*d_search).second);
                                                                int d_pdg = abs(daughter.PdgCode());
                                                                if ((pdg == 321) && (d_pdg == 13))
                                                                {
                                                                        //std::cout << "---> K decay to mu" << std::endl;
                                                                        flagsEnd = nnet::TrainingDataAlg::kDecay;
                                                                        break;
                                                                }
                                                                if ((pdg == 211) && (d_pdg == 13))
                                                                {
                                                                        //std::cout << "---> pi decay to mu" << std::endl;
                                                                        flagsEnd = nnet::TrainingDataAlg::kDecay;
                                                                        break;
                                                                }
                                                        }
                                                }
                                        }

                                        if ((particle.EndProcess() == "Decay") && (ekEnd > 200.0)) // decay in flight
                                        {
                                                unsigned int nSec = particle.NumberDaughters();
                                                for (size_t d = 0; d < nSec; ++d)
                                                {
                                                        auto d_search = particleMap.find(particle.Daughter(d));
                                                        if (d_search != particleMap.end())
                                                        {
                                                                auto const & daughter = *((*d_search).second);
                                                                int d_pdg = abs(daughter.PdgCode());
                                                                if ((pdg == 321) && (d_pdg == 13))
                                                                {
                                                                        //std::cout << "---> in-flight K decay to mu" << std::endl;
                                                                        flagsEnd = nnet::TrainingDataAlg::kHadr;
                                                                        break;
                                                                }
                                                                if ((pdg == 211) && (d_pdg == 13))
                                                                {
                                                                        //std::cout << "---> in-flight pi decay to mu" << std::endl;
                                                                        flagsEnd = nnet::TrainingDataAlg::kHadr;
                                                                        break;
                                                                }
                                                        }
                                                }
                                        }
                                }
                                break;

                        default: continue;
                }
                
                if (particle.Process() == "primary")
                {
                        flagsStart |= nnet::TrainingDataAlg::kNuPri;
                }
                
                
                if (flagsStart != nnet::TrainingDataAlg::kNone)
                {
                        auto wd = getProjection(particle.Position(), plane);
                        
                        if ((wd.TPC == (int)fTPC) && (wd.Cryo == (int)fCryo))
                        {
                                wireToDriftToVtxFlags[wd.Wire][wd.Drift] |= flagsStart;
                                // std::cout << "---> flagsStart:" << flagsStart << " plane:" << plane << " wire:" << wd.Wire << " drift:" << wd.Drift << std::endl;
                        }
                        // else std::cout << "---> not in current TPC" << std::endl;
                }
                if (flagsEnd != nnet::TrainingDataAlg::kNone)
                {
                        auto wd = getProjection(particle.EndPosition(), plane);
                        if ((wd.TPC == (int)fTPC) && (wd.Cryo == (int)fCryo))
                        {
                            //if (flagsEnd == nnet::TrainingDataAlg::kElectronEnd) { std::cout << "---> clear electron endpoint" << std::endl; }
                                wireToDriftToVtxFlags[wd.Wire][wd.Drift] |= flagsEnd;
                                //if (flagsEnd == nnet::TrainingDataAlg::kElectronEnd)
                                //    std::cout << "---> flagsEnd:" << flagsEnd << " plane:" << plane << " wire:" << wd.Wire << " drift:" << wd.Drift << std::endl;
                        }
                        // else std::cout << "---> not in current TPC" << std::endl;
                }

                //if (ekStart > 30.0)
                //{
                //      std::cout << particle.PdgCode() << ", " << ekStart << ": "
                //              << particle.Process() << " --> " << particle.EndProcess()
                //              << " " << ekEnd << std::endl;
                //}
        }
}
// ------------------------------------------------------

bool nnet::TrainingDataAlg::setDataEventData(const art::Event& event,
        unsigned int plane, unsigned int tpc, unsigned int cryo)
{

  art::Handle< std::vector<recob::Wire> > wireHandle;
  std::vector< art::Ptr<recob::Wire> > Wirelist;

  if(event.getByLabel(fWireProducerLabel, wireHandle))
    art::fill_ptr_vector(Wirelist, wireHandle);

  if(!setWireDriftData(*wireHandle, plane, tpc, cryo)) {
    mf::LogError("TrainingDataAlg") << "Wire data not set.";
    return false;
  }

  // Hit info
  art::Handle< std::vector<recob::Hit> > HitHandle;
  std::vector< art::Ptr<recob::Hit> > Hitlist;

  if(event.getByLabel(fHitProducerLabel, HitHandle))
    art::fill_ptr_vector(Hitlist, HitHandle);

  // Track info
  art::Handle< std::vector<recob::Track> > TrackHandle;
  std::vector< art::Ptr<recob::Track> > Tracklist;

  if(event.getByLabel(fTrackModuleLabel, TrackHandle))
    art::fill_ptr_vector(Tracklist, TrackHandle);

  art::FindManyP<recob::Track> ass_trk_hits(HitHandle,   event, fTrackModuleLabel);

  // Loop over wires (sorry about hard coded value) to fill in 1) pdg and 2) charge depo
  for (size_t widx = 0; widx < 240; ++widx) {
    
    std::vector< float > labels_deposit(fNDrifts, 0);  // full-drift-length buffers
    std::vector< int > labels_pdg(fNDrifts, 0);

    // First, the charge depo
    for(size_t subwidx = 0; subwidx < Wirelist.size(); ++subwidx) {
      if(widx+240 == Wirelist[subwidx]->Channel()) {    
        labels_deposit = Wirelist[subwidx]->Signal();
        break;
      }
    }
   
    // Second, the pdg code
    // This code finds the angle of the track and records
    //  events based on its angle to try to get an isometric sample
    //  instead of just a bunch of straight tracks

    // Meta code:
    // For each hit:
    //  find farthest hit from point
    //  then find farthest hit from THAT one
    //  should be start and end of track, then just use trig

    for(size_t iHit = 0; iHit < Hitlist.size(); ++iHit) {

      if(Hitlist[iHit]->Channel() != widx+240) { continue; }
      if(Hitlist[iHit]->View() != 1) { continue; }

      // Make sure there is a track association
      if(ass_trk_hits.at(iHit).size() == 0) { continue; }
      
      // Not sure about this
      // Cutting on length to not just get a bunch of shower stubs
      // Might add a lot of bias though
      if(ass_trk_hits.at(iHit)[0]->Length() < 5) { continue; }

      // Search for farest hit from this one
      int far_index = 0;
      double far_dist = 0;

      for(size_t jHit = 0; jHit < Hitlist.size(); ++jHit) {
        if(jHit == iHit) { continue; }
        if(Hitlist[jHit]->View() != 1) { continue; }

        if(ass_trk_hits.at(jHit).size() == 0) { continue; }
        if(ass_trk_hits.at(jHit)[0]->ID() != 
           ass_trk_hits.at(iHit)[0]->ID()) { continue; }

        double dist = sqrt((Hitlist[iHit]->Channel()-Hitlist[jHit]->Channel()) *
                           (Hitlist[iHit]->Channel()-Hitlist[jHit]->Channel()) +
                           (Hitlist[iHit]->PeakTime()-Hitlist[jHit]->PeakTime()) * 
                           (Hitlist[iHit]->PeakTime()-Hitlist[jHit]->PeakTime()));

        if(far_dist < dist){
          far_dist = dist;
          far_index = jHit;
        }
      }

      // Search for the other end of the track
      int other_end = 0;
      int other_dist = 0;

      for(size_t jHit = 0; jHit < Hitlist.size(); ++jHit) {
        if(jHit == iHit or int(jHit) == far_index) { continue; }
        if(Hitlist[jHit]->View() != 1) { continue; }

        if(ass_trk_hits.at(jHit).size() == 0) { continue; }
        if(ass_trk_hits.at(jHit)[0]->ID() != 
           ass_trk_hits.at(iHit)[0]->ID()) { continue; }

        double dist = sqrt((Hitlist[far_index]->Channel()-Hitlist[jHit]->Channel()) *
                           (Hitlist[far_index]->Channel()-Hitlist[jHit]->Channel()) +
                           (Hitlist[far_index]->PeakTime()-Hitlist[jHit]->PeakTime()) * 
                           (Hitlist[far_index]->PeakTime()-Hitlist[jHit]->PeakTime()));

        if(other_dist < dist){
          other_dist = dist;
          other_end = jHit;
        }
      }

      // We have the end points now
      double del_wire = double(Hitlist[other_end]->Channel() - Hitlist[far_index]->Channel());
      double del_time = double(Hitlist[other_end]->PeakTime() - Hitlist[far_index]->PeakTime());
      double hypo = sqrt(del_wire * del_wire + del_time * del_time);      

      if(hypo == 0) { continue; } // Should never happen, but doing it anyway

      double cosser = TMath::Abs(del_wire / hypo);
      double norm_ang = TMath::ACos(cosser) * 2 / TMath::Pi();

      // Using fEventsPerBin to keep track of number of hits per angle (normalized to 0 to 1)

      int binner = int(norm_ang * fEventsPerBin.size());
      if(binner >= (int)fEventsPerBin.size()) { binner = fEventsPerBin.size() - 1; } // Dealing with rounding errors

      // So we should get a total of 5000 * 100 = 50,000 if we use the whole set
      if(fEventsPerBin[binner] > 5000) { continue; }
      fEventsPerBin[binner]++;
      
      // If survives everything, saves the pdg
      labels_pdg[Hitlist[iHit]->PeakTime()] = 211; // Same as pion for now

    }

    setWireEdepsAndLabels(labels_deposit, labels_pdg, widx);

  } // for each Wire

  /*
  for(size_t i = 0; i < fEventsPerBin.size(); i ++) {
    std::cout << i << ") " << fEventsPerBin[i] << " - ";
  }
  */

  return true;

}

bool nnet::TrainingDataAlg::setEventData(const art::Event& event,
        unsigned int plane, unsigned int tpc, unsigned int cryo)
{
        art::ValidHandle< std::vector<recob::Wire> > wireHandle
                = event.getValidHandle< std::vector<recob::Wire> >(fWireProducerLabel);

        if (!setWireDriftData(*wireHandle, plane, tpc, cryo))
        {
                mf::LogError("TrainingDataAlg") << "Wire data not set.";
                return false;
        }

    if (!fSaveSimInfo || event.isRealData())
    {
        mf::LogInfo("TrainingDataAlg") << "Skip MC simulation info.";
        return true;
    }

        art::ServiceHandle<sim::LArG4Parameters> larParameters;
        double electronsToGeV = 1. / larParameters->GeVToElectrons();

        auto particleHandle = event.getValidHandle< std::vector<simb::MCParticle> >(fSimulationProducerLabel);

        auto simChannelHandle = event.getValidHandle< std::vector<sim::SimChannel> >(fSimulationProducerLabel);

    std::unordered_map< int, const simb::MCParticle* > particleMap;
        for (auto const & particle : *particleHandle)
    {
                particleMap[particle.TrackId()] = &particle;
        }

        std::unordered_map< size_t, std::unordered_map< int, int > > wireToDriftToVtxFlags;
        if (fSaveVtxFlags) collectVtxFlags(wireToDriftToVtxFlags, particleMap, plane);

        fEdepTot = 0;

        std::map< int, int > trackToPDG;
        for (size_t widx = 0; widx < fNWires; ++widx)
        {
                auto wireChannelNumber = fWireChannels[widx];
                if (wireChannelNumber == raw::InvalidChannelID) continue;

                std::vector< float > labels_deposit(fNDrifts, 0);         // full-drift-length buffers,
                std::vector< int > labels_pdg(labels_deposit.size(), 0);  // both of the same size,
                int labels_size = labels_deposit.size();                  // cached as int for comparisons below

                std::map< int, std::map< int, double > > timeToTrackToCharge;
                for (auto const & channel : *simChannelHandle)
                {
                        if (channel.Channel() != wireChannelNumber) continue;           

                        auto const & timeSlices = channel.TDCIDEMap();
                        for (auto const & timeSlice : timeSlices)
                        {
                                int time = timeSlice.first;

                                auto const & energyDeposits = timeSlice.second;
                                for (auto const & energyDeposit : energyDeposits)
                                {
                                        int pdg = 0;
                                        int tid = energyDeposit.trackID;
                                        if (tid < 0) // negative tid means it is EM activity, and -tid is the mother
                                        {
                                            pdg = 11; tid = -tid;
                                            
                                                auto search = particleMap.find(tid);
                                            if (search == particleMap.end())
                                            {
                                                    mf::LogWarning("TrainingDataAlg") << "PARTICLE NOT FOUND";
                                                    continue;
                                            }
                                            auto const & mother = *((*search).second); // mother particle of this EM
                                            int mPdg = abs(mother.PdgCode());
                                            if ((mPdg == 13) || (mPdg == 211) || (mPdg == 2212))
                                            {
                                                if (energyDeposit.numElectrons > 10) pdg |= nnet::TrainingDataAlg::kDelta; // tag delta ray
                                            }
                                        }
                                        else
                                        {
                                          auto search = particleMap.find(tid);
                                          if (search == particleMap.end())
                                          {
                                                   mf::LogWarning("TrainingDataAlg") << "PARTICLE NOT FOUND";
                                                   continue;
                                           }
                                           auto const & particle = *((*search).second);
                                           pdg = abs(particle.PdgCode());
                                           
                                           if (particle.Process() == "primary")
                                           {
                                                        if (pdg == 11)
                                                        {
                                                                pdg |= nnet::TrainingDataAlg::kPriEl; // tag primary
                                                        }
                                                        else if (pdg == 13)
                                                        {
                                                                pdg |= nnet::TrainingDataAlg::kPriMu; // tag primary
                                                        }
                                           }

            auto msearch = particleMap.find(particle.Mother());
                                if (msearch != particleMap.end())
                                {
                                        auto const & mother = *((*msearch).second);
              if (pdg == 11) // electron, check if it is Michel
              {
                                if (nnet::TrainingDataAlg::isMuonDecaying(mother, particleMap))
                              {
                        pdg |= nnet::TrainingDataAlg::kMichel; // tag Michel
                              }
              }
            }
                                        }

                                        trackToPDG[energyDeposit.trackID] = pdg;

                                        double energy = energyDeposit.numElectrons * electronsToGeV;
                                        timeToTrackToCharge[time][energyDeposit.trackID] += energy;
                                        fEdepTot += energy;

                        } // loop over energy deposits
                } // loop over time slices
        } // for each SimChannel

        int type_pdg_mask = nnet::TrainingDataAlg::kTypeMask | nnet::TrainingDataAlg::kPdgMask;
                for (auto const & ttc : timeToTrackToCharge)
                {
                        float max_deposit = 0.0;
                        int max_pdg = 0;
                        for (auto const & tc : ttc.second) {

                                if( tc.second > max_deposit ) 
                                {
                                        max_deposit = tc.second;
                                        max_pdg = trackToPDG[tc.first];
                                }                       
                        }

                        if (ttc.first < labels_size)
                        {
                            int tick_idx = ttc.first + fAdcDelay;
                                if (tick_idx < labels_size)
                                {
                                    labels_deposit[tick_idx] = max_deposit;
                                    labels_pdg[tick_idx] = max_pdg & type_pdg_mask;
                                }
                        }
                }

                for (auto const & drift_flags : wireToDriftToVtxFlags[widx])
                {
                        int drift = drift_flags.first, flags = drift_flags.second;
                        if ((drift >= 0) && (drift < labels_size))
                        {
                                labels_pdg[drift] |= flags;
                        }
                }

                setWireEdepsAndLabels(labels_deposit, labels_pdg, widx);

        } // for each Wire

        return true;
}
// ------------------------------------------------------

bool nnet::TrainingDataAlg::findCrop(float max_e_cut, unsigned int & w0, unsigned int & w1, unsigned int & d0, unsigned int & d1) const
{
    if (fWireDriftEdep.empty() || fWireDriftEdep.front().empty()) return false;

    float max_cut = 0.25 * max_e_cut;

    w0 = 0;
    float cut = 0;
    while (w0 < fWireDriftEdep.size())
    {
        for (auto const d : fWireDriftEdep[w0]) cut += d;
        if (cut < max_cut) w0++;
        else break;
    }
    w1 = fWireDriftEdep.size() - 1;
    cut = 0;
    while (w1 > w0)
    {
        for (auto const d : fWireDriftEdep[w1]) cut += d;
        if (cut < max_cut) w1--;
        else break;
    }
    w1++;

    d0 = 0;
    cut = 0;
    while (d0 < fWireDriftEdep.front().size())
    {
        for (size_t i = w0; i < w1; ++i) cut += fWireDriftEdep[i][d0];
        if (cut < max_cut) d0++;
        else break;
    }
    d1 = fWireDriftEdep.front().size() - 1;
    cut = 0;
    while (d1 > d0)
    {
        for (size_t i = w0; i < w1; ++i) cut += fWireDriftEdep[i][d1];
        if (cut < max_cut) d1--;
        else break;
    }
    d1++;

    unsigned int margin = 20;
    if ((w1 - w0 > 8) && (d1 - d0 > 8))
    {
        if (w0 < margin) w0 = 0;
        else w0 -= margin;

        if (w1 > fWireDriftEdep.size() - margin) w1 = fWireDriftEdep.size();
        else w1 += margin;
        
        if (d0 < margin) d0 = 0;
        else d0 -= margin;
        
        if (d1 > fWireDriftEdep.front().size() - margin) d1 = fWireDriftEdep.front().size();
        else d1 += margin;
        
        return true;
    }
    else return false;
}
// ------------------------------------------------------