/**
  \class    TSGForOIDNN
  \brief    Create L3MuonTrajectorySeeds from L2 Muons in an outside-in manner
  \author   Dmitry Kondratyev, Arnab Purohit, Jan-Frederik Schulte (Purdue University, West Lafayette, USA)
 */

#include "DataFormats/TrackReco/interface/Track.h"
#include "FWCore/Framework/interface/global/EDProducer.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/ESHandle.h"
#include "FWCore/Framework/interface/MakerMacros.h"
#include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "Geometry/CommonDetUnit/interface/GlobalTrackingGeometry.h"
#include "Geometry/Records/interface/GlobalTrackingGeometryRecord.h"
#include "MagneticField/Engine/interface/MagneticField.h"
#include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
#include "RecoTracker/MeasurementDet/interface/MeasurementTrackerEvent.h"
#include "TrackingTools/GeomPropagators/interface/Propagator.h"
#include "TrackingTools/KalmanUpdators/interface/Chi2MeasurementEstimator.h"
#include "TrackingTools/KalmanUpdators/interface/KFUpdator.h"
#include "TrackingTools/MeasurementDet/interface/MeasurementDet.h"
#include "TrackingTools/PatternTools/interface/TrajMeasLessEstim.h"
#include "TrackingTools/Records/interface/TrackingComponentsRecord.h"
#include "TrackingTools/TrajectoryState/interface/TrajectoryStateTransform.h"
#include "TrackingTools/TrajectoryState/interface/TrajectoryStateOnSurface.h"
#include "TrackingTools/GeomPropagators/interface/StateOnTrackerBound.h"
#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
#include "TrackingTools/DetLayers/interface/NavigationSchool.h"
#include "RecoTracker/Record/interface/NavigationSchoolRecord.h"
#include "DataFormats/Math/interface/deltaR.h"
#include "Geometry/TrackerGeometryBuilder/interface/TrackerGeometry.h"
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/json_parser.hpp>
#include <memory>
namespace pt = boost::property_tree;

class TSGForOIDNN : public edm::global::EDProducer<> {
public:
  explicit TSGForOIDNN(const edm::ParameterSet& iConfig);
  ~TSGForOIDNN() override;
  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
  void produce(edm::StreamID sid, edm::Event& iEvent, const edm::EventSetup& iSetup) const override;

private:
  /// Labels for input collections
  const edm::EDGetTokenT<reco::TrackCollection> src_;
  /// Tokens for ESHandle
  const edm::ESGetToken<Chi2MeasurementEstimatorBase, TrackingComponentsRecord> t_estimatorH_;
  const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> t_magfieldH_;
  const edm::ESGetToken<Propagator, TrackingComponentsRecord> t_propagatorAlongH_;
  const edm::ESGetToken<Propagator, TrackingComponentsRecord> t_propagatorOppositeH_;
  const edm::ESGetToken<TrackerGeometry, TrackerDigiGeometryRecord> t_tmpTkGeometryH_;
  const edm::ESGetToken<GlobalTrackingGeometry, GlobalTrackingGeometryRecord> t_geometryH_;
  const edm::ESGetToken<NavigationSchool, NavigationSchoolRecord> t_navSchool_;
  const edm::ESGetToken<Propagator, TrackingComponentsRecord> t_SHPOpposite_;
  const std::unique_ptr<TrajectoryStateUpdator> updator_;
  const edm::EDGetTokenT<MeasurementTrackerEvent> measurementTrackerTag_;
  const std::string theCategory_;

  /// Maximum number of seeds for each L2
  const unsigned int maxSeeds_;
  /// Maximum number of hitbased seeds for each L2
  const unsigned int maxHitSeeds_;
  /// Maximum number of hitless seeds for each L2
  const unsigned int maxHitlessSeeds_;
  /// How many layers to try
  const unsigned int numOfLayersToTry_;
  /// How many hits to try per layer
  const unsigned int numOfHitsToTry_;
  /// Rescale L2 parameter uncertainties (fixed error vs pT, eta)
  const double fixedErrorRescalingForHitless_;
  /// Estimator used to find dets and TrajectoryMeasurements
  const std::string estimatorName_;
  /// Minimum eta value to activate searching in the TEC
  const double minEtaForTEC_;
  /// Maximum eta value to activate searching in the TOB
  const double maxEtaForTOB_;

  /// IP refers to TSOS at interaction point,
  /// MuS refers to TSOS at muon system
  const unsigned int maxHitlessSeedsIP_;
  const unsigned int maxHitlessSeedsMuS_;
  const unsigned int maxHitDoubletSeeds_;

  /// Get number of seeds to use from DNN output instead of "max..Seeds" parameters
  const bool getStrategyFromDNN_;
  /// Whether to use DNN regressor (if false, will use classifier)
  const bool useRegressor_;

  /// Settings for classifier
  const double etaSplitForDnn_;
  std::string dnnModelPath_barrel_;
  std::string dnnModelPath_endcap_;
  std::unique_ptr<tensorflow::GraphDef> graphDef_barrel_;
  tensorflow::Session* tf_session_barrel_;
  std::unique_ptr<tensorflow::GraphDef> graphDef_endcap_;
  tensorflow::Session* tf_session_endcap_;

  /// Settings for regressor
  std::string dnnModelPath_HB_;
  std::string dnnModelPath_HLIP_;
  std::string dnnModelPath_HLMuS_;
  std::unique_ptr<tensorflow::GraphDef> graphDef_HB_;
  tensorflow::Session* tf_session_HB_;
  std::unique_ptr<tensorflow::GraphDef> graphDef_HLIP_;
  tensorflow::Session* tf_session_HLIP_;
  std::unique_ptr<tensorflow::GraphDef> graphDef_HLMuS_;
  tensorflow::Session* tf_session_HLMuS_;

  /// DNN metadata
  const std::string dnnMetadataPath_;
  pt::ptree metadata_;

  /// Create seeds without hits on a given layer (TOB or TEC)
  void makeSeedsWithoutHits(const GeometricSearchDet& layer,
                            const TrajectoryStateOnSurface& tsos,
                            const Propagator& propagatorAlong,
                            const Chi2MeasurementEstimatorBase& estimator,
                            double errorSF,
                            unsigned int& hitlessSeedsMade,
                            unsigned int& numSeedsMade,
                            std::vector<TrajectorySeed>& out) const;

  /// Find hits on a given layer (TOB or TEC) and create seeds from updated TSOS with hit
  void makeSeedsFromHits(const GeometricSearchDet& layer,
                         const TrajectoryStateOnSurface& tsos,
                         const Propagator& propagatorAlong,
                         const Chi2MeasurementEstimatorBase& estimator,
                         const MeasurementTrackerEvent& measurementTracker,
                         unsigned int& hitSeedsMade,
                         unsigned int& numSeedsMade,
                         const unsigned int& maxHitSeeds,
                         unsigned int& layerCount,
                         std::vector<TrajectorySeed>& out) const;

  /// Similar to makeSeedsFromHits, but seed is created only if there are compatible hits on two adjacent layers
  void makeSeedsFromHitDoublets(const GeometricSearchDet& layer,
                                const TrajectoryStateOnSurface& tsos,
                                const Propagator& propagatorAlong,
                                const Chi2MeasurementEstimatorBase& estimator,
                                const MeasurementTrackerEvent& measurementTracker,
                                const NavigationSchool& navSchool,
                                unsigned int& hitDoubletSeedsMade,
                                unsigned int& numSeedsMade,
                                const unsigned int& maxHitDoubletSeeds,
                                unsigned int& layerCount,
                                std::vector<TrajectorySeed>& out) const;

  /// Update dictionary of inputs for DNN
  void updateFeatureMap(std::unordered_map<std::string, float>& the_map,
                        const reco::Track& l2,
                        const TrajectoryStateOnSurface& tsos_IP,
                        const TrajectoryStateOnSurface& tsos_MuS) const;

  /// Container for DNN outupts
  struct StrategyParameters {
    int nHBd, nHLIP, nHLMuS, sf;
  };

  /// Evaluate DNN classifier
  void evaluateClassifier(const std::unordered_map<std::string, float>& feature_map,
                          tensorflow::Session* session,
                          const pt::ptree& metadata,
                          StrategyParameters& out,
                          bool& dnnSuccess) const;

  /// Evaluate DNN regressor
  void evaluateRegressor(const std::unordered_map<std::string, float>& feature_map,
                         tensorflow::Session* session_HB,
                         const pt::ptree& metadata_HB,
                         tensorflow::Session* session_HLIP,
                         const pt::ptree& metadata_HLIP,
                         tensorflow::Session* session_HLMuS,
                         const pt::ptree& metadata_HLMuS,
                         StrategyParameters& out,
                         bool& dnnSuccess) const;
};

TSGForOIDNN::TSGForOIDNN(const edm::ParameterSet& iConfig)
    : src_(consumes(iConfig.getParameter<edm::InputTag>("src"))),
      t_estimatorH_(esConsumes(edm::ESInputTag("", iConfig.getParameter<std::string>("estimator")))),
      t_magfieldH_(esConsumes()),
      t_propagatorAlongH_(esConsumes(edm::ESInputTag("", iConfig.getParameter<std::string>("propagatorName")))),
      t_propagatorOppositeH_(esConsumes(edm::ESInputTag("", iConfig.getParameter<std::string>("propagatorName")))),
      t_tmpTkGeometryH_(esConsumes()),
      t_geometryH_(esConsumes()),
      t_navSchool_(esConsumes(edm::ESInputTag("", "SimpleNavigationSchool"))),
      t_SHPOpposite_(esConsumes(edm::ESInputTag("", "hltESPSteppingHelixPropagatorOpposite"))),
      updator_(new KFUpdator()),
      measurementTrackerTag_(consumes(iConfig.getParameter<edm::InputTag>("MeasurementTrackerEvent"))),
      theCategory_(std::string("Muon|RecoMuon|TSGForOIDNN")),
      maxSeeds_(iConfig.getParameter<uint32_t>("maxSeeds")),
      maxHitSeeds_(iConfig.getParameter<uint32_t>("maxHitSeeds")),
      maxHitlessSeeds_(iConfig.getParameter<uint32_t>("maxHitlessSeeds")),
      numOfLayersToTry_(iConfig.getParameter<int32_t>("layersToTry")),
      numOfHitsToTry_(iConfig.getParameter<int32_t>("hitsToTry")),
      fixedErrorRescalingForHitless_(iConfig.getParameter<double>("fixedErrorRescaleFactorForHitless")),
      minEtaForTEC_(iConfig.getParameter<double>("minEtaForTEC")),
      maxEtaForTOB_(iConfig.getParameter<double>("maxEtaForTOB")),
      maxHitlessSeedsIP_(iConfig.getParameter<uint32_t>("maxHitlessSeedsIP")),
      maxHitlessSeedsMuS_(iConfig.getParameter<uint32_t>("maxHitlessSeedsMuS")),
      maxHitDoubletSeeds_(iConfig.getParameter<uint32_t>("maxHitDoubletSeeds")),
      getStrategyFromDNN_(iConfig.getParameter<bool>("getStrategyFromDNN")),
      useRegressor_(iConfig.getParameter<bool>("useRegressor")),
      etaSplitForDnn_(iConfig.getParameter<double>("etaSplitForDnn")),
      dnnMetadataPath_(iConfig.getParameter<std::string>("dnnMetadataPath")) {
  if (getStrategyFromDNN_) {
    edm::FileInPath dnnMetadataPath(dnnMetadataPath_);
    pt::read_json(dnnMetadataPath.fullPath(), metadata_);
    tensorflow::setLogging("3");

    if (useRegressor_) {
      // use regressor
      dnnModelPath_HB_ = metadata_.get<std::string>("HB.dnnmodel_path");
      edm::FileInPath dnnPath_HB(dnnModelPath_HB_);
      graphDef_HB_ = std::unique_ptr<tensorflow::GraphDef>(tensorflow::loadGraphDef(dnnPath_HB.fullPath()));
      tf_session_HB_ = tensorflow::createSession(graphDef_HB_.get());

      dnnModelPath_HLIP_ = metadata_.get<std::string>("HLIP.dnnmodel_path");
      edm::FileInPath dnnPath_HLIP(dnnModelPath_HLIP_);
      graphDef_HLIP_ = std::unique_ptr<tensorflow::GraphDef>(tensorflow::loadGraphDef(dnnPath_HLIP.fullPath()));
      tf_session_HLIP_ = tensorflow::createSession(graphDef_HLIP_.get());

      dnnModelPath_HLMuS_ = metadata_.get<std::string>("HLMuS.dnnmodel_path");
      edm::FileInPath dnnPath_HLMuS(dnnModelPath_HLMuS_);
      graphDef_HLMuS_ = std::unique_ptr<tensorflow::GraphDef>(tensorflow::loadGraphDef(dnnPath_HLMuS.fullPath()));
      tf_session_HLMuS_ = tensorflow::createSession(graphDef_HLMuS_.get());
    } else {
      // use classifier (default)
      dnnModelPath_barrel_ = metadata_.get<std::string>("barrel.dnnmodel_path");
      edm::FileInPath dnnPath_barrel(dnnModelPath_barrel_);
      graphDef_barrel_ = std::unique_ptr<tensorflow::GraphDef>(tensorflow::loadGraphDef(dnnPath_barrel.fullPath()));
      tf_session_barrel_ = tensorflow::createSession(graphDef_barrel_.get());

      dnnModelPath_endcap_ = metadata_.get<std::string>("endcap.dnnmodel_path");
      edm::FileInPath dnnPath_endcap(dnnModelPath_endcap_);
      graphDef_endcap_ = std::unique_ptr<tensorflow::GraphDef>(tensorflow::loadGraphDef(dnnPath_endcap.fullPath()));
      tf_session_endcap_ = tensorflow::createSession(graphDef_endcap_.get());
    }
  }
  produces<std::vector<TrajectorySeed> >();
}

TSGForOIDNN::~TSGForOIDNN() {
  if (getStrategyFromDNN_) {
    if (useRegressor_) {
      tensorflow::closeSession(tf_session_HB_);
      tensorflow::closeSession(tf_session_HLIP_);
      tensorflow::closeSession(tf_session_HLMuS_);
    } else {
      tensorflow::closeSession(tf_session_barrel_);
      tensorflow::closeSession(tf_session_endcap_);
    }
  }
}

//
// Produce seeds
//
void TSGForOIDNN::produce(edm::StreamID sid, edm::Event& iEvent, edm::EventSetup const& iEventSetup) const {
  // Initialize variables
  unsigned int numSeedsMade = 0;
  unsigned int layerCount = 0;
  unsigned int hitlessSeedsMadeIP = 0;
  unsigned int hitlessSeedsMadeMuS = 0;
  unsigned int hitSeedsMade = 0;
  unsigned int hitDoubletSeedsMade = 0;

  // Container for DNN inputs
  std::unordered_map<std::string, float> feature_map;

  // Container for DNN outputs
  StrategyParameters strPars;

  // Surface used to make a TSOS at the PCA to the beamline
  Plane::PlanePointer dummyPlane = Plane::build(Plane::PositionType(), Plane::RotationType());

  // Get setup objects
  const MagneticField& magfield = iEventSetup.getData(t_magfieldH_);
  const Chi2MeasurementEstimatorBase& estimator = iEventSetup.getData(t_estimatorH_);
  const Propagator& tmpPropagatorAlong = iEventSetup.getData(t_propagatorAlongH_);
  const Propagator& tmpPropagatorOpposite = iEventSetup.getData(t_propagatorOppositeH_);
  const TrackerGeometry& tmpTkGeometry = iEventSetup.getData(t_tmpTkGeometryH_);
  const GlobalTrackingGeometry& geometry = iEventSetup.getData(t_geometryH_);
  const NavigationSchool& navSchool = iEventSetup.getData(t_navSchool_);
  auto const& measurementTracker = iEvent.get(measurementTrackerTag_);

  // Read L2 track collection
  auto const& l2TrackCol = iEvent.get(src_);

  // The product
  std::unique_ptr<std::vector<TrajectorySeed> > result(new std::vector<TrajectorySeed>());

  // Get vector of Detector layers
  auto const* gsTracker = measurementTracker.geometricSearchTracker();
  std::vector<BarrelDetLayer const*> const& tob = gsTracker->tobLayers();
  std::vector<ForwardDetLayer const*> const& tecPositive =
      tmpTkGeometry.isThere(GeomDetEnumerators::P2OTEC) ? gsTracker->posTidLayers() : gsTracker->posTecLayers();
  std::vector<ForwardDetLayer const*> const& tecNegative =
      tmpTkGeometry.isThere(GeomDetEnumerators::P2OTEC) ? gsTracker->negTidLayers() : gsTracker->negTecLayers();

  // Get suitable propagators
  std::unique_ptr<Propagator> propagatorAlong = SetPropagationDirection(tmpPropagatorAlong, alongMomentum);
  std::unique_ptr<Propagator> propagatorOpposite = SetPropagationDirection(tmpPropagatorOpposite, oppositeToMomentum);

  // Stepping Helix Propagator for propogation from muon system to tracker
  edm::ESHandle<Propagator> shpOpposite = iEventSetup.getHandle(t_SHPOpposite_);

  // Loop over the L2's and make seeds for all of them
  LogTrace(theCategory_) << "TSGForOIDNN::produce: Number of L2's: " << l2TrackCol.size();

  for (auto const& l2 : l2TrackCol) {
    // Container of Seeds
    std::vector<TrajectorySeed> out;
    LogTrace("TSGForOIDNN") << "TSGForOIDNN::produce: L2 muon pT, eta, phi --> " << l2.pt() << " , " << l2.eta()
                            << " , " << l2.phi();

    FreeTrajectoryState fts = trajectoryStateTransform::initialFreeState(l2, &magfield);

    dummyPlane->move(fts.position() - dummyPlane->position());
    TrajectoryStateOnSurface tsosAtIP = TrajectoryStateOnSurface(fts, *dummyPlane);
    LogTrace("TSGForOIDNN") << "TSGForOIDNN::produce: Created TSOSatIP: " << tsosAtIP;

    // Get the TSOS on the innermost layer of the L2
    TrajectoryStateOnSurface tsosAtMuonSystem = trajectoryStateTransform::innerStateOnSurface(l2, geometry, &magfield);
    LogTrace("TSGForOIDNN") << "TSGForOIDNN::produce: Created TSOSatMuonSystem: " << tsosAtMuonSystem;

    LogTrace("TSGForOIDNN")
        << "TSGForOIDNN::produce: Check the error of the L2 parameter and use hit seeds if big errors";

    StateOnTrackerBound fromInside(propagatorAlong.get());
    TrajectoryStateOnSurface outerTkStateInside = fromInside(fts);

    StateOnTrackerBound fromOutside(&*shpOpposite);
    TrajectoryStateOnSurface outerTkStateOutside = fromOutside(tsosAtMuonSystem);

    // Check if the two positions (using updated and not-updated TSOS) agree withing certain extent.
    // If both TSOSs agree, use only the one at vertex, as it uses more information. If they do not agree, search for seeds based on both.
    double l2muonEta = l2.eta();
    double absL2muonEta = std::abs(l2muonEta);

    // make non-const copies of parameters, so they can be overriden for individual L2 muons
    unsigned int maxHitSeeds = maxHitSeeds_;
    unsigned int maxHitDoubletSeeds = maxHitDoubletSeeds_;
    unsigned int maxHitlessSeedsIP = maxHitlessSeedsIP_;
    unsigned int maxHitlessSeedsMuS = maxHitlessSeedsMuS_;

    float errorSFHitless = fixedErrorRescalingForHitless_;

    // update strategy parameters by evaluating DNN
    if (getStrategyFromDNN_) {
      bool dnnSuccess = false;

      // Update feature map with parameters of the current muon
      updateFeatureMap(feature_map, l2, tsosAtIP, outerTkStateOutside);

      if (useRegressor_) {
        // Use regressor
        evaluateRegressor(feature_map,
                          tf_session_HB_,
                          metadata_.get_child("HB"),
                          tf_session_HLIP_,
                          metadata_.get_child("HLIP"),
                          tf_session_HLMuS_,
                          metadata_.get_child("HLMuS"),
                          strPars,
                          dnnSuccess);
      } else {
        // Use classifier
        bool isBarrel = absL2muonEta < etaSplitForDnn_;
        const pt::ptree& tr = isBarrel ? metadata_.get_child("barrel") : metadata_.get_child("endcap");
        tensorflow::Session* tf = isBarrel ? tf_session_barrel_ : tf_session_endcap_;
        evaluateClassifier(feature_map, tf, tr, strPars, dnnSuccess);
      }
      if (!dnnSuccess)
        break;

      maxHitSeeds = 0;
      maxHitDoubletSeeds = strPars.nHBd;
      maxHitlessSeedsIP = strPars.nHLIP;
      maxHitlessSeedsMuS = strPars.nHLMuS;
      errorSFHitless = strPars.sf;
    }

    numSeedsMade = 0;
    hitlessSeedsMadeIP = 0;
    hitlessSeedsMadeMuS = 0;
    hitSeedsMade = 0;
    hitDoubletSeedsMade = 0;

    auto createSeeds = [&](auto const& layers) {
      for (auto const& layer : layers) {
        if (hitlessSeedsMadeIP < maxHitlessSeedsIP && numSeedsMade < maxSeeds_)
          makeSeedsWithoutHits(*layer,
                               tsosAtIP,
                               *(propagatorAlong.get()),
                               estimator,
                               errorSFHitless,
                               hitlessSeedsMadeIP,
                               numSeedsMade,
                               out);

        if (outerTkStateInside.isValid() && outerTkStateOutside.isValid() && hitlessSeedsMadeMuS < maxHitlessSeedsMuS &&
            numSeedsMade < maxSeeds_)
          makeSeedsWithoutHits(*layer,
                               outerTkStateOutside,
                               *(propagatorOpposite.get()),
                               estimator,
                               errorSFHitless,
                               hitlessSeedsMadeMuS,
                               numSeedsMade,
                               out);

        if (hitSeedsMade < maxHitSeeds && numSeedsMade < maxSeeds_)
          makeSeedsFromHits(*layer,
                            tsosAtIP,
                            *(propagatorAlong.get()),
                            estimator,
                            measurementTracker,
                            hitSeedsMade,
                            numSeedsMade,
                            maxHitSeeds,
                            layerCount,
                            out);

        if (hitDoubletSeedsMade < maxHitDoubletSeeds && numSeedsMade < maxSeeds_)
          makeSeedsFromHitDoublets(*layer,
                                   tsosAtIP,
                                   *(propagatorAlong.get()),
                                   estimator,
                                   measurementTracker,
                                   navSchool,
                                   hitDoubletSeedsMade,
                                   numSeedsMade,
                                   maxHitDoubletSeeds,
                                   layerCount,
                                   out);
      };
    };

    // BARREL
    if (absL2muonEta < maxEtaForTOB_) {
      layerCount = 0;
      createSeeds(tob);
      LogTrace("TSGForOIDNN") << "TSGForOIDNN:::produce: NumSeedsMade = " << numSeedsMade
                              << " , layerCount = " << layerCount;
    }

    // Reset number of seeds if in overlap region
    if (absL2muonEta > minEtaForTEC_ && absL2muonEta < maxEtaForTOB_) {
      numSeedsMade = 0;
      hitlessSeedsMadeIP = 0;
      hitlessSeedsMadeMuS = 0;
      hitSeedsMade = 0;
      hitDoubletSeedsMade = 0;
    }

    // ENDCAP+
    if (l2muonEta > minEtaForTEC_) {
      layerCount = 0;
      createSeeds(tecPositive);
      LogTrace("TSGForOIDNN") << "TSGForOIDNN:::produce: NumSeedsMade = " << numSeedsMade
                              << " , layerCount = " << layerCount;
    }

    // ENDCAP-
    if (l2muonEta < -minEtaForTEC_) {
      layerCount = 0;
      createSeeds(tecNegative);
      LogTrace("TSGForOIDNN") << "TSGForOIDNN:::produce: NumSeedsMade = " << numSeedsMade
                              << " , layerCount = " << layerCount;
    }

    for (std::vector<TrajectorySeed>::iterator it = out.begin(); it != out.end(); ++it) {
      result->push_back(*it);
    }

  }  // L2Collection

  edm::LogInfo(theCategory_) << "TSGForOIDNN::produce: number of seeds made: " << result->size();

  iEvent.put(std::move(result));
}

//
// Create seeds without hits on a given layer (TOB or TEC)
//
void TSGForOIDNN::makeSeedsWithoutHits(const GeometricSearchDet& layer,
                                       const TrajectoryStateOnSurface& tsos,
                                       const Propagator& propagatorAlong,
                                       const Chi2MeasurementEstimatorBase& estimator,
                                       double errorSF,
                                       unsigned int& hitlessSeedsMade,
                                       unsigned int& numSeedsMade,
                                       std::vector<TrajectorySeed>& out) const {
  // create hitless seeds
  LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsWithoutHits: Start hitless";
  std::vector<GeometricSearchDet::DetWithState> dets;
  layer.compatibleDetsV(tsos, propagatorAlong, estimator, dets);
  if (!dets.empty()) {
    auto const& detOnLayer = dets.front().first;
    auto& tsosOnLayer = dets.front().second;
    LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsWithoutHits: tsosOnLayer " << tsosOnLayer;
    if (!tsosOnLayer.isValid()) {
      edm::LogInfo(theCategory_) << "ERROR!: Hitless TSOS is not valid!";
    } else {
      tsosOnLayer.rescaleError(errorSF);
      PTrajectoryStateOnDet const& ptsod =
          trajectoryStateTransform::persistentState(tsosOnLayer, detOnLayer->geographicalId().rawId());
      TrajectorySeed::RecHitContainer rHC;
      out.push_back(TrajectorySeed(ptsod, rHC, oppositeToMomentum));
      LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsWithoutHits: TSOS (Hitless) done ";
      hitlessSeedsMade++;
      numSeedsMade++;
    }
  }
}

//
// Find hits on a given layer (TOB or TEC) and create seeds from updated TSOS with hit
//
void TSGForOIDNN::makeSeedsFromHits(const GeometricSearchDet& layer,
                                    const TrajectoryStateOnSurface& tsos,
                                    const Propagator& propagatorAlong,
                                    const Chi2MeasurementEstimatorBase& estimator,
                                    const MeasurementTrackerEvent& measurementTracker,
                                    unsigned int& hitSeedsMade,
                                    unsigned int& numSeedsMade,
                                    const unsigned int& maxHitSeeds,
                                    unsigned int& layerCount,
                                    std::vector<TrajectorySeed>& out) const {
  if (layerCount > numOfLayersToTry_)
    return;

  const TrajectoryStateOnSurface& onLayer(tsos);

  std::vector<GeometricSearchDet::DetWithState> dets;
  layer.compatibleDetsV(onLayer, propagatorAlong, estimator, dets);

  // Find Measurements on each DetWithState
  LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsFromHits: Find measurements on each detWithState  " << dets.size();
  std::vector<TrajectoryMeasurement> meas;
  for (auto const& detI : dets) {
    MeasurementDetWithData det = measurementTracker.idToDet(detI.first->geographicalId());
    if (det.isNull())
      continue;
    if (!detI.second.isValid())
      continue;  // Skip if TSOS is not valid

    std::vector<TrajectoryMeasurement> mymeas =
        det.fastMeasurements(detI.second, onLayer, propagatorAlong, estimator);  // Second TSOS is not used
    for (auto const& measurement : mymeas) {
      if (measurement.recHit()->isValid())
        meas.push_back(measurement);  // Only save those which are valid
    }
  }

  // Update TSOS using TMs after sorting, then create Trajectory Seed and put into vector
  LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsFromHits: Update TSOS using TMs after sorting, then create "
                             "Trajectory Seed, number of TM = "
                          << meas.size();
  std::sort(meas.begin(), meas.end(), TrajMeasLessEstim());

  unsigned int found = 0;
  for (auto const& measurement : meas) {
    if (hitSeedsMade >= maxHitSeeds)
      return;
    TrajectoryStateOnSurface updatedTSOS = updator_->update(measurement.forwardPredictedState(), *measurement.recHit());
    LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsFromHits: TSOS for TM " << found;
    if (not updatedTSOS.isValid())
      continue;

    edm::OwnVector<TrackingRecHit> seedHits;
    seedHits.push_back(*measurement.recHit()->hit());
    PTrajectoryStateOnDet const& pstate =
        trajectoryStateTransform::persistentState(updatedTSOS, measurement.recHit()->geographicalId().rawId());
    LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsFromHits: Number of seedHits: " << seedHits.size();
    TrajectorySeed seed(pstate, std::move(seedHits), oppositeToMomentum);
    out.push_back(seed);
    found++;
    numSeedsMade++;
    hitSeedsMade++;
    if (found == numOfHitsToTry_)
      break;
  }

  if (found)
    layerCount++;
}

//
// Find hits compatible with L2 trajectory on two adjacent layers; if found, create a seed using both hits
//
void TSGForOIDNN::makeSeedsFromHitDoublets(const GeometricSearchDet& layer,
                                           const TrajectoryStateOnSurface& tsos,
                                           const Propagator& propagatorAlong,
                                           const Chi2MeasurementEstimatorBase& estimator,
                                           const MeasurementTrackerEvent& measurementTracker,
                                           const NavigationSchool& navSchool,
                                           unsigned int& hitDoubletSeedsMade,
                                           unsigned int& numSeedsMade,
                                           const unsigned int& maxHitDoubletSeeds,
                                           unsigned int& layerCount,
                                           std::vector<TrajectorySeed>& out) const {
  // This method is similar to makeSeedsFromHits, but the seed is created
  // only when in addition to a hit on a given layer, there are more compatible hits
  // on next layers (going from outside inwards), compatible with updated TSOS.
  // If that's the case, multiple compatible hits are used to create a single seed.

  // Configured to only check the immideately adjacent layer and add one more hit
  int max_addtnl_layers = 1;  // max number of additional layers to scan
  int max_meas = 1;           // number of measurements to consider on each additional layer

  // // // First, regular procedure to find a compatible hit - like in makeSeedsFromHits // // //

  const TrajectoryStateOnSurface& onLayer(tsos);

  // Find dets compatible with original TSOS
  std::vector<GeometricSearchDet::DetWithState> dets;
  layer.compatibleDetsV(onLayer, propagatorAlong, estimator, dets);

  LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsFromHitDoublets: Find measurements on each detWithState  "
                          << dets.size();
  std::vector<TrajectoryMeasurement> meas;

  // Loop over dets
  for (auto const& detI : dets) {
    MeasurementDetWithData det = measurementTracker.idToDet(detI.first->geographicalId());

    if (det.isNull())
      continue;  // skip if det does not exist
    if (!detI.second.isValid())
      continue;  // skip if TSOS is invalid

    // Find measurements on this det
    std::vector<TrajectoryMeasurement> mymeas = det.fastMeasurements(detI.second, onLayer, propagatorAlong, estimator);

    // Save valid measurements
    for (auto const& measurement : mymeas) {
      if (measurement.recHit()->isValid())
        meas.push_back(measurement);
    }  // end loop over meas
  }    // end loop over dets

  LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsFromHitDoublets: Update TSOS using TMs after sorting, then create "
                             "Trajectory Seed, number of TM = "
                          << meas.size();

  // sort valid measurements found on the first layer
  std::sort(meas.begin(), meas.end(), TrajMeasLessEstim());

  unsigned int found = 0;
  int hit_num = 0;

  // Loop over all valid measurements compatible with original TSOS
  //for (std::vector<TrajectoryMeasurement>::const_iterator mea = meas.begin(); mea != meas.end(); ++mea) {
  for (auto const& measurement : meas) {
    if (hitDoubletSeedsMade >= maxHitDoubletSeeds)
      return;  // abort if enough seeds created

    hit_num++;

    // Update TSOS with measurement on first considered layer
    TrajectoryStateOnSurface updatedTSOS = updator_->update(measurement.forwardPredictedState(), *measurement.recHit());

    LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsFromHitDoublets: TSOS for TM " << found;
    if (not updatedTSOS.isValid())
      continue;  // Skip if updated TSOS is invalid

    edm::OwnVector<TrackingRecHit> seedHits;

    // Save hit on first layer
    seedHits.push_back(*measurement.recHit()->hit());
    const DetLayer* detLayer = dynamic_cast<const DetLayer*>(&layer);

    // // // Now for this measurement we will loop over additional layers and try to update the TSOS again // // //

    // find layers compatible with updated TSOS
    auto const& compLayers = navSchool.nextLayers(*detLayer, *updatedTSOS.freeState(), alongMomentum);

    int addtnl_layers_scanned = 0;
    int found_compatible_on_next_layer = 0;
    int det_id = 0;

    // Copy updated TSOS - we will update it again with a measurement from the next layer, if we find it
    TrajectoryStateOnSurface updatedTSOS_next(updatedTSOS);

    // loop over layers compatible with updated TSOS
    for (auto compLayer : compLayers) {
      int nmeas = 0;

      if (addtnl_layers_scanned >= max_addtnl_layers)
        break;  // break if we already looped over enough layers
      if (found_compatible_on_next_layer > 0)
        break;  // break if we already found additional hit

      // find dets compatible with updated TSOS
      std::vector<GeometricSearchDet::DetWithState> dets_next;
      TrajectoryStateOnSurface onLayer_next(updatedTSOS);

      compLayer->compatibleDetsV(onLayer_next, propagatorAlong, estimator, dets_next);

      //if (!detWithState.size()) continue;
      std::vector<TrajectoryMeasurement> meas_next;

      // find measurements on dets_next and save the valid ones
      for (auto const& detI_next : dets_next) {
        MeasurementDetWithData det = measurementTracker.idToDet(detI_next.first->geographicalId());

        if (det.isNull())
          continue;  // skip if det does not exist
        if (!detI_next.second.isValid())
          continue;  // skip if TSOS is invalid

        // Find measurements on this det
        std::vector<TrajectoryMeasurement> mymeas_next =
            det.fastMeasurements(detI_next.second, onLayer_next, propagatorAlong, estimator);

        for (auto const& mea_next : mymeas_next) {
          // save valid measurements
          if (mea_next.recHit()->isValid())
            meas_next.push_back(mea_next);

        }  // end loop over mymeas_next
      }    // end loop over dets_next

      // sort valid measurements found on this layer
      std::sort(meas_next.begin(), meas_next.end(), TrajMeasLessEstim());

      // loop over valid measurements compatible with updated TSOS (TSOS updated with a hit on the first layer)
      for (auto const& mea_next : meas_next) {
        if (nmeas >= max_meas)
          break;  // skip if we already found enough hits

        // try to update TSOS again, with an additional hit
        updatedTSOS_next = updator_->update(mea_next.forwardPredictedState(), *mea_next.recHit());

        if (not updatedTSOS_next.isValid())
          continue;  // skip if TSOS updated with additional hit is not valid

        // If there was a compatible hit on this layer, we end up here.
        // An additional compatible hit is saved.
        seedHits.push_back(*mea_next.recHit()->hit());
        det_id = mea_next.recHit()->geographicalId().rawId();
        nmeas++;
        found_compatible_on_next_layer++;

      }  // end loop over meas_next

      addtnl_layers_scanned++;

    }  // end loop over compLayers (additional layers scanned after the original layer)

    if (found_compatible_on_next_layer == 0)
      continue;
    // only consider the hit if there was a compatible hit on one of the additional scanned layers

    // Create a seed from two saved hits
    PTrajectoryStateOnDet const& pstate = trajectoryStateTransform::persistentState(updatedTSOS_next, det_id);
    TrajectorySeed seed(pstate, std::move(seedHits), oppositeToMomentum);

    LogTrace("TSGForOIDNN") << "TSGForOIDNN::makeSeedsFromHitDoublets: Number of seedHits: " << seedHits.size();
    out.push_back(seed);

    found++;
    numSeedsMade++;
    hitDoubletSeedsMade++;

    if (found == numOfHitsToTry_)
      break;  // break if enough measurements scanned

  }  // end loop over measurements compatible with original TSOS

  if (found)
    layerCount++;
}

//
// Update the dictionary of variables to use as input features for DNN
//
void TSGForOIDNN::updateFeatureMap(std::unordered_map<std::string, float>& the_map,
                                   const reco::Track& l2,
                                   const TrajectoryStateOnSurface& tsos_IP,
                                   const TrajectoryStateOnSurface& tsos_MuS) const {
  the_map["pt"] = l2.pt();
  the_map["eta"] = l2.eta();
  the_map["phi"] = l2.phi();
  the_map["validHits"] = l2.found();
  if (tsos_IP.isValid()) {
    the_map["tsos_IP_eta"] = tsos_IP.globalPosition().eta();
    the_map["tsos_IP_phi"] = tsos_IP.globalPosition().phi();
    the_map["tsos_IP_pt"] = tsos_IP.globalMomentum().perp();
    the_map["tsos_IP_pt_eta"] = tsos_IP.globalMomentum().eta();
    the_map["tsos_IP_pt_phi"] = tsos_IP.globalMomentum().phi();
    const AlgebraicSymMatrix55& matrix_IP = tsos_IP.curvilinearError().matrix();
    the_map["err0_IP"] = sqrt(matrix_IP[0][0]);
    the_map["err1_IP"] = sqrt(matrix_IP[1][1]);
    the_map["err2_IP"] = sqrt(matrix_IP[2][2]);
    the_map["err3_IP"] = sqrt(matrix_IP[3][3]);
    the_map["err4_IP"] = sqrt(matrix_IP[4][4]);
    the_map["tsos_IP_valid"] = 1.0;
  } else {
    the_map["tsos_IP_eta"] = -999;
    the_map["tsos_IP_phi"] = -999;
    the_map["tsos_IP_pt"] = -999;
    the_map["tsos_IP_pt_eta"] = -999;
    the_map["tsos_IP_pt_phi"] = -999;
    the_map["err0_IP"] = -999;
    the_map["err1_IP"] = -999;
    the_map["err2_IP"] = -999;
    the_map["err3_IP"] = -999;
    the_map["err4_IP"] = -999;
    the_map["tsos_IP_valid"] = 0.0;
  }
  if (tsos_MuS.isValid()) {
    the_map["tsos_MuS_eta"] = tsos_MuS.globalPosition().eta();
    the_map["tsos_MuS_phi"] = tsos_MuS.globalPosition().phi();
    the_map["tsos_MuS_pt"] = tsos_MuS.globalMomentum().perp();
    the_map["tsos_MuS_pt_eta"] = tsos_MuS.globalMomentum().eta();
    the_map["tsos_MuS_pt_phi"] = tsos_MuS.globalMomentum().phi();
    const AlgebraicSymMatrix55& matrix_MuS = tsos_MuS.curvilinearError().matrix();
    the_map["err0_MuS"] = sqrt(matrix_MuS[0][0]);
    the_map["err1_MuS"] = sqrt(matrix_MuS[1][1]);
    the_map["err2_MuS"] = sqrt(matrix_MuS[2][2]);
    the_map["err3_MuS"] = sqrt(matrix_MuS[3][3]);
    the_map["err4_MuS"] = sqrt(matrix_MuS[4][4]);
    the_map["tsos_MuS_valid"] = 1.0;
  } else {
    the_map["tsos_MuS_eta"] = -999;
    the_map["tsos_MuS_phi"] = -999;
    the_map["tsos_MuS_pt"] = -999;
    the_map["tsos_MuS_pt_eta"] = -999;
    the_map["tsos_MuS_pt_phi"] = -999;
    the_map["err0_MuS"] = -999;
    the_map["err1_MuS"] = -999;
    the_map["err2_MuS"] = -999;
    the_map["err3_MuS"] = -999;
    the_map["err4_MuS"] = -999;
    the_map["tsos_MuS_valid"] = 0.0;
  }
}

//
// Obtain seeding strategy parameters by evaluating DNN classifier for a given muon
//
void TSGForOIDNN::evaluateClassifier(const std::unordered_map<std::string, float>& feature_map,
                                     tensorflow::Session* session,
                                     const pt::ptree& metadata,
                                     StrategyParameters& out,
                                     bool& dnnSuccess) const {
  int n_features = metadata.get<int>("n_features", 0);

  // Prepare tensor for DNN inputs
  tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, n_features});
  std::string fname;
  int i_feature = 0;
  for (const pt::ptree::value_type& feature : metadata.get_child("feature_names")) {
    fname = feature.second.data();
    if (feature_map.find(fname) == feature_map.end()) {
      // don't evaluate DNN if any input feature is missing
      dnnSuccess = false;
    } else {
      input.matrix<float>()(0, i_feature) = float(feature_map.at(fname));
      i_feature++;
    }
  }

  // Prepare tensor for DNN outputs
  std::vector<tensorflow::Tensor> outputs;

  // Evaluate DNN and put results in output tensor
  std::string inputLayer = metadata.get<std::string>("input_layer");
  std::string outputLayer = metadata.get<std::string>("output_layer");

  tensorflow::run(session, {{inputLayer, input}}, {outputLayer}, &outputs);
  tensorflow::Tensor out_tensor = outputs[0];
  tensorflow::TTypes<float, 1>::Matrix dnn_outputs = out_tensor.matrix<float>();

  // Find output with largest prediction
  int imax = -1;
  float out_max = 0;
  for (long long int i = 0; i < out_tensor.dim_size(1); i++) {
    float ith_output = dnn_outputs(0, i);
    if (ith_output > out_max) {
      imax = i;
      out_max = ith_output;
    }
  }

  // Decode output
  const std::string label = "output_labels.label_" + std::to_string(imax);
  out.nHBd = metadata.get<int>(label + ".nHBd");
  out.nHLIP = metadata.get<int>(label + ".nHLIP");
  out.nHLMuS = metadata.get<int>(label + ".nHLMuS");
  out.sf = metadata.get<int>(label + ".SF");

  dnnSuccess = true;
}

//
// Obtain seeding strategy parameters by evaluating DNN regressor for a given muon
//
void TSGForOIDNN::evaluateRegressor(const std::unordered_map<std::string, float>& feature_map,
                                    tensorflow::Session* session_HB,
                                    const pt::ptree& metadata_HB,
                                    tensorflow::Session* session_HLIP,
                                    const pt::ptree& metadata_HLIP,
                                    tensorflow::Session* session_HLMuS,
                                    const pt::ptree& metadata_HLMuS,
                                    StrategyParameters& out,
                                    bool& dnnSuccess) const {
  int n_features = metadata_HB.get<int>("n_features", 0);

  // Prepare tensor for DNN inputs
  tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, n_features});
  std::string fname;
  int i_feature = 0;
  for (const pt::ptree::value_type& feature : metadata_HB.get_child("feature_names")) {
    fname = feature.second.data();
    if (feature_map.find(fname) == feature_map.end()) {
      // don't evaluate DNN if any input feature is missing
      dnnSuccess = false;
    } else {
      input.matrix<float>()(0, i_feature) = float(feature_map.at(fname));
      i_feature++;
    }
  }

  // Prepare tensor for DNN outputs
  std::vector<tensorflow::Tensor> outputs_HB;
  // Evaluate DNN and put results in output tensor
  std::string inputLayer_HB = metadata_HB.get<std::string>("input_layer");
  std::string outputLayer_HB = metadata_HB.get<std::string>("output_layer");
  tensorflow::run(session_HB, {{inputLayer_HB, input}}, {outputLayer_HB}, &outputs_HB);
  tensorflow::Tensor out_tensor_HB = outputs_HB[0];
  tensorflow::TTypes<float, 1>::Matrix dnn_outputs_HB = out_tensor_HB.matrix<float>();

  // Prepare tensor for DNN outputs
  std::vector<tensorflow::Tensor> outputs_HLIP;
  // Evaluate DNN and put results in output tensor
  std::string inputLayer_HLIP = metadata_HLIP.get<std::string>("input_layer");
  std::string outputLayer_HLIP = metadata_HLIP.get<std::string>("output_layer");
  tensorflow::run(session_HLIP, {{inputLayer_HLIP, input}}, {outputLayer_HLIP}, &outputs_HLIP);
  tensorflow::Tensor out_tensor_HLIP = outputs_HLIP[0];
  tensorflow::TTypes<float, 1>::Matrix dnn_outputs_HLIP = out_tensor_HLIP.matrix<float>();

  // Prepare tensor for DNN outputs
  std::vector<tensorflow::Tensor> outputs_HLMuS;
  // Evaluate DNN and put results in output tensor
  std::string inputLayer_HLMuS = metadata_HLMuS.get<std::string>("input_layer");
  std::string outputLayer_HLMuS = metadata_HLMuS.get<std::string>("output_layer");
  tensorflow::run(session_HLMuS, {{inputLayer_HLMuS, input}}, {outputLayer_HLMuS}, &outputs_HLMuS);
  tensorflow::Tensor out_tensor_HLMuS = outputs_HLMuS[0];
  tensorflow::TTypes<float, 1>::Matrix dnn_outputs_HLMuS = out_tensor_HLMuS.matrix<float>();

  // Decode output
  out.nHBd = round(dnn_outputs_HB(0, 0));
  out.nHLIP = round(dnn_outputs_HLIP(0, 0));
  out.sf = round(dnn_outputs_HLIP(0, 1));
  out.nHLMuS = round(dnn_outputs_HLMuS(0, 0));

  // Prevent prediction of negative number of seeds or too many seeds
  out.nHBd = std::clamp(out.nHBd, 0, 10);
  out.nHLIP = std::clamp(out.nHLIP, 0, 10);
  out.nHLMuS = std::clamp(out.nHLMuS, 0, 10);

  // Prevent prediction of 0 seeds in total
  if (out.nHBd == 0 && out.nHLIP == 0 && out.nHLMuS == 0) {
    // default strategy, similar to Run 2
    out.nHBd = 1;
    out.nHLIP = 5;
  }

  // Prevent extreme predictions for scale factors
  // (on average SF=2 was found to be optimal)
  if (out.sf <= 0)
    out.sf = 2;
  if (out.sf > 10)
    out.sf = 10;

  dnnSuccess = true;
}

//
// Default values of configuration parameters
//
void TSGForOIDNN::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
  edm::ParameterSetDescription desc;
  desc.add<edm::InputTag>("src", edm::InputTag("hltL2Muons", "UpdatedAtVtx"));
  desc.add<int>("layersToTry", 2);
  desc.add<double>("fixedErrorRescaleFactorForHitless", 2.0);
  desc.add<int>("hitsToTry", 1);
  desc.add<edm::InputTag>("MeasurementTrackerEvent", edm::InputTag("hltSiStripClusters"));
  desc.add<std::string>("estimator", "hltESPChi2MeasurementEstimator100");
  desc.add<double>("maxEtaForTOB", 1.8);
  desc.add<double>("minEtaForTEC", 0.7);
  desc.addUntracked<bool>("debug", false);
  desc.add<unsigned int>("maxSeeds", 20);
  desc.add<unsigned int>("maxHitlessSeeds", 5);
  desc.add<unsigned int>("maxHitSeeds", 1);
  desc.add<std::string>("propagatorName", "PropagatorWithMaterialParabolicMf");
  desc.add<unsigned int>("maxHitlessSeedsIP", 5);
  desc.add<unsigned int>("maxHitlessSeedsMuS", 0);
  desc.add<unsigned int>("maxHitDoubletSeeds", 0);
  desc.add<bool>("getStrategyFromDNN", false);
  desc.add<bool>("useRegressor", false);
  desc.add<double>("etaSplitForDnn", 1.0);
  desc.add<std::string>("dnnMetadataPath", "");
  descriptions.add("tsgForOIDNN", desc);
}

DEFINE_FWK_MODULE(TSGForOIDNN);


//==================================================================================================
/**
  Copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
**/
//==================================================================================================

#ifndef BOOST_SIMD_FUNCTION_SIMD_CORRECT_FMA_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_SIMD_CORRECT_FMA_HPP_INCLUDED

#include <boost/simd/function/scalar/correct_fma.hpp>
#include <boost/simd/arch/common/generic/function/autodispatcher.hpp>
#include <boost/simd/arch/common/simd/function/correct_fma.hpp>

#if defined(BOOST_HW_SIMD_X86_OR_AMD_AVAILABLE)
#  if BOOST_HW_SIMD_X86_AMD_FMA4
#    include <boost/simd/arch/x86/fma4/simd/function/correct_fma.hpp>
#  endif
#  if BOOST_HW_SIMD_X86_FMA3
#    include <boost/simd/arch/x86/fma3/simd/function/correct_fma.hpp>
#  endif
#endif

#if defined(BOOST_HW_SIMD_PPC_AVAILABLE)
#  if BOOST_HW_SIMD_PPC >= BOOST_HW_SIMD_PPC_VMX_VERSION
#    include <boost/simd/arch/ppc/vmx/simd/function/correct_fma.hpp>
#  endif
#endif

#endif


//Copyright (c) 2008-2016 Emil Dotchevski and Reverge Studios, Inc.

//Distributed under the Boost Software License, Version 1.0. (See accompanying
//file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#include <boost/qvm/scalar_traits.hpp>
#include <boost/qvm/vec.hpp>
#include <boost/qvm/mat.hpp>
#include <boost/qvm/quat.hpp>

template <bool>
struct tester;

template <>
struct
tester<true>
    {
    };

using namespace boost::qvm;
tester<is_scalar<char>::value> t1;
tester<is_scalar<signed char>::value> t2;
tester<is_scalar<unsigned char>::value> t3;
tester<is_scalar<signed short>::value> t4;
tester<is_scalar<unsigned short>::value> t5;
tester<is_scalar<signed int>::value> t6;
tester<is_scalar<unsigned int>::value> t7;
tester<is_scalar<signed long>::value> t8;
tester<is_scalar<unsigned long>::value> t9;
tester<is_scalar<float>::value> t10;
tester<is_scalar<double>::value> t11;
tester<is_scalar<long double>::value> t13;
tester<!is_scalar<vec<float,4> >::value> t14;
tester<!is_scalar<mat<float,4,4> >::value> t15;
tester<!is_scalar<quat<float> >::value> t16;


#include <drake/systems/plants/RigidBodyIK.h>
#include <drake/systems/plants/RigidBodyTree.h>
#include <drake/systems/plants/constraint/RigidBodyConstraint.h>

#include <drake/systems/plants/IKoptions.h>
#include <iostream>
#include <cstdlib>
#include <limits>
#include <boost/shared_ptr.hpp>

using namespace std;
using namespace Eigen;

#include "lcmtypes/bot_core/robot_state_t.hpp"
#include "lcmtypes/drc/robot_plan_w_keyframes_t.hpp"
#include <lcm/lcm-cpp.hpp>

// Find the joint position indices corresponding to 'name'
vector<int> getJointPositionVectorIndices(const RigidBodyTree &model, const std::string &name) {
  shared_ptr<RigidBody> joint_parent_body = model.findJoint(name);
  int num_positions = joint_parent_body->getJoint().getNumPositions();
  vector<int> ret(static_cast<size_t>(num_positions));

  // fill with sequentially increasing values, starting at joint_parent_body->position_num_start:
  iota(ret.begin(), ret.end(), joint_parent_body->position_num_start);
  return ret;
}

void findJointAndInsert(const RigidBodyTree &model, const std::string &name, vector<int> &position_list) {
  auto position_indices = getJointPositionVectorIndices(model, name);

  position_list.insert(position_list.end(), position_indices.begin(), position_indices.end());
}

class App{
  public:
    App(boost::shared_ptr<lcm::LCM> &lcm_);

    ~App(){
    }

    void getRobotState(bot_core::robot_state_t& robot_state_msg, int64_t utime_in, Eigen::VectorXd q, std::vector<std::string> jointNames);

  private:
    boost::shared_ptr<lcm::LCM> lcm_;

};

App::App(boost::shared_ptr<lcm::LCM> &lcm_):
                       lcm_(lcm_){

}

Eigen::Quaterniond euler_to_quat(double roll, double pitch, double yaw) {

  // This conversion function introduces a NaN in Eigen Rotations when:
  // roll == pi , pitch,yaw =0    ... or other combinations.
  // cos(pi) ~=0 but not exactly 0
  // Post DRC Trails: replace these with Eigen's own conversions
  if ( ((roll==M_PI) && (pitch ==0)) && (yaw ==0)){
    return  Eigen::Quaterniond(0,1,0,0);
  }else if( ((pitch==M_PI) && (roll ==0)) && (yaw ==0)){
    return  Eigen::Quaterniond(0,0,1,0);
  }else if( ((yaw==M_PI) && (roll ==0)) && (pitch ==0)){
    return  Eigen::Quaterniond(0,0,0,1);
  }

  double sy = sin(yaw*0.5);
  double cy = cos(yaw*0.5);
  double sp = sin(pitch*0.5);
  double cp = cos(pitch*0.5);
  double sr = sin(roll*0.5);
  double cr = cos(roll*0.5);
  double w = cr*cp*cy + sr*sp*sy;
  double x = sr*cp*cy - cr*sp*sy;
  double y = cr*sp*cy + sr*cp*sy;
  double z = cr*cp*sy - sr*sp*cy;
  return Eigen::Quaterniond(w,x,y,z);
}

void quat_to_euler(Eigen::Quaterniond q, double& roll, double& pitch, double& yaw) {
  const double q0 = q.w();
  const double q1 = q.x();
  const double q2 = q.y();
  const double q3 = q.z();
  roll = atan2(2*(q0*q1+q2*q3), 1-2*(q1*q1+q2*q2));
  pitch = asin(2*(q0*q2-q3*q1));
  yaw = atan2(2*(q0*q3+q1*q2), 1-2*(q2*q2+q3*q3));
}

void App::getRobotState(bot_core::robot_state_t& robot_state_msg, int64_t utime_in, Eigen::VectorXd q, std::vector<std::string> jointNames){
  robot_state_msg.utime = utime_in;

  // Pelvis Pose:
  robot_state_msg.pose.translation.x =q(0);
  robot_state_msg.pose.translation.y =q(1);
  robot_state_msg.pose.translation.z =q(2);

  Eigen::Quaterniond quat = euler_to_quat( q(3), q(4), q(5));

  robot_state_msg.pose.rotation.w = quat.w();
  robot_state_msg.pose.rotation.x = quat.x();
  robot_state_msg.pose.rotation.y = quat.y();
  robot_state_msg.pose.rotation.z = quat.z();

  robot_state_msg.twist.linear_velocity.x  = 0;
  robot_state_msg.twist.linear_velocity.y  = 0;
  robot_state_msg.twist.linear_velocity.z  = 0;
  robot_state_msg.twist.angular_velocity.x = 0;
  robot_state_msg.twist.angular_velocity.y = 0;
  robot_state_msg.twist.angular_velocity.z = 0;

  // Joint States:
  for (size_t i = 0; i < jointNames.size(); i++)  {
    robot_state_msg.joint_name.push_back( jointNames[i] );
    robot_state_msg.joint_position.push_back( q(i) );
    robot_state_msg.joint_velocity.push_back( 0);
    robot_state_msg.joint_effort.push_back( 0 );
  }
  robot_state_msg.num_joints = robot_state_msg.joint_position.size();
}

int main(int argc, char *argv[])
{

  std::string file_path = "fname";
  if ((argc > 1)){
    file_path = argv[1];
  }else{
    std::cout << "you need to provide the path to atlas_minimal_contact.urdf\n";
    return 1;
  }

  boost::shared_ptr<lcm::LCM> lcm(new lcm::LCM() );
  App app(lcm);

  RigidBodyTree model(file_path);

  Vector2d tspan;
  tspan << 0, 1;

  // Default Atlas v5 posture:
  VectorXd qstar(model.num_positions);
  qstar <<
  -0.0260, 0, 0.8440, 0, 0, 0, 0, 0, 0, 0.2700, 0, 0.0550, -0.5700, 1.1300, -0.5500, -0.0550, -1.3300, 2.1530, 0.5000, 0.0985, 0, 0.0008, -0.2700, 0, -0.0550, -0.5700, 1.1300, -0.5500, 0.0550, 1.3300, 2.1530, -0.5000, 0.0985, 0, 0.0008, 0.2564;

  // 1 Back Posture Constraint
  PostureConstraint kc_posture_back(&model, tspan);
  std::vector<int> back_idx;
  findJointAndInsert(model, "back_bkz", back_idx);
  findJointAndInsert(model, "back_bky", back_idx);
  findJointAndInsert(model, "back_bkx", back_idx);
  VectorXd back_lb = VectorXd::Zero(3);
  VectorXd back_ub = VectorXd::Zero(3);
  kc_posture_back.setJointLimits(3, back_idx.data(), back_lb, back_ub);

  // 2 Knees Constraint
  PostureConstraint kc_posture_knees(&model, tspan);
  std::vector<int> knee_idx;
  findJointAndInsert(model, "l_leg_kny", knee_idx);
  findJointAndInsert(model, "r_leg_kny", knee_idx);
  VectorXd knee_lb = VectorXd::Zero(2);
  knee_lb(0) = 1.0; // usually use 0.6
  knee_lb(1) = 1.0; // usually use 0.6
  VectorXd knee_ub = VectorXd::Zero(2);
  knee_ub(0) = 2.5;
  knee_ub(1) = 2.5;
  kc_posture_knees.setJointLimits(2, knee_idx.data(), knee_lb, knee_ub);

  // 3 Left Arm Constraint
  PostureConstraint kc_posture_larm(&model, tspan);
  std::vector<int> larm_idx;
  findJointAndInsert(model, "l_arm_shz", larm_idx);
  findJointAndInsert(model, "l_arm_shx", larm_idx);
  findJointAndInsert(model, "l_arm_ely", larm_idx);
  findJointAndInsert(model, "l_arm_elx", larm_idx);
  findJointAndInsert(model, "l_arm_uwy", larm_idx);
  findJointAndInsert(model, "l_arm_mwx", larm_idx);
  findJointAndInsert(model, "l_arm_lwy", larm_idx);
  VectorXd larm_lb = VectorXd::Zero(7);
  larm_lb(0) = 0.27;
  larm_lb(1) = -1.33;
  larm_lb(2) = 2.153;
  larm_lb(3) = 0.500;
  larm_lb(4) = 0.0985;
  larm_lb(5) = 0;
  larm_lb(6) = 0.0008;
  VectorXd larm_ub = larm_lb;
  kc_posture_larm.setJointLimits(7, larm_idx.data(), larm_lb, larm_ub);

  // 4 Left Foot Position and Orientation Constraints
  int l_foot = model.findLinkId("l_foot");
  Vector3d l_foot_pt = Vector3d::Zero();
  Vector3d lfoot_pos0;
  lfoot_pos0(0) = 0;
  lfoot_pos0(1) = 0.13;
  lfoot_pos0(2) = 0.08;
  Vector3d lfoot_pos_lb = lfoot_pos0;
  lfoot_pos_lb(0) += 0.001;
  lfoot_pos_lb(1) += 0.001;
  lfoot_pos_lb(2) += 0.001;
  Vector3d lfoot_pos_ub = lfoot_pos_lb;
  lfoot_pos_ub(2) += 0.01;
  // std::cout << lfoot_pos0.transpose() << " lfoot\n" ;
  WorldPositionConstraint kc_lfoot_pos(&model, l_foot, l_foot_pt, lfoot_pos_lb, lfoot_pos_ub, tspan);
  Eigen::Vector4d quat_des(1, 0, 0, 0);
  double tol = 0.0017453292519943296;
  WorldQuatConstraint kc_lfoot_quat(&model, l_foot, quat_des, tol, tspan);

  // 5 Right Foot Position and Orientation Constraints
  int r_foot = model.findLinkId("r_foot");
  Vector3d r_foot_pt = Vector3d::Zero();
  Vector3d rfoot_pos0;
  rfoot_pos0(0) = 0;
  rfoot_pos0(1) = -0.13;
  rfoot_pos0(2) = 0.08;
  Vector3d rfoot_pos_lb = rfoot_pos0;
  rfoot_pos_lb(0) += 0.001;
  rfoot_pos_lb(1) += 0.001;
  rfoot_pos_lb(2) += 0.001;
  Vector3d rfoot_pos_ub = rfoot_pos_lb;
  rfoot_pos_ub(2) += 0.001;
  //std::cout << rfoot_pos0.transpose() << " lfoot\n" ;
  WorldPositionConstraint kc_rfoot_pos(&model, r_foot, r_foot_pt, rfoot_pos_lb, rfoot_pos_ub, tspan);
  WorldQuatConstraint kc_rfoot_quat(&model, r_foot, quat_des, tol, tspan);

  // 6 Right Position Constraints (actual reaching constraint)
  int r_hand = model.findLinkId("r_hand");
  Vector3d r_hand_pt = Vector3d::Zero();
  Vector3d rhand_pos0;
  //Vector3d rhand_pos0 = model.forwardKin(r_hand_pt, r_hand, 0, 0);
  rhand_pos0(0) = 0.2;
  rhand_pos0(1) = -0.5;
  rhand_pos0(2) = 0.4;
  Vector3d rhand_pos_lb = rhand_pos0;
  rhand_pos_lb(0) += 0.05;
  rhand_pos_lb(1) += 0.05;
  rhand_pos_lb(2) += 0.05;
  Vector3d rhand_pos_ub = rhand_pos_lb;
  rhand_pos_ub(2) += 0.05;
  //std::cout << rhand_pos_ub.transpose() << " rhand\n" ;
  WorldPositionConstraint kc_rhand(&model, r_hand, r_hand_pt, rhand_pos_lb, rhand_pos_ub, tspan);

  // 7 QuasiStatic Constraints
  QuasiStaticConstraint kc_quasi(&model, tspan);
  kc_quasi.setShrinkFactor(0.2);
  kc_quasi.setActive(true);
  Eigen::Matrix3Xd l_foot_pts = Eigen::Matrix3Xd::Zero(3, 4);
  l_foot_pts << -0.0820, -0.0820, 0.1780, 0.1780,
      0.0624, -0.0624, 0.0624, -0.0624,
      -0.0811, -0.0811, -0.0811, -0.0811;
  kc_quasi.addContact(1, &l_foot, &l_foot_pts);
  Eigen::Matrix3Xd r_foot_pts = Eigen::Matrix3Xd::Zero(3, 4);
  r_foot_pts << -0.0820, -0.0820, 0.1780, 0.1780,
      0.0624, -0.0624, 0.0624, -0.0624,
      -0.0811, -0.0811, -0.0811, -0.0811;
  kc_quasi.addContact(1, &r_foot, &r_foot_pts);

  std::vector<RigidBodyConstraint *> constraint_array;
  constraint_array.push_back(&kc_quasi);
  constraint_array.push_back(&kc_posture_knees);
  constraint_array.push_back(&kc_lfoot_pos);
  constraint_array.push_back(&kc_lfoot_quat);
  constraint_array.push_back(&kc_rfoot_pos);
  constraint_array.push_back(&kc_rfoot_quat);
  constraint_array.push_back(&kc_rhand);
  constraint_array.push_back(&kc_posture_larm);
  constraint_array.push_back(&kc_posture_back);

  IKoptions ikoptions(&model);
  VectorXd q_sol(model.num_positions);
  int info;
  vector<string> infeasible_constraint;
  inverseKin(&model, qstar, qstar, constraint_array.size(), constraint_array.data(), q_sol, info, infeasible_constraint, ikoptions);
  printf("INFO = %d\n", info);
  if (info != 1) {
    return 1;
  }

  /////////////////////////////////////////
  KinematicsCache<double> cache = model.doKinematics(q_sol);
  Vector3d com = model.centerOfMass(cache);
  printf("%5.2f\n%5.2f\n%5.2f\n",com(0),com(1),com(2));

  bot_core::robot_state_t robot_state_msg;
  std::vector<string> jointNames;
  for (int i=0 ; i <model.num_positions ; i++){
    // std::cout << model.getPositionName(i) << " " << i << "\n";
    jointNames.push_back( model.getPositionName(i) ) ;
  }

  app.getRobotState(robot_state_msg, 0*1E6, q_sol , jointNames);

  lcm->publish("EST_ROBOT_STATE",&robot_state_msg);
  return 0;
}


// This file is part of VIO - Semi-direct Visual Odometry.
//
// Copyright (C) 2014 Christian Forster <forster at ifi dot uzh dot ch>
// (Robotics and Perception Group, University of Zurich, Switzerland).
//
// VIO is free software: you can redistribute it and/or modify it under the
// terms of the GNU General Public License as published by the Free Software
// Foundation, either version 3 of the License, or any later version.
//
// VIO is distributed in the hope that it will be useful, but WITHOUT ANY
// WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#include <vio/abstract_camera.h>
#include <stdlib.h>
#include <Eigen/StdVector>
#include <boost/bind.hpp>
#include <fstream>
#include <vio/frame_handler_base.h>
#include <vio/config.h>
#include <vio/feature.h>
#include <vio/matcher.h>
#include <vio/map.h>
#include <vio/point.h>

namespace vio
{

// definition of global and static variables which were declared in the header

FrameHandlerBase::FrameHandlerBase() :
  stage_(STAGE_PAUSED),
  set_reset_(false),
  set_start_(false),
  num_obs_last_(0)
{
}

FrameHandlerBase::~FrameHandlerBase()
{
}

bool FrameHandlerBase::startFrameProcessingCommon(const double timestamp)
{
  if(set_start_)
  {
    resetAll();
    stage_ = STAGE_FIRST_FRAME;
  }

  if(stage_ == STAGE_PAUSED)
    return false;

  map_.emptyTrash();
  return true;
}

int FrameHandlerBase::finishFrameProcessingCommon(
    const size_t update_id,
    const UpdateResult dropout,
    const size_t num_observations)
{

  if(stage_ == STAGE_DEFAULT_FRAME)
  num_obs_last_ = num_observations;

  if(dropout == RESULT_FAILURE &&
      (stage_ == STAGE_DEFAULT_FRAME || stage_ == STAGE_RELOCALIZING ))
  {
    //stage_ = STAGE_RELOCALIZING;
  }
  if(set_reset_)
    resetAll();

  return 0;
}

void FrameHandlerBase::resetCommon()
{
  map_.reset();
  stage_ = STAGE_PAUSED;
  set_reset_ = false;
  set_start_ = false;
  num_obs_last_ = 0;
}

bool ptLastOptimComparator(Point* lhs, Point* rhs)
{
  return (lhs->last_structure_optim_ < rhs->last_structure_optim_);
}

void FrameHandlerBase::optimizeStructure(
    FramePtr frame,
    size_t max_n_pts,
    int max_iter)
{
  for(auto&& it:frame->fts_){
      if(it->point==NULL)continue;
      if(it->point->obs_.size()<2){
          it->point->last_frame_overlap_id_= frame->id_;
          continue;
      }
      it->point->optimize(max_iter);
      it->point->last_frame_overlap_id_= frame->id_;
  }
}
void FrameHandlerBase::posEdit(
            FramePtr frame)
    {
        for(auto&& it:frame->fts_){
            if(it->point!=NULL && !it->point->pos_.hasNaN() && it->point->pos_.norm() !=0.){
                if(it->point->obs_.size()==2){
                    Eigen::Matrix<double,4,4> A,frame_a,frame_b;
                    frame_a=it->point->obs_.front()->frame->se3().matrix();
                    frame_b=it->point->obs_.back()->frame->se3().matrix();
                    A.row(0) = it->point->obs_.front()->f(0) * frame_a.row(2) - it->point->obs_.front()->f(2) * frame_a.row(0);
                    A.row(1) = it->point->obs_.front()->f(1) * frame_a.row(2) - it->point->obs_.front()->f(2) * frame_a.row(1);
                    A.row(2) = it->point->obs_.back()->f(0) * frame_b.row(2) - it->point->obs_.back()->f(2) * frame_b.row(0);
                    A.row(3) = it->point->obs_.back()->f(1) * frame_b.row(2) - it->point->obs_.back()->f(2) * frame_b.row(1);
                    // Aを特異値分解する (A = U S Vt)
                    // https://eigen.tuxfamily.org/dox/classEigen_1_1JacobiSVD.html
                    Eigen::JacobiSVD<Eigen::Matrix<double,4,4>> svd(A, Eigen::ComputeFullU | Eigen::ComputeFullV);
                    const Eigen::Matrix<double,4,1> singular_vector = svd.matrixV().block<4, 1>(0, 3);
                    it->point->pos_=singular_vector.block<3, 1>(0, 0) / singular_vector(3);
                }
                //if(frame->w2f(it->point->pos_).z()<1e-9)map_.safeDeletePoint(it->point);
            }else{
                map_.safeDeletePoint(it->point);
            }
        }
    }

} // namespace vio


/*
 * This file is a part of
 *
 * ============================================
 * ###   Pteros molecular modeling library  ###
 * ============================================
 *
 * (C) 2009-2018, Semen Yesylevskyy
 *
 * All works, which use Pteros, should cite the following papers:
 *
 *  1.  Semen O. Yesylevskyy, "Pteros 2.0: Evolution of the fast parallel
 *      molecular analysis library for C++ and python",
 *      Journal of Computational Chemistry, 2015, 36(19), 1480–1488.
 *      doi: 10.1002/jcc.23943.
 *
 *  2.  Semen O. Yesylevskyy, "Pteros: Fast and easy to use open-source C++
 *      library for molecular analysis",
 *      Journal of Computational Chemistry, 2012, 33(19), 1632–1636.
 *      doi: 10.1002/jcc.22989.
 *
 * This is free software distributed under Artistic License:
 * http://www.opensource.org/licenses/artistic-license-2.0.php
 *
*/

#include "selection_parser.h"
#include "pteros/core/system.h"
#include "pteros/core/selection.h"
#include "pteros/core/logging.h"
#include "pteros/core/pteros_error.h"
#include "pteros/core/distance_search.h"
#include <Eigen/Core>
#include <boost/algorithm/string.hpp>
#include <boost/range/counting_range.hpp>
#include <unordered_set>
#include <regex>
#include <list>

using namespace std;
using namespace pteros;
using namespace boost;

//===============================================

// We derive from normal peg::parser
class Pteros_PEG_parser: public peg::parser {
public:
    Pteros_PEG_parser(const char* s): peg::parser(s) {
        enable_ast<MyAst>();
        enable_packrat_parsing();
        log = [&](size_t ln, size_t col, const string& msg) {
            error_message = fmt::format("{}:{}",col,msg);
        };
    }

    virtual ~Pteros_PEG_parser(){}

    string error_message;
};

// Instance of the parser itself
Pteros_PEG_parser _parser(R"(
        LOGICAL_EXPR       <-  LOGICAL_OPERAND (LOGICAL_OPERATOR LOGICAL_OPERAND)*
        LOGICAL_OPERATOR   <-  < 'or' / 'and' >
        LOGICAL_OPERAND    <-  (NOT / BY)? ( '(' LOGICAL_EXPR ')' / ALL / NUM_COMPARISON / KEYWORD_EXPR / WITHIN )
        ALL                <-  < 'all' >
        NOT                <-  < 'not' >
        BY                 <-  'by' < ('residue' / 'chain' / 'mol') > / 'same' < ('residue' / 'chain' / 'mol') > 'as'

        NUM_COMPARISON     <-  NUM_EXPR COMPARISON_OPERATOR NUM_EXPR (COMPARISON_OPERATOR NUM_EXPR)?
        COMPARISON_OPERATOR <- <  '==' / '<=' / '>=' / '<>' / '!=' / '<' / '>' / '=' >
        NUM_EXPR           <- NUM_TERM (PLUS_MINUS NUM_TERM)*
        NUM_TERM           <- NUM_POWER (DIV_MUL NUM_POWER)*
        NUM_POWER          <- NUM_FACTOR (POW NUM_FACTOR)?
        NUM_FACTOR         <- UNARY_MINUS? ( '(' NUM_EXPR ')' / X / Y / Z / BETA / OCC / RESINDEX / INDEX
                                                              / RESID / DIST / MASS / CHARGE ) / FLOAT
        PLUS_MINUS         <- < '+' / '-' >
        DIV_MUL            <- < '*' / '/' >
        POW                <- < '^' / '**' >
        UNARY_MINUS        <- < '-' >

        X                  <-  ('X' / 'x') ('of' COM)?
        Y                  <-  ('Y' / 'y') ('of' COM)?
        Z                  <-  ('Z' / 'z') ('of' COM)?
        BETA               <- < 'beta' >
        OCC                <- < 'occupancy' / 'occ' >
        RESINDEX           <- < 'resindex' >
        INDEX              <- < 'index' >
        RESID              <- < 'resid' >
        MASS               <- < 'mass' >
        CHARGE             <- < 'charge' >

        COM                <- COM_TYPE PBC? 'of' LOGICAL_OPERAND
        COM_TYPE           <- < 'com' / 'cog' >

        DIST               <- ('dist' / 'distance') (POINT / VECTOR / PLANE)
        POINT              <- 'point' PBC? (FLOAT FLOAT FLOAT / COM)
        VECTOR             <- 'vector' PBC? (FLOAT FLOAT FLOAT / COM) FLOAT FLOAT FLOAT
        PLANE              <- 'plane' PBC? (FLOAT FLOAT FLOAT / COM) FLOAT FLOAT FLOAT

        FLOAT              <- < INTEGER ('.' [0-9]+ )? ( ('e' / 'E' ) INTEGER )? >
        INTEGER            <- < ('-' / '+')? [0-9]+ >

        KEYWORD_EXPR       <- STR_KEYWORD_EXPR / INT_KEYWORD_EXPR
        STR_KEYWORD_EXPR   <- STR_KEYWORD (STR / REGEX)+
        STR_KEYWORD        <- < 'name' / 'resname' / 'tag' / 'chain' / 'type' >
        STR                <- !('or'/'and') < [a-zA-Z0-9]+ >

        INT_KEYWORD_EXPR   <- INT_KEYWORD (RANGE / INTEGER)+
        INT_KEYWORD        <- < 'resindex' / 'index' / 'resid' >
        RANGE              <- INTEGER ('-'/'to'/':') INTEGER

        WITHIN             <- 'within' FLOAT (PBC SELF / SELF PBC / PBC / SELF)? 'of' LOGICAL_OPERAND

        PBC                <- < 'pbc' / 'nopbc' / 'periodic' / 'nonperiodic' >
        SELF               <- < 'self' / 'noself' >

        %whitespace      <-  [ \t\r\n]*

        REGEX              <- '"' <(!'"' .)*> '"' / "'" <(!"'" .)*> "'"
    )");

//===============================================

bool is_node_coordinate_dependent(const std::shared_ptr<MyAst>& node){
    if(node->name == "X" || node->name == "Y" || node->name == "Z" || node->name == "WITHIN"
            || node->name == "POINT" || node->name == "PLANE" || node->name == "VECTOR" || node->name == "COM"
      ){
        return true;
    } else {
        return false;
    }
}

Selection_parser::Selection_parser(std::vector<int> *subset):
    has_coord(false),
    starting_subset(subset),
    sys(nullptr)
{

}

Selection_parser::~Selection_parser(){}

void set_coord_dependence(const std::shared_ptr<MyAst>& node){
    node->is_coord_dependent = is_node_coordinate_dependent(node);
    if(node->nodes.size()){
        // tree
        node->is_coord_dependent = is_node_coordinate_dependent(node);
        for(int i=0;i<node->nodes.size();++i){
            set_coord_dependence(node->nodes[i]);
            if(node->nodes[i]->is_coord_dependent) node->is_coord_dependent = true;
        }
    }

    //cout << node->name << " :: " << node->is_coord_dependent << endl;
}

void Selection_parser::optimize(std::shared_ptr<MyAst>& node){
    // optimize arithmetics
    if(node->name == "NUM_EXPR" || node->name == "NUM_TERM" || node->name == "NUM_POWER") {
        // If not-coord dependent just optimize
        if(!node->is_coord_dependent){
            // Replace with float node
            node = std::make_shared<MyAst>("",0,0,"FLOAT", fmt::format("{}", get_numeric(node)(0)));
            node->is_coord_dependent = false; // Keep correct flag just in case
        }
    }

    // Convert chained logic into a tree
    if(node->name == "LOGICAL_EXPR" && node->nodes.size()>3){
        // second operand and all after it should become new node
        vector<std::shared_ptr<MyAst>> ch;
        copy(node->nodes.begin()+2,node->nodes.end(),back_inserter(ch));
        auto operand2 = std::make_shared<MyAst>("",0,0,"LOGICAL_EXPR",ch);
        // Inherit coord dependence for this new node (important in order not to mislead precomputer)
        operand2->is_coord_dependent = node->is_coord_dependent;
        // keep only 3 nodes
        node->nodes.resize(3);
        // Put new second operand
        node->nodes[2] = operand2;
    }

    // Recurse into children
    for(int i=0;i<node->nodes.size();++i) optimize(node->nodes[i]);
}

void Selection_parser::precompute(std::shared_ptr<MyAst>& node){
    if(    node->name!="PRE"
        && node->name!="NUM_COMPARISON"
        && node->name!="STR_KEYWORD_EXPR"
        && node->name!="INT_KEYWORD_EXPR"
        && node->name!="LOGICAL_EXPR"
        && node->name!="LOGICAL_OPERAND"
        && node->name!="ALL"
        && node->name!="WITHIN"
        && node->name!="BYRES") return;

    if(!node->is_coord_dependent){
        auto ast = std::make_shared<MyAst>("",0,0,"PRE","");
        eval_node(node, ast->precomputed);
        node = ast;
    } else if(!node->nodes.empty()) {
        for(int i=0;i<node->nodes.size();++i) precompute(node->nodes[i]);
    }
}

void Selection_parser::create_ast(string& sel_str, System* system){
    if (_parser.parse(sel_str.c_str(), tree)) {
        tree = peg::AstOptimizer(true,{"POINT","X","Y","Z"}).optimize(tree);

        //cout << sel_str << endl;
        //cout << peg::ast_to_s(tree) << endl;

        set_coord_dependence(tree);
    } else {
        throw Pteros_error(_parser.error_message);
    }

    if(tree->is_coord_dependent) has_coord = true; // Global coord dependence

    sys = system;
    Natoms = sys->num_atoms();

    if(starting_subset)
        current_subset = starting_subset;
    else
        current_subset = nullptr;

    // Optimize tree
    optimize(tree);

    //cout << peg::ast_to_s(tree) << endl;

    // proceed with optimizing pure nodes to precomputed if needed
    if(has_coord) precompute(tree);
}

void Selection_parser::apply_ast(size_t fr, vector<int>& result){
    frame = fr;
    eval_node(tree,result);
}

void Selection_parser::eval_node(const std::shared_ptr<MyAst> &node, std::vector<int>& result){

    result.clear();

    // Here starts evaluation

    if(node->name == "PRE"){ // Precomputed nodes are created during optimization stage
        result = node->precomputed;
    }

    else if(node->name == "NUM_COMPARISON")
    {
        vector<string> c; // comparison operators
        vector<std::function<float(int)>> op; // comparison operands
        vector<std::function<bool(float,float)>> comparison; // Function(s) to evaluate

        if(node->nodes.size() == 3){ // simple
            c.resize(1);
            op.resize(2);
            comparison.resize(1);

            c[0]  = node->nodes[1]->token;
            op[0] = get_numeric(node->nodes[0]);
            op[1] = get_numeric(node->nodes[2]);

        } else { // chained
            c.resize(2);
            op.resize(3);
            comparison.resize(2);

            c[0]  = node->nodes[1]->token;
            c[1]  = node->nodes[3]->token;
            op[0] = get_numeric(node->nodes[0]);
            op[1] = get_numeric(node->nodes[2]);
            op[2] = get_numeric(node->nodes[4]);
        }

        for(int i=0;i<c.size();++i){
            if(c[i] == "=" || c[i]== "=="){
                comparison[i] = [](float a, float b){ return a==b; };
            } else if (c[i] == "!=" || c[i]=="<>"){
                comparison[i] = [](float a, float b){ return a!=b; };
            } else if (c[i] == "<"){
                comparison[i] = [](float a, float b){ return a<b; };
            } else if (c[i] == ">"){
                comparison[i] = [](float a, float b){ return a>b; };
            } else if (c[i] == "<="){
                comparison[i] = [](float a, float b){ return a<=b; };
            } else if (c[i] == ">="){
                comparison[i] = [](float a, float b){ return a>=b; };
            }
        }

        if(!current_subset) {
            if(node->nodes.size() == 3){ // simple
                for(int at=0;at<Natoms;++at)
                    if( comparison[0](op[0](at),op[1](at)) ) result.push_back(at);
            } else { // chained
                for(int at=0;at<Natoms;++at)
                    if( comparison[0](op[0](at),op[1](at)) && comparison[1](op[1](at),op[2](at)) ) result.push_back(at);
            }
        } else {
            if(node->nodes.size() == 3){ // simple
                for(int at: *current_subset)
                    if( comparison[0](op[0](at),op[1](at)) ) result.push_back(at);
            } else { // chained
                for(int at: *current_subset)
                    if( comparison[0](op[0](at),op[1](at)) && comparison[1](op[1](at),op[2](at)) ) result.push_back(at);
            }
        }
    }

    else if(node->name == "STR_KEYWORD_EXPR")
    {
        std::function<bool(int,const string&)> comp_func_str;
        std::function<bool(int,const std::regex&)> comp_func_regex;

        const string& keyword = node->nodes[0]->token;

        if(keyword == "name"){
            comp_func_str = [this](int at, const string& str){ return sys->atoms[at].name == str; };
            comp_func_regex = [this](int at, const std::regex& reg){ return std::regex_match(sys->atoms[at].name.c_str(),reg); };
        } else if(keyword == "type"){
            comp_func_str = [this](int at, const string& str){ return sys->atoms[at].type_name == str; };
            comp_func_regex = [this](int at, const std::regex& reg){ return std::regex_match(sys->atoms[at].type_name.c_str(),reg); };
        } else if(keyword == "resname"){
            comp_func_str = [this](int at, const string& str){ return sys->atoms[at].resname == str; };
            comp_func_regex = [this](int at, const std::regex& reg){ return std::regex_match(sys->atoms[at].resname.c_str(),reg); };
        } else if(keyword == "tag"){
            comp_func_str = [this](int at, const string& str){ return sys->atoms[at].tag== str; };
            comp_func_regex = [this](int at, const std::regex& reg){ return std::regex_match(sys->atoms[at].tag.c_str(),reg); };
        } else if(keyword == "chain"){
            comp_func_str = [this](int at, const string& str){ return sys->atoms[at].chain == str[0]; };
            comp_func_regex = [this](int at, const std::regex& reg){
                string s(" ");
                s[0] = sys->atoms[at].chain;
                return std::regex_match(s.c_str(),reg);
            };
        }

        list<string> str_values;
        list<std::regex> regex_values;
        for(int i=1;i<node->nodes.size();++i){
            if(node->nodes[i]->name=="STR")
                str_values.emplace_back(node->nodes[i]->token);
            else
                regex_values.emplace_back(node->nodes[i]->token);
        }

        // Loop body
        auto body = [&](int at){
            // Cycle over regex values
            bool reg_found = false;
            for(const auto& reg: regex_values){
                if(comp_func_regex(at,reg)){
                    result.push_back(at);
                    reg_found = true;
                    break;
                }
            }

            // If at least one regex matched no need to proceed with strings
            if(reg_found) return;

            // Cycle over string values
            for(const auto& str: str_values){
                if(comp_func_str(at,str)){
                    result.push_back(at);
                    break;
                }
            }
        };

        if(!current_subset){
            for(int at=0;at<Natoms;++at) body(at);
        } else {
            for(int at: *current_subset) body(at);
        }

        // Sort unique
        sort(result.begin(),result.end());
        vector<int>::iterator it = unique(result.begin(), result.end());
        result.resize( it - result.begin() );
    }

    //---------------------------------------------------------------------------
    else if(node->name == "INT_KEYWORD_EXPR")
    {
        const string& keyword = node->nodes[0]->token;
        int Nchildren = node->nodes.size(); // Get number of children

        // If starting subset is present than this is a subselection and
        // we have to interpret indexes as local indexes!
        if(keyword == "index") {
            // Cycle over children
            for(int i=1;i<Nchildren;++i){
                if(node->nodes[i]->name == "INTEGER") {
                    int k = stoi(node->nodes[i]->token);
                    // Shift to local index for subselection if needed
                    if(starting_subset) k+=(*starting_subset)[0];
                    // We have to check the range here
                    if(k>=0 && k<Natoms)
                        result.push_back(k);
                } else {
                    // this is a range, not an integer
                    int i1 = stoi(node->nodes[i]->nodes[0]->token);
                    int i2 = stoi(node->nodes[i]->nodes[1]->token);
                    // Shift to local index for subselection if needed
                    if(starting_subset){
                        i1+=(*starting_subset)[0];
                        i2+=(*starting_subset)[0];
                    }
                    for(int k=i1;k<=i2;++k)
                        // We have to check the range here
                        if(k>=0 && k<Natoms)
                            result.push_back(k);
                }
            }

        } else { // resid or resindex

            // Make lists
            vector<int> int_list;
            vector<int> range_list;
            for(int i=1;i<Nchildren;++i){
                if(node->nodes[i]->name == "INTEGER") {
                    int_list.push_back(stoi(node->nodes[i]->token));
                } else {
                    range_list.push_back(stoi(node->nodes[i]->nodes[0]->token));
                    range_list.push_back(stoi(node->nodes[i]->nodes[1]->token));
                }
            }

            // Comparison function
            std::function<bool(int,int)> comp_func;
            if(keyword == "resid"){
                comp_func = [this](int at, int k){ return sys->atoms[at].resid == k; };
            } else if(keyword == "resindex"){
                comp_func = [this](int at, int k){ return sys->atoms[at].resindex == k; };
            }

            // Loop body
            auto body = [&](int at){
                // Individual numbers
                for(int k: int_list)
                    if(comp_func(at,k)) result.push_back(at);
                // Ranges
                for(int i=0;i<range_list.size();i+=2){ // Itarage by pair
                    for(int k=range_list[i];k<=range_list[i+1];++k){ // Inside range
                        if(comp_func(at,k)) result.push_back(at);
                    }
                }
            };

            // Do loop
            if(!current_subset){
                for(int at=0;at<Natoms;++at) body(at);
            } else {
                for(int at: *current_subset) body(at);
            }
        } // index if

        sort(result.begin(),result.end());
        vector<int>::iterator it = unique(result.begin(), result.end());
        result.resize( it - result.begin() );
    }

    //---------------------------------------------------------------------------

    else if(node->name == "LOGICAL_EXPR")
    {
        if(node->nodes[1]->token == "or") {
            vector<int> res1,res2;
            eval_node(node->nodes[0],res1);
            eval_node(node->nodes[2],res2);
            std::set_union(res1.begin(),res1.end(),res2.begin(),res2.end(),back_inserter(result));

        } else if(node->nodes[1]->token == "and") {
            // Optimize to put pure node first
            bool pure1 = !node->nodes[0]->is_coord_dependent;
            bool pure2 = !node->nodes[2]->is_coord_dependent;

            if(pure2 && !pure1) std::swap(node->nodes[0],node->nodes[2]);

            vector<int> res1,res2;
            eval_node(node->nodes[0],res1);
            current_subset = &res1; // Set subset for second
            eval_node(node->nodes[2],res2); // Is using filled current subset

            std::set_intersection(res1.begin(),res1.end(),res2.begin(),res2.end(),back_inserter(result));

            // Reset subset
            if(starting_subset){
                current_subset = starting_subset;
            } else {
                current_subset = nullptr;
            }
        }
    }

    else if(node->name == "LOGICAL_OPERAND")
    {
        vector<int> res;
        eval_node(node->nodes[1],res);

        if(node->nodes[0]->name == "NOT"){
            if(!current_subset){
                auto r = boost::counting_range(0,Natoms);
                std::set_difference(r.begin(),r.end(), res.begin(),res.end(), back_inserter(result));
            } else {
                // For subset
                std::set_difference(current_subset->begin(),current_subset->end(), res.begin(),res.end(), back_inserter(result));
            }

        } else if(node->nodes[0]->name == "BY"){
            if(node->nodes[0]->token == "residue"){
                // First make a set of resids we need to search
                std::unordered_set<int> resind;
                for(auto at: res) resind.insert(sys->atoms[at].resindex);

                // Now cycle over all atoms in the starting subset if present (not current subset!!!)
                if(starting_subset){
                    for(int at: *starting_subset) // over starting subset
                        if(resind.count(sys->atoms[at].resindex)) result.push_back(at);
                } else {
                    for(int at=0;at<Natoms;++at) // over all atoms
                        if(resind.count(sys->atoms[at].resindex)) result.push_back(at);
                }

            } else if(node->nodes[0]->token == "chain") {
                // First make a set of chains we need to search
                std::unordered_set<char> chains;
                for(auto at: res) chains.insert(sys->atoms[at].chain);

                // Now cycle over all atoms in the starting subset if present (not current subset!!!)
                if(starting_subset){
                    for(int at: *starting_subset) // over starting subset
                        if(chains.count(sys->atoms[at].chain)) result.push_back(at);
                } else {
                    for(int at=0;at<Natoms;++at) // over all atoms
                        if(chains.count(sys->atoms[at].chain)) result.push_back(at);
                }

            } else if(node->nodes[0]->token == "mol") {
                if(!sys->force_field.ready) throw Pteros_error("Can't select by molecule: no topology!");

                // set of mols we need to include
                std::unordered_set<int> mols;

                // go over res and add mol is we need it
                for(auto at: res){
                    // Cycle over all molecules in the system
                    for(int j=0; j<sys->force_field.molecules.size(); ++j){ // Over molecules
                        if(at>=sys->force_field.molecules[j](0) && at<=sys->force_field.molecules[j](1)){
                            mols.insert(j);
                            break;
                        }
                    }
                }

                // Now add all chosen molecules
                for(int m: mols)
                    for(int at=sys->force_field.molecules[m](0); at<=sys->force_field.molecules[m](1); ++at)
                        result.push_back(at);

                // Sort and restrict to starting subset (!) if needed
                sort(result.begin(),result.end());
                if(starting_subset)
                    std::set_intersection(result.begin(),result.end(),
                                          starting_subset->begin(),starting_subset->end(),
                                          back_inserter(result));

            } // mol
        } //BY
    }

    else if(node->name == "ALL")
    {
        result.resize(Natoms);
        for(int at=0;at<Natoms;++at) result[at] = at;
    }

    else if(node->name == "WITHIN")
    {
        float cutoff = stof(node->nodes[0]->token);
        bool periodic = false;
        bool include_self = true;

        int eval_ind;

        if(node->nodes.size() == 4){ // Both pbc and self
            if(node->nodes[1]->name == "PBC" && (node->nodes[1]->token == "pbc" || node->nodes[1]->token == "periodic"))
                periodic = true;

            if(node->nodes[2]->name == "PBC" && (node->nodes[2]->token == "pbc" || node->nodes[2]->token == "periodic"))
                periodic = true;

            if(node->nodes[1]->name == "SELF" && node->nodes[1]->token == "noself")
                include_self = false;

            if(node->nodes[2]->name == "SELF" && node->nodes[2]->token == "noself")
                include_self = false;

            eval_ind = 3;

        } else if(node->nodes.size() == 3){ // Either pbc or self
            if(node->nodes[1]->name == "PBC" && (node->nodes[1]->token == "pbc" || node->nodes[1]->token == "periodic"))
                periodic = true;

            if(node->nodes[1]->name == "SELF" && node->nodes[1]->token == "noself")
                include_self = false;

            eval_ind = 2;

        } else { // Neither pbc nor self
            eval_ind = 1;
        }

        Selection dum1(*sys), dum2(*sys);
        // Result is returned directly into the index array of selection dum2
        // thus no additional copying
        eval_node(node->nodes[eval_ind],dum2._index);

        // Prepare selection dum1
        if(!current_subset){
            // We are NOT limited by subspace
            dum1._index.resize(Natoms);
            for(int i=0;i<Natoms;++i) dum1._index[i] = i;
        } else {
            // We are limited by subspace
            dum1._index = *current_subset;
        }

        // Set frame for both selections
        dum1.set_frame(frame);
        dum2.set_frame(frame);

        search_within(cutoff,dum1,dum2,result,include_self,periodic);
    }

    else
    {
        throw Pteros_error("Unknown node {}!",node->name);
    }
}

//returns a 3-vector
Eigen::Vector3f Selection_parser::get_vector(const std::shared_ptr<MyAst> &node){
    if(node->name == "COM")
    {
        bool with_mass = (node->nodes[0]->token == "com") ? true : false;

        auto pbc = noPBC;
        if(node->nodes[1]->name == "PBC" && (node->nodes[1]->token == "pbc" || node->nodes[1]->token == "periodic")) pbc = fullPBC;

        Selection sel(*sys);
        // We have to ignore current subset here!
        // Reset subset
        auto old_subset = current_subset;
        current_subset = nullptr;

        eval_node(node->nodes.back(), sel._index);
        sel.set_frame(frame);

        current_subset = old_subset;

        return sel.center(with_mass,pbc);
    }
}

// Returns callable, which returns value for numeric node for atom at
std::function<float(int)> Selection_parser::get_numeric(const std::shared_ptr<MyAst> &node){

    std::function<float(int)> res;

    // terminals
    if(node->name == "INTEGER"){
        float val = stol(node->token);
        res = [val](int at){ return val; };
    } else if(node->name == "FLOAT"){
        float val = stof(node->token);
        res = [val](int at){ return val; };
    } else if(node->name == "X"){
        if(node->nodes.empty())
            res =[this](int at){ return sys->traj[frame].coord[at](0); };
        else {
            float x = get_vector(node->nodes[0])[0];
            res =[x](int at){ return x; };
        }
    } else if(node->name == "Y"){
        if(node->nodes.empty())
            res =[this](int at){ return sys->traj[frame].coord[at](1); };
        else {
            float y = get_vector(node->nodes[0])[1];
            res =[y](int at){ return y; };
        }
    } else if(node->name == "Z"){
        if(node->nodes.empty())
            res =[this](int at){ return sys->traj[frame].coord[at](2); };
        else {
            float z = get_vector(node->nodes[0])[2];
            res =[z](int at){ return z; };
        }
    } else if(node->name == "BETA"){
        res = [this](int at){ return sys->atoms[at].beta; };
    } else if(node->name == "OCC"){
        res = [this](int at){ return sys->atoms[at].occupancy; };
    } else if(node->name == "INDEX"){
        res = [](int at){ return at; };
    } else if(node->name == "RESINDEX"){
        res = [this](int at){ return sys->atoms[at].resindex; };
    } else if(node->name == "RESID"){
        res = [this](int at){ return sys->atoms[at].resid; };
    } else if(node->name == "MASS"){
        res = [this](int at){ return sys->atoms[at].mass; };
    } else if(node->name == "CHARGE"){
        res = [this](int at){ return sys->atoms[at].charge; };
    // Compounds
    } else if(node->name == "UNARY_MINUS"){
        auto func = get_numeric(node->nodes[0]);
        res = [func](int at){ return -func(at); };

    } else if(node->name == "NUM_EXPR" || node->name == "NUM_TERM" || node->name == "NUM_POWER") {

        int N = node->nodes.size();

        vector<std::function<float(int)>> operands((N-1)/2+1);
        vector<std::function<float(float,float)>> operators((N-1)/2);

        for(int i=0;i<operands.size();++i){
            operands[i] = get_numeric(node->nodes[i*2]);
        }

        for(int i=0;i<operators.size();++i){
            auto op = node->nodes[i*2+1]->token;
            if     (op=="+") {
                operators[i] = [](float a,float b){ return a+b; };
            } else if(op=="-") {
                operators[i] = [](float a,float b){ return a-b; };
            } else if(op=="*") {
                operators[i] = [](float a,float b){ return a*b; };
            } else if(op=="/") {
                operators[i] = [](float a,float b){
                    if(b==0.0) throw Pteros_error("Division by zero in selection!");
                    return a/b;
                };
            } else if(op=="^" || op=="**") {
                operators[i] = [](float a,float b){ return std::pow(a,b); };
            }
        }

        res = [operands,operators](int at){
            float result = operands[0](at);
            for(int i=0;i<operators.size();++i) result = operators[i](result,operands[i+1](at));
            return result;
        };

    } else if(node->name == "POINT") {
        Eigen::Vector3f p;

        int N = node->nodes.size();
        int offset = 0;

        bool pbc = false;
        if(node->nodes[0]->token == "pbc" || node->nodes[0]->token == "periodic"){
            pbc = true;
            offset = 1;
        }

        if(N==1 || N==2){ // com and possibly pbc
            p = get_vector(node->nodes[0+offset]);
        } else { // 3 floats and possibly pbc
            p(0) = get_numeric(node->nodes[0+offset])(0);
            p(1) = get_numeric(node->nodes[1+offset])(0);
            p(2) = get_numeric(node->nodes[2+offset])(0);
        }

        // Return distance
        if(pbc){
            res = [this,p](int at){
                return sys->box(frame).distance(p, sys->traj[frame].coord[at]);
            };
        } else {
            res = [this,p](int at){
                return (p - sys->traj[frame].coord[at]).norm();
            };
        }

    } else if(node->name == "VECTOR" || node->name == "PLANE") {
        Eigen::Vector3f p,dir;

        int N = node->nodes.size();
        int offset = 0;

        bool pbc = false;
        if(node->nodes[0]->token == "pbc" || node->nodes[0]->token == "periodic"){
            pbc = true;
            offset = 1;
        }

        if(N==4 || N==5){ // com + 3 floats
            p = get_vector(node->nodes[0+offset]);
            dir(0) = get_numeric(node->nodes[1+offset])(0);
            dir(1) = get_numeric(node->nodes[2+offset])(0);
            dir(2) = get_numeric(node->nodes[3+offset])(0);
        } if(N==6){ // 6 floats
            p(0) = get_numeric(node->nodes[0+offset])(0);
            p(1) = get_numeric(node->nodes[1+offset])(0);
            p(2) = get_numeric(node->nodes[2+offset])(0);
            dir(0) = get_numeric(node->nodes[3+offset])(0);
            dir(1) = get_numeric(node->nodes[4+offset])(0);
            dir(2) = get_numeric(node->nodes[5+offset])(0);
        }

        if(node->name == "VECTOR"){
            // For vector
            res = [this,p,dir,pbc](int at){
                Eigen::Vector3f atom = sys->traj[frame].coord[at];
                // Get vector from p to current atom
                Eigen::Vector3f v = atom - p;
                // Project v onto dir
                v = (v.dot(dir)/dir.squaredNorm())*dir;
                // Get the end point of projection
                v += p;
                // Return distance between atom and v
                if(pbc){
                    return sys->box(frame).distance(atom, v);
                } else {
                    return (atom-v).norm();
                }
            };
        } else {
            // For plane
            res = [this,p,dir,pbc](int at){
                Eigen::Vector3f atom = sys->traj[frame].coord[at];
                // Get vector from p to current atom
                Eigen::Vector3f v = atom - p;
                // Project v onto dir
                v = (v.dot(dir)/dir.squaredNorm())*dir;
                // Get closest point on a plane to atom
                v = atom-v;
                // Return distance between atom and v
                if(pbc){
                    return sys->box(frame).distance(atom, v);
                } else {
                    return (atom-v).norm();
                }
            };
        }

    }  else {
        throw Pteros_error("Wrong numeric node!");
    }

    return res;
}