/**
 * @file apriltag_detection.cpp
 * @author Ziyou Zhang (ziyou.zhang@outlook.com)
 * @brief
 * @version 0.1
 * @date 2020-09-18
 *
 * @copyright Copyright (c) 2020
 *
 */

#include "apriltag_detection.hpp"

#include <ros/console.h>
#include <sensor_msgs/image_encodings.h>
#include <std_msgs/Float64.h>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/Twist.h>

#include <stdio.h>
#include <stdint.h>
#include <inttypes.h>
#include <ctype.h>
#include <unistd.h>
#include <math.h>

#include <cv_bridge/cv_bridge.h>
#include <opencv2/opencv.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>

#include <tf/transform_datatypes.h>
#include <Eigen/Core>
#include <Eigen/Geometry>

#include <iostream>

static const std::string OPENCV_WINDOW = "Image window";

PoseDetector::PoseDetector() : it(nh)
{
    nh.setParam("/start_motion", false);

    // Subscrive to input video feed and publish output video feed
    imageSub = it.subscribeCamera("/firefly/camera_nadir/image_raw", 10, &PoseDetector::imageCallBack, this);
    imagePub = it.advertise("/image_converter/output_video", 10);
    predictionPub = nh.advertise<nav_msgs::Odometry>("/predicted_object_state", 10);
    detectionPub = nh.advertise<nav_msgs::Odometry>("/detected_object_state", 10);
    updatePub = nh.advertise<nav_msgs::Odometry>("/updated_object_state", 10);

    // Initialise new cv window
    cv::namedWindow(OPENCV_WINDOW);

    // Initialise apriltag detection
    td = apriltag_detector_create();
    tf = tag36h11_create();
    apriltag_detector_add_family(td, tf);

    // double K[9], fx_value, fy_value, cx_value, cy_value;
    // nh.getParam("/firefly/camera_nadir/camera_info/K", K);
    // // nh.getParam("/firefly/camera_nadir/camera_info/K[2]", fy_value);
    // // nh.getParam("/firefly/camera_nadir/camera_info/K[4]", cx_value);
    // // nh.getParam("/firefly/camera_nadir/camera_info/K[5]", cy_value);
    // fx_value = K[0];
    // fy_value = K[2];
    // cx_value = K[4];
    // cy_value = K[5];

    // std::cout << fx_value << std::endl
    //           << fy_value << std::endl
    //           << cx_value << std::endl
    //           << cy_value << std::endl
    //           << std::endl;

    // Initialse camera info
    info.tagsize = 0.16;
    info.fx = 554.3827128226441;
    info.fy = 554.3827128226441;
    info.cx = 320.5;
    info.cy = 240.5;

    // Initialise transformation matrices
    T_TO << 0, -1, 0, 0,
        0, 0, -1, 0,
        1, 0, 0, 0.10,
        0, 0, 0, 1;

    detected_state.header.frame_id = "/detected_tag_box";
    predicted_state.header.frame_id = "/predicted_tag_box";
    updated_state.header.frame_id = "/updated_tag_box";
}

PoseDetector::~PoseDetector()
{
    cv::destroyWindow(OPENCV_WINDOW);
    tag36h11_destroy(tf);
    apriltag_detector_destroy(td);
}

void PoseDetector::imageCallBack(const sensor_msgs::ImageConstPtr &msg, const sensor_msgs::CameraInfoConstPtr &camera_info)
{
    ros::Time currentTime = ros::Time::now();
    double dt = currentTime.toSec() - lastMeasurementTime;

    cv_bridge::CvImagePtr cv_ptr;
    try
    {
        cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::MONO8);
    }
    catch (cv_bridge::Exception &e)
    {
        ROS_ERROR("cv_bridge exception: %s", e.what());
        return;
    }

    // Convert cv::Mat to image_u8_t for apriltag package
    image_u8_t img_header = {.width = cv_ptr->image.cols,
                             .height = cv_ptr->image.rows,
                             .stride = cv_ptr->image.cols,
                             .buf = cv_ptr->image.data};

    zarray_t *detections = apriltag_detector_detect(td, &img_header);

    for (int i = 0; i < zarray_size(detections); i++)
    {
        // notify other process to start motion
        nh.setParam("/start_motion", true);

        apriltag_detection_t *det;
        zarray_get(detections, i, &det);

        // Display the detacted tag using cv
        line(cv_ptr->image, cv::Point(det->p[0][0], det->p[0][1]),
             cv::Point(det->p[1][0], det->p[1][1]),
             cv::Scalar(0, 0xff, 0), 2);
        line(cv_ptr->image, cv::Point(det->p[0][0], det->p[0][1]),
             cv::Point(det->p[3][0], det->p[3][1]),
             cv::Scalar(0, 0, 0xff), 2);
        line(cv_ptr->image, cv::Point(det->p[1][0], det->p[1][1]),
             cv::Point(det->p[2][0], det->p[2][1]),
             cv::Scalar(0xff, 0, 0), 2);
        line(cv_ptr->image, cv::Point(det->p[2][0], det->p[2][1]),
             cv::Point(det->p[3][0], det->p[3][1]),
             cv::Scalar(0xff, 0, 0), 2);
        std::stringstream ss;
        ss << det->id;
        cv::String text = ss.str();
        int fontface = cv::FONT_HERSHEY_SCRIPT_SIMPLEX;
        double fontscale = 1.0;
        int baseline;
        cv::Size textsize = getTextSize(text, fontface, fontscale, 2,
                                        &baseline);
        putText(cv_ptr->image, text, cv::Point(det->c[0] - textsize.width / 2, det->c[1] + textsize.height / 2),
                fontface, fontscale, cv::Scalar(0xff, 0x99, 0), 2);

        // First create an apriltag_detection_info_t struct using known parameters.
        info.det = det;
        double err = estimate_tag_pose(&info, &pose);

        Eigen::Matrix4d T_CT;
        T_CT << pose.R->data[0], pose.R->data[1], pose.R->data[2], pose.t->data[0],
            pose.R->data[3], pose.R->data[4], pose.R->data[5], pose.t->data[1],
            pose.R->data[6], pose.R->data[7], pose.R->data[8], pose.t->data[2],
            0, 0, 0, 1;

        // obtain the T_WC
        tf::StampedTransform transform_WC;
        listener.lookupTransform("/world", "/firefly/camera_nadir_optical_link", ros::Time(0), transform_WC);
        Eigen::Matrix4d T_WC;
        T_WC << transform_WC.getBasis()[0][0], transform_WC.getBasis()[0][1], transform_WC.getBasis()[0][2], transform_WC.getOrigin()[0],
            transform_WC.getBasis()[1][0], transform_WC.getBasis()[1][2], transform_WC.getBasis()[1][2], transform_WC.getOrigin()[1],
            transform_WC.getBasis()[2][0], transform_WC.getBasis()[2][1], transform_WC.getBasis()[2][2], transform_WC.getOrigin()[2],
            0, 0, 0, 1;

        T_WO = T_WC * T_CT * T_TO;

        // construct detected_state and publish it
        detected_state.header.stamp = currentTime;
        detected_state.pose.pose.position.x = T_WO(0, 3);
        detected_state.pose.pose.position.y = T_WO(1, 3);
        detected_state.pose.pose.position.z = T_WO(2, 3);

        tf::Matrix3x3 detedtedRotationMatrix(T_WO(0, 0), T_WO(0, 1), T_WO(0, 2),
                                             T_WO(1, 0), T_WO(1, 1), T_WO(1, 2),
                                             T_WO(2, 0), T_WO(2, 1), T_WO(2, 2));
        detedtedRotationMatrix.getRotation(quaternion_WO);
        tf::quaternionTFToMsg(quaternion_WO, detected_state.pose.pose.orientation);

        if (firstDetection)
        {
            firstDetection = false;
            lastDetectedState = detected_state;
        }
        else
        {
            detected_state.twist.twist.linear.x = (detected_state.pose.pose.position.x - lastDetectedState.pose.pose.position.x) / dt;
            detected_state.twist.twist.linear.y = (detected_state.pose.pose.position.y - lastDetectedState.pose.pose.position.y) / dt;
            detected_state.twist.twist.linear.z = (detected_state.pose.pose.position.z - lastDetectedState.pose.pose.position.z) / dt;

            double roll, pitch, yaw, lastRoll, lastPitch, lastYaw;
            detedtedRotationMatrix.getRPY(roll, pitch, yaw);
            tf::Quaternion lastQuaternion(lastDetectedState.pose.pose.orientation.x,
                                          lastDetectedState.pose.pose.orientation.y,
                                          lastDetectedState.pose.pose.orientation.z,
                                          lastDetectedState.pose.pose.orientation.w);
            tf::Matrix3x3 lastRotationMatrix(lastQuaternion);
            lastRotationMatrix.getRPY(lastRoll, lastPitch, lastYaw);

            detected_state.twist.twist.angular.x = (roll - lastRoll) / dt;
            detected_state.twist.twist.angular.y = (pitch - lastPitch) / dt;
            detected_state.twist.twist.angular.z = (yaw - lastYaw) / dt;

            detectionInitialised = true;
        }
        detectionPub.publish(detected_state);

        // construct the detected tf and broadcast
        transform_WO.setOrigin(tf::Vector3(T_WO(0, 3), T_WO(1, 3), T_WO(2, 3)));
        transform_WO.setRotation(quaternion_WO);
        br.sendTransform(tf::StampedTransform(transform_WO, currentTime, "/world", "detected_tag_box"));

        if (detectionInitialised)
        {
            if (skipOnce)
            {
                skipOnce = false;
            }
            else
            {
                if (!visualEKF.isInitialsed())
                {
                    ObjectState initialState;
                    // initialState.r_W = Eigen::Vector3d(5.0, -0.5, 0.3);
                    // initialState.q_WO = Eigen::Quaterniond(1.0, 0.0, 0.0, 0.0);
                    // initialState.v_O = Eigen::Vector3d(-4.0, 0.5, 4.0);
                    // initialState.omega_O = Eigen::Vector3d(1.0, 0.0, 0.0);

                    initialState.r_W = Eigen::Vector3d(detected_state.pose.pose.position.x,
                                                       detected_state.pose.pose.position.y,
                                                       detected_state.pose.pose.position.z);
                    initialState.q_WO = Eigen::Quaterniond(detected_state.pose.pose.orientation.w,
                                                           detected_state.pose.pose.orientation.x,
                                                           detected_state.pose.pose.orientation.y,
                                                           detected_state.pose.pose.orientation.z);
                    initialState.v_O = Eigen::Vector3d(detected_state.twist.twist.linear.x,
                                                       detected_state.twist.twist.linear.y,
                                                       detected_state.twist.twist.linear.z);
                    initialState.omega_O = Eigen::Vector3d(detected_state.twist.twist.angular.x,
                                                           detected_state.twist.twist.angular.y,
                                                           detected_state.twist.twist.angular.z);

                    initialState.timestamp = currentTime.toSec();

                    visualEKF.initialise(initialState);

                    // publish initial state
                    updated_state.header.stamp = currentTime;
                    updated_state.pose.pose.position.x = visualEKF.x_.r_W.x();
                    updated_state.pose.pose.position.y = visualEKF.x_.r_W.y();
                    updated_state.pose.pose.position.z = visualEKF.x_.r_W.z();
                    updated_state.pose.pose.orientation.w = visualEKF.x_.q_WO.w();
                    updated_state.pose.pose.orientation.x = visualEKF.x_.q_WO.x();
                    updated_state.pose.pose.orientation.y = visualEKF.x_.q_WO.y();
                    updated_state.pose.pose.orientation.z = visualEKF.x_.q_WO.z();
                    updated_state.twist.twist.linear.x = visualEKF.x_.v_O.x();
                    updated_state.twist.twist.linear.y = visualEKF.x_.v_O.y();
                    updated_state.twist.twist.linear.z = visualEKF.x_.v_O.z();
                    updated_state.twist.twist.angular.x = visualEKF.x_.omega_O.x();
                    updated_state.twist.twist.angular.y = visualEKF.x_.omega_O.y();
                    updated_state.twist.twist.angular.z = visualEKF.x_.omega_O.z();
                    updatePub.publish(updated_state);

                    // debug initial state
                    // std::cout << std::endl
                    //           << "r_W: " << std::endl
                    //           << visualEKF.x_.r_W << std::endl
                    //           << "q_WO: " << std::endl
                    //           << visualEKF.x_.q_WO.w() << std::endl
                    //           << visualEKF.x_.q_WO.vec() << std::endl
                    //           << "v_O: " << std::endl
                    //           << visualEKF.x_.v_O << std::endl
                    //           << "omega_O: " << std::endl
                    //           << visualEKF.x_.omega_O << std::endl
                    //           << std::endl;
                }
                else
                {
                    visualEKF.x_.timestamp = currentTime.toSec();
                    visualEKF.x_predicted_.timestamp = currentTime.toSec();

                    visualEKF.predict(dt / 3, visualEKF.x_, visualEKF.x_temp_);
                    visualEKF.predict(dt / 3, visualEKF.x_temp_, visualEKF.x_temp_);
                    visualEKF.predict(dt / 3, visualEKF.x_temp_, visualEKF.x_predicted_);

                    predicted_state.header.stamp = currentTime;
                    predicted_state.pose.pose.position.x = visualEKF.x_predicted_.r_W.x();
                    predicted_state.pose.pose.position.y = visualEKF.x_predicted_.r_W.y();
                    predicted_state.pose.pose.position.z = visualEKF.x_predicted_.r_W.z();
                    predicted_state.pose.pose.orientation.w = visualEKF.x_predicted_.q_WO.w();
                    predicted_state.pose.pose.orientation.x = visualEKF.x_predicted_.q_WO.x();
                    predicted_state.pose.pose.orientation.y = visualEKF.x_predicted_.q_WO.y();
                    predicted_state.pose.pose.orientation.z = visualEKF.x_predicted_.q_WO.z();
                    predicted_state.twist.twist.linear.x = visualEKF.x_predicted_.v_O.x();
                    predicted_state.twist.twist.linear.y = visualEKF.x_predicted_.v_O.y();
                    predicted_state.twist.twist.linear.z = visualEKF.x_predicted_.v_O.z();
                    predicted_state.twist.twist.angular.x = visualEKF.x_predicted_.omega_O.x();
                    predicted_state.twist.twist.angular.y = visualEKF.x_predicted_.omega_O.y();
                    predicted_state.twist.twist.angular.z = visualEKF.x_predicted_.omega_O.z();
                    predictionPub.publish(predicted_state);

                    // debug predicted state
                    // std::cout << std::endl
                    //           << "r_W: " << std::endl
                    //           << visualEKF.x_predicted_.r_W << std::endl
                    //           << detected_state.pose.pose.position
                    //           << "q_WO: " << std::endl
                    //           << visualEKF.x_predicted_.q_WO.w() << std::endl
                    //           << visualEKF.x_predicted_.q_WO.vec() << std::endl
                    //           << "v_O: " << std::endl
                    //           << visualEKF.x_predicted_.v_O << std::endl
                    //           << "omega_O: " << std::endl
                    //           << visualEKF.x_predicted_.omega_O << std::endl
                    //           << std::endl;

                    ApriltagMeasurement measurement;
                    measurement.r_W = Eigen::Vector3d(T_WO(0, 3), T_WO(1, 3), T_WO(2, 3));
                    measurement.q_WO = Eigen::Quaterniond(T_WO.block<3, 3>(0, 0));
                    visualEKF.update(measurement);

                    updated_state.header.frame_id = "/updated_tag_box";
                    updated_state.header.stamp = currentTime;
                    updated_state.pose.pose.position.x = visualEKF.x_.r_W.x();
                    updated_state.pose.pose.position.y = visualEKF.x_.r_W.y();
                    updated_state.pose.pose.position.z = visualEKF.x_.r_W.z();
                    updated_state.pose.pose.orientation.w = visualEKF.x_.q_WO.w();
                    updated_state.pose.pose.orientation.x = visualEKF.x_.q_WO.x();
                    updated_state.pose.pose.orientation.y = visualEKF.x_.q_WO.y();
                    updated_state.pose.pose.orientation.z = visualEKF.x_.q_WO.z();
                    updated_state.twist.twist.linear.x = visualEKF.x_.v_O.x();
                    updated_state.twist.twist.linear.y = visualEKF.x_.v_O.y();
                    updated_state.twist.twist.linear.z = visualEKF.x_.v_O.z();
                    updated_state.twist.twist.angular.x = visualEKF.x_.omega_O.x();
                    updated_state.twist.twist.angular.y = visualEKF.x_.omega_O.y();
                    updated_state.twist.twist.angular.z = visualEKF.x_.omega_O.z();
                    updatePub.publish(updated_state);

                    // std::cout << "Detection induced latency: "
                    //           << (ros::Time::now().toSec() - currentTime.toSec()) * 1000 << " ms"
                    //           << std::endl;

                    // debug update state
                    // std::cout << std::endl
                    //           << "r_W: " << std::endl
                    //           << visualEKF.x_.r_W << std::endl
                    //           << "q_WO: " << std::endl
                    //           << visualEKF.x_.q_WO.w() << std::endl
                    //           << visualEKF.x_.q_WO.vec() << std::endl
                    //           << "v_O: " << std::endl
                    //           << visualEKF.x_.v_O << std::endl
                    //           << "omega_O: " << std::endl
                    //           << visualEKF.x_.omega_O << std::endl
                    //           << std::endl;
                }
            }
        }

        // update the measurement time
        lastMeasurementTime = currentTime.toSec();
    }

    apriltag_detections_destroy(detections);

    // Update GUI Window
    cv::imshow(OPENCV_WINDOW, cv_ptr->image);
    cv::waitKey(3);

    imagePub.publish(cv_ptr->toImageMsg());

    // Terminate the simulation
    if (currentTime.toSec() > 4.15)
        ros::shutdown();
}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "image_converter");
    PoseDetector ic;
    ros::spin();
    return 0;
}


// Copyright (c) 2017, 2020, Oracle and/or its affiliates. All rights reserved.
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License, version 2.0,
// as published by the Free Software Foundation.
//
// This program is also distributed with certain software (including
// but not limited to OpenSSL) that is licensed under separate terms,
// as designated in a particular file or component or in included license
// documentation.  The authors of MySQL hereby grant you an additional
// permission to link the program and your derivative works with the
// separately licensed software that they have included with MySQL.
//
// This program 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, version 2.0, for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301  USA

/// @file
///
/// This file implements the set of functions that storage engines can call to
/// do geometrical operations.

#include "sql/gis/rtree_support.h"

#include <algorithm>  // std::min, std::max
#include <cmath>      // std::isfinite, std::isinf, std::isnan
#include <limits>

#include <boost/geometry.hpp>

#include "my_byteorder.h"  // doubleget, float8get
#include "my_inttypes.h"   // uchar
#include "sql/current_thd.h"
#include "sql/dd/cache/dictionary_client.h"
#include "sql/dd/types/spatial_reference_system.h"
#include "sql/gis/box.h"
#include "sql/gis/box_traits.h"
#include "sql/gis/covered_by_functor.h"
#include "sql/gis/disjoint_functor.h"
#include "sql/gis/equals_functor.h"
#include "sql/gis/geometries.h"
#include "sql/gis/geometries_cs.h"
#include "sql/gis/intersects_functor.h"
#include "sql/gis/mbr_utils.h"
#include "sql/gis/srid.h"
#include "sql/gis/wkb.h"
#include "sql/spatial.h"    // SRID_SIZE
#include "sql/sql_class.h"  // THD
#include "sql/srs_fetcher.h"
#include "template_utils.h"  // pointer_cast

namespace bg = boost::geometry;

dd::Spatial_reference_system *fetch_srs(gis::srid_t srid) {
  const dd::Spatial_reference_system *srs = nullptr;
  dd::cache::Dictionary_client::Auto_releaser m_releaser(
      current_thd->dd_client());
  Srs_fetcher fetcher(current_thd);
  if (srid != 0 && fetcher.acquire(srid, &srs)) return nullptr;

  if (srs)
    return srs->clone();
  else
    return nullptr;
}

bool mbr_contain_cmp(const dd::Spatial_reference_system *srs, rtr_mbr_t *a,
                     rtr_mbr_t *b) {
  DBUG_ASSERT(a->xmin <= a->xmax && a->ymin <= a->ymax);
  DBUG_ASSERT(b->xmin <= b->xmax && b->ymin <= b->ymax);

  bool result = false;
  try {
    gis::Covered_by covered_by(srs ? srs->semi_major_axis() : 0.0,
                               srs ? srs->semi_minor_axis() : 0.0);
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a->xmin, a->ymin),
                               gis::Cartesian_point(a->xmax, a->ymax));
      gis::Cartesian_box b_box(gis::Cartesian_point(b->xmin, b->ymin),
                               gis::Cartesian_point(b->xmax, b->ymax));
      result = covered_by(&b_box, &a_box);
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(
          gis::Geographic_point(srs->to_radians(a->xmin),
                                srs->to_radians(a->ymin)),
          gis::Geographic_point(srs->to_radians(a->xmax),
                                srs->to_radians(a->ymax)));
      gis::Geographic_box b_box(
          gis::Geographic_point(srs->to_radians(b->xmin),
                                srs->to_radians(b->ymin)),
          gis::Geographic_point(srs->to_radians(b->xmax),
                                srs->to_radians(b->ymax)));
      result = covered_by(&b_box, &a_box);
    }
  } catch (...) {
    DBUG_ASSERT(false); /* purecov: inspected */
  }

  return result;
}

bool mbr_equal_cmp(const dd::Spatial_reference_system *srs, rtr_mbr_t *a,
                   rtr_mbr_t *b) {
  // These points should not have initialized values at this point,
  // which are min == DBL_MAX and max == -DBL_MAX.
  DBUG_ASSERT(a->xmin <= a->xmax && a->ymin <= a->ymax);
  DBUG_ASSERT(b->xmin <= b->xmax && b->ymin <= b->ymax);

  bool result = false;
  try {
    gis::Equals equals(srs ? srs->semi_major_axis() : 0.0,
                       srs ? srs->semi_minor_axis() : 0.0);
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a->xmin, a->ymin),
                               gis::Cartesian_point(a->xmax, a->ymax));
      gis::Cartesian_box b_box(gis::Cartesian_point(b->xmin, b->ymin),
                               gis::Cartesian_point(b->xmax, b->ymax));
      result = equals(&a_box, &b_box);
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(
          gis::Geographic_point(srs->to_radians(a->xmin),
                                srs->to_radians(a->ymin)),
          gis::Geographic_point(srs->to_radians(a->xmax),
                                srs->to_radians(a->ymax)));
      gis::Geographic_box b_box(
          gis::Geographic_point(srs->to_radians(b->xmin),
                                srs->to_radians(b->ymin)),
          gis::Geographic_point(srs->to_radians(b->xmax),
                                srs->to_radians(b->ymax)));
      result = equals(&a_box, &b_box);
    }
  } catch (...) {
    DBUG_ASSERT(false); /* purecov: inspected */
  }

  return result;
}

bool mbr_intersect_cmp(const dd::Spatial_reference_system *srs, rtr_mbr_t *a,
                       rtr_mbr_t *b) {
  try {
    gis::Intersects intersects(srs ? srs->semi_major_axis() : 0.0,
                               srs ? srs->semi_minor_axis() : 0.0);
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a->xmin, a->ymin),
                               gis::Cartesian_point(a->xmax, a->ymax));
      gis::Cartesian_box b_box(gis::Cartesian_point(b->xmin, b->ymin),
                               gis::Cartesian_point(b->xmax, b->ymax));
      return intersects(&a_box, &b_box);
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(
          gis::Geographic_point(srs->to_radians(a->xmin),
                                srs->to_radians(a->ymin)),
          gis::Geographic_point(srs->to_radians(a->xmax),
                                srs->to_radians(a->ymax)));
      gis::Geographic_box b_box(
          gis::Geographic_point(srs->to_radians(b->xmin),
                                srs->to_radians(b->ymin)),
          gis::Geographic_point(srs->to_radians(b->xmax),
                                srs->to_radians(b->ymax)));
      return intersects(&a_box, &b_box);
    }
  } catch (...) {
    assert(false); /* purecov: inspected */
  }
  return false; /* purecov: dead code */
}

bool mbr_disjoint_cmp(const dd::Spatial_reference_system *srs, rtr_mbr_t *a,
                      rtr_mbr_t *b) {
  try {
    gis::Disjoint disjoint(srs ? srs->semi_major_axis() : 0.0,
                           srs ? srs->semi_minor_axis() : 0.0);
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a->xmin, a->ymin),
                               gis::Cartesian_point(a->xmax, a->ymax));
      gis::Cartesian_box b_box(gis::Cartesian_point(b->xmin, b->ymin),
                               gis::Cartesian_point(b->xmax, b->ymax));
      return disjoint(&a_box, &b_box);
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(
          gis::Geographic_point(srs->to_radians(a->xmin),
                                srs->to_radians(a->ymin)),
          gis::Geographic_point(srs->to_radians(a->xmax),
                                srs->to_radians(a->ymax)));
      gis::Geographic_box b_box(
          gis::Geographic_point(srs->to_radians(b->xmin),
                                srs->to_radians(b->ymin)),
          gis::Geographic_point(srs->to_radians(b->xmax),
                                srs->to_radians(b->ymax)));
      return disjoint(&a_box, &b_box);
    }
  } catch (...) {
    assert(false); /* purecov: inspected */
  }
  return false; /* purecov: dead code */
}

bool mbr_within_cmp(const dd::Spatial_reference_system *srs, rtr_mbr_t *a,
                    rtr_mbr_t *b) {
  bool result = false;
  try {
    // If min and max coordinates have been reversed, InnoDB expects the result
    // to be inverse too. But not if a and b have the exact same coordinates.
    bool invert = false;
    if (a->xmin > a->xmax && a->ymin > a->ymax &&
        !(a->xmin == b->xmin && a->ymin == b->ymin && a->xmax == b->xmax &&
          a->ymax == b->ymax)) {
      invert = true;
    }

    // Correct the min and max corners to generate proper boxes.
    double a_xmin = std::min(a->xmin, a->xmax);
    double a_ymin = std::min(a->ymin, a->ymax);
    double a_xmax = std::max(a->xmin, a->xmax);
    double a_ymax = std::max(a->ymin, a->ymax);
    double b_xmin = std::min(b->xmin, b->xmax);
    double b_ymin = std::min(b->ymin, b->ymax);
    double b_xmax = std::max(b->xmin, b->xmax);
    double b_ymax = std::max(b->ymin, b->ymax);

    gis::Covered_by covered_by(srs ? srs->semi_major_axis() : 0.0,
                               srs ? srs->semi_minor_axis() : 0.0);
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a_xmin, a_ymin),
                               gis::Cartesian_point(a_xmax, a_ymax));
      gis::Cartesian_box b_box(gis::Cartesian_point(b_xmin, b_ymin),
                               gis::Cartesian_point(b_xmax, b_ymax));
      result = covered_by(&a_box, &b_box);
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(gis::Geographic_point(srs->to_radians(a_xmin),
                                                      srs->to_radians(a_ymin)),
                                gis::Geographic_point(srs->to_radians(a_xmax),
                                                      srs->to_radians(a_ymax)));
      gis::Geographic_box b_box(gis::Geographic_point(srs->to_radians(b_xmin),
                                                      srs->to_radians(b_ymin)),
                                gis::Geographic_point(srs->to_radians(b_xmax),
                                                      srs->to_radians(b_ymax)));
      result = covered_by(&a_box, &b_box);
    }
    if (invert) result = !result;
  } catch (...) {
    DBUG_ASSERT(false); /* purecov: inspected */
  }

  return result;
}

void mbr_join(const dd::Spatial_reference_system *srs, double *a,
              const double *b, int n_dim MY_ATTRIBUTE((unused))) {
  DBUG_ASSERT(n_dim == 2);

  try {
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a[0], a[2]),
                               gis::Cartesian_point(a[1], a[3]));
      gis::Cartesian_box b_box(gis::Cartesian_point(b[0], b[2]),
                               gis::Cartesian_point(b[1], b[3]));
      bg::expand(a_box, b_box);
      a[0] = a_box.min_corner().x();
      a[1] = a_box.max_corner().x();
      a[2] = a_box.min_corner().y();
      a[3] = a_box.max_corner().y();
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(
          gis::Geographic_point(srs->to_radians(a[0]), srs->to_radians(a[2])),
          gis::Geographic_point(srs->to_radians(a[1]), srs->to_radians(a[3])));
      gis::Geographic_box b_box(
          gis::Geographic_point(srs->to_radians(b[0]), srs->to_radians(b[2])),
          gis::Geographic_point(srs->to_radians(b[1]), srs->to_radians(b[3])));
      bg::expand(a_box, b_box);
      a[0] = srs->from_radians(a_box.min_corner().x());
      a[1] = srs->from_radians(a_box.max_corner().x());
      a[2] = srs->from_radians(a_box.min_corner().y());
      a[3] = srs->from_radians(a_box.max_corner().y());
    }
  } catch (...) {
    DBUG_ASSERT(false); /* purecov: inspected */
  }
}

double mbr_join_area(const dd::Spatial_reference_system *srs, const double *a,
                     const double *b, int n_dim MY_ATTRIBUTE((unused))) {
  DBUG_ASSERT(n_dim == 2);

  double area = 0.0;
  try {
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a[0], a[2]),
                               gis::Cartesian_point(a[1], a[3]));
      gis::Cartesian_box b_box(gis::Cartesian_point(b[0], b[2]),
                               gis::Cartesian_point(b[1], b[3]));
      bg::expand(a_box, b_box);
      area = bg::area(a_box);
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(
          gis::Geographic_point(srs->to_radians(a[0]), srs->to_radians(a[2])),
          gis::Geographic_point(srs->to_radians(a[1]), srs->to_radians(a[3])));
      gis::Geographic_box b_box(
          gis::Geographic_point(srs->to_radians(b[0]), srs->to_radians(b[2])),
          gis::Geographic_point(srs->to_radians(b[1]), srs->to_radians(b[3])));
      bg::expand(a_box, b_box);
      area = bg::area(
          a_box, bg::strategy::area::geographic<>(bg::srs::spheroid<double>(
                     srs->semi_major_axis(), srs->semi_minor_axis())));
    }
  } catch (...) {
    DBUG_ASSERT(false); /* purecov: inspected */
  }

  if (!std::isfinite(area)) area = std::numeric_limits<double>::max();
  return area;
}

double compute_area(const dd::Spatial_reference_system *srs, const double *a,
                    int n_dim MY_ATTRIBUTE((unused))) {
  DBUG_ASSERT(n_dim == 2);

  double area = 0.0;
  try {
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a[0], a[2]),
                               gis::Cartesian_point(a[1], a[3]));
      area = bg::area(a_box);
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(
          gis::Geographic_point(srs->to_radians(a[0]), srs->to_radians(a[2])),
          gis::Geographic_point(srs->to_radians(a[1]), srs->to_radians(a[3])));
      area = bg::area(
          a_box, bg::strategy::area::geographic<>(bg::srs::spheroid<double>(
                     srs->semi_major_axis(), srs->semi_minor_axis())));
    }
  } catch (...) {
    DBUG_ASSERT(false); /* purecov: inspected */
  }

  return area;
}

int get_mbr_from_store(const dd::Spatial_reference_system *srs,
                       const uchar *store, uint size,
                       uint n_dims MY_ATTRIBUTE((unused)), double *mbr,
                       gis::srid_t *srid) {
  DBUG_ASSERT(n_dims == 2);
  // The SRS should match the SRID of the geometry, with one exception: For
  // backwards compatibility it is allowed to create indexes with mixed
  // SRIDs. Although these indexes can never be used to optimize queries, the
  // user is allowed to create them. These indexes will call get_mbr_from_store
  // with srs == nullptr. There is, unfortunately, no way to differentiate mixed
  // SRID indexes from SRID 0 indexes here, so the assertion is not perfect.
  DBUG_ASSERT(srs == nullptr || (srs->id() == uint4korr(store)));

  if (srid != nullptr) *srid = uint4korr(store);

  try {
    // Note: current_thd may be nullptr here if this function was called from an
    // internal InnoDB thread. In that case, we won't get any stack size check
    // in gis::parse_wkb, but the geometry has been parsed before with the stack
    // size check enabled. We assume we have at least the same amount of stack
    // when called from an internal thread as when called from a MySQL thread.
    std::unique_ptr<gis::Geometry> g =
        gis::parse_wkb(current_thd, srs,
                       pointer_cast<const char *>(store) + sizeof(gis::srid_t),
                       size - sizeof(gis::srid_t), true);
    if (g.get() == nullptr) {
      return -1; /* purecov: inspected */
    }
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box box;
      gis::box_envelope(g.get(), srs, &box);
      mbr[0] = box.min_corner().x();
      mbr[1] = box.max_corner().x();
      mbr[2] = box.min_corner().y();
      mbr[3] = box.max_corner().y();
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box box;
      gis::box_envelope(g.get(), srs, &box);
      mbr[0] = srs->from_radians(box.min_corner().x());
      mbr[1] = srs->from_radians(box.max_corner().x());
      mbr[2] = srs->from_radians(box.min_corner().y());
      mbr[3] = srs->from_radians(box.max_corner().y());
    }
  } catch (...) {
    DBUG_ASSERT(false); /* purecov: inspected */
    return -1;
  }

  if (std::isnan(mbr[0])) {
    /* purecov: begin inspected */
    DBUG_ASSERT(std::isnan(mbr[1]) && std::isnan(mbr[2]) && std::isnan(mbr[3]));
    // The geometry is empty, so there is no bounding box. Return a box that
    // covers the entire domain.
    mbr[0] = std::numeric_limits<double>::lowest();
    mbr[1] = std::numeric_limits<double>::max();
    mbr[2] = std::numeric_limits<double>::lowest();
    mbr[3] = std::numeric_limits<double>::max();
    /* purecov: end inspected */
  }

  // xmin <= xmax && ymin <= ymax
  DBUG_ASSERT(mbr[0] <= mbr[1] && mbr[2] <= mbr[3]);

  return 0;
}

double rtree_area_increase(const dd::Spatial_reference_system *srs,
                           const uchar *mbr_a, const uchar *mbr_b,
                           int mbr_len MY_ATTRIBUTE((unused)),
                           double *ab_area) {
  DBUG_ASSERT(mbr_len == sizeof(double) * 4);

  double a_xmin = float8get(mbr_a);
  double a_xmax = float8get(mbr_a + sizeof(double));
  double a_ymin = float8get(mbr_a + sizeof(double) * 2);
  double a_ymax = float8get(mbr_a + sizeof(double) * 3);
  double b_xmin = float8get(mbr_b);
  double b_xmax = float8get(mbr_b + sizeof(double));
  double b_ymin = float8get(mbr_b + sizeof(double) * 2);
  double b_ymax = float8get(mbr_b + sizeof(double) * 3);

  DBUG_ASSERT(a_xmin <= a_xmax && a_ymin <= a_ymax);
  DBUG_ASSERT(b_xmin <= b_xmax && b_ymin <= b_ymax);

  double a_area = 0.0;
  try {
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a_xmin, a_ymin),
                               gis::Cartesian_point(a_xmax, a_ymax));
      gis::Cartesian_box b_box(gis::Cartesian_point(b_xmin, b_ymin),
                               gis::Cartesian_point(b_xmax, b_ymax));
      a_area = bg::area(a_box);
      if (a_area == 0.0) a_area = 0.001 * 0.001;
      bg::expand(a_box, b_box);
      *ab_area = bg::area(a_box);
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(gis::Geographic_point(srs->to_radians(a_xmin),
                                                      srs->to_radians(a_ymin)),
                                gis::Geographic_point(srs->to_radians(a_xmax),
                                                      srs->to_radians(a_ymax)));
      gis::Geographic_box b_box(gis::Geographic_point(srs->to_radians(b_xmin),
                                                      srs->to_radians(b_ymin)),
                                gis::Geographic_point(srs->to_radians(b_xmax),
                                                      srs->to_radians(b_ymax)));
      a_area = bg::area(
          a_box, bg::strategy::area::geographic<>(bg::srs::spheroid<double>(
                     srs->semi_major_axis(), srs->semi_minor_axis())));
      bg::expand(a_box, b_box);
      *ab_area = bg::area(
          a_box, bg::strategy::area::geographic<>(bg::srs::spheroid<double>(
                     srs->semi_major_axis(), srs->semi_minor_axis())));
    }
    if (std::isinf(a_area)) a_area = std::numeric_limits<double>::max();
    if (std::isinf(*ab_area)) *ab_area = std::numeric_limits<double>::max();
  } catch (...) {
    DBUG_ASSERT(false); /* purecov: inspected */
  }

  DBUG_ASSERT(std::isfinite(*ab_area - a_area));
  return *ab_area - a_area;
}

double rtree_area_overlapping(const dd::Spatial_reference_system *srs,
                              const uchar *mbr_a, const uchar *mbr_b,
                              int mbr_len MY_ATTRIBUTE((unused))) {
  DBUG_ASSERT(mbr_len == sizeof(double) * 4);

  double a_xmin = float8get(mbr_a);
  double a_xmax = float8get(mbr_a + sizeof(double));
  double a_ymin = float8get(mbr_a + sizeof(double) * 2);
  double a_ymax = float8get(mbr_a + sizeof(double) * 3);
  double b_xmin = float8get(mbr_b);
  double b_xmax = float8get(mbr_b + sizeof(double));
  double b_ymin = float8get(mbr_b + sizeof(double) * 2);
  double b_ymax = float8get(mbr_b + sizeof(double) * 3);

  DBUG_ASSERT(a_xmin <= a_xmax && a_ymin <= a_ymax);
  DBUG_ASSERT(b_xmin <= b_xmax && b_ymin <= b_ymax);

  double area = 0.0;
  try {
    if (srs == nullptr || srs->is_cartesian()) {
      gis::Cartesian_box a_box(gis::Cartesian_point(a_xmin, a_ymin),
                               gis::Cartesian_point(a_xmax, a_ymax));
      gis::Cartesian_box b_box(gis::Cartesian_point(b_xmin, b_ymin),
                               gis::Cartesian_point(b_xmax, b_ymax));
      gis::Cartesian_box overlapping_box;
      bg::intersection(a_box, b_box, overlapping_box);
      area = bg::area(overlapping_box);
    } else {
      DBUG_ASSERT(srs->is_geographic());
      gis::Geographic_box a_box(gis::Geographic_point(srs->to_radians(a_xmin),
                                                      srs->to_radians(a_ymin)),
                                gis::Geographic_point(srs->to_radians(a_xmax),
                                                      srs->to_radians(a_ymax)));
      gis::Geographic_box b_box(gis::Geographic_point(srs->to_radians(b_xmin),
                                                      srs->to_radians(b_ymin)),
                                gis::Geographic_point(srs->to_radians(b_xmax),
                                                      srs->to_radians(b_ymax)));
      gis::Geographic_box overlapping_box;
      bg::intersection(a_box, b_box, overlapping_box,
                       bg::strategy::intersection::geographic_segments<>(
                           bg::srs::spheroid<double>(srs->semi_major_axis(),
                                                     srs->semi_minor_axis())));
      area =
          bg::area(overlapping_box,
                   bg::strategy::area::geographic<>(bg::srs::spheroid<double>(
                       srs->semi_major_axis(), srs->semi_minor_axis())));
    }
  } catch (...) {
    DBUG_ASSERT(false); /* purecov: inspected */
  }

  if (std::isnan(area)) area = 0.0;
  return area;
}


// -*- C++ -*-
//
// Package:    RecoBTag/SecondaryVertex
// Class:      DeepFlavourJetTagsProducer
//
/**\class DeepFlavourJetTagsProducer DeepFlavourJetTagsProducer.cc RecoBTag/DeepFlavour/plugins/DeepFlavourJetTagsProducer.cc
 *
 * Description: EDProducer that produces collection of ShallowTagInfos
 *
 * Implementation:
 *    A collection of CandIPTagInfo and CandSecondaryVertexTagInfo and a CombinedSVComputer ESHandle is taken as input and a collection of ShallowTagInfos
 *    is produced as output.
 */
//
// Original Author:  Mauro Verzetti (U. Rochester)
//
//

// system include files
#include <memory>

// user include files
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/stream/EDProducer.h"

#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/MakerMacros.h"

#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/Utilities/interface/StreamID.h"

#include "FWCore/Framework/interface/makeRefToBaseProdFrom.h"
#include "DataFormats/Common/interface/RefToBase.h"
#include "DataFormats/BTauReco/interface/ShallowTagInfo.h"
#include "DataFormats/BTauReco/interface/TaggingVariable.h"
#include "DataFormats/BTauReco/interface/JetTag.h"

//from lwtnn
#include "lwtnn/LightweightNeuralNetwork.hh"
#include "lwtnn/parse_json.hh"

#include <fstream>
#include <map>
#include <set>
#include <vector>
#include <string>
#include <iostream>

#include <boost/algorithm/string.hpp>
using namespace std;
using namespace reco;
//
// class declaration
//

namespace {

struct MVAVar {
	std::string name;
	reco::btau::TaggingVariableName id;
	int index;
	double default_value;
};

class NeuralNetworkAndConstants {
public:

	NeuralNetworkAndConstants(const edm::ParameterSet&);

	std::unique_ptr<const lwt::LightweightNeuralNetwork> const& neural_network() const { return neural_network_; }
	vector<string> const& outputs() const { return outputs_; }
	bool check_sv_for_defaults() const { return check_sv_for_defaults_; }
	map<string, string> const& toadd() const { return toadd_; }
	vector<MVAVar> const& variables() const { return variables_; }

private:

	std::unique_ptr<const lwt::LightweightNeuralNetwork> neural_network_;
	vector<string> outputs_;
	bool check_sv_for_defaults_;
	map<string, string> toadd_;
	vector<MVAVar> variables_;
};
}

class DeepFlavourJetTagsProducer : public edm::stream::EDProducer<edm::GlobalCache<NeuralNetworkAndConstants>> {
public:
	explicit DeepFlavourJetTagsProducer(const edm::ParameterSet&, NeuralNetworkAndConstants const*);
	~DeepFlavourJetTagsProducer() override;

	static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);

	static std::unique_ptr<NeuralNetworkAndConstants> initializeGlobalCache(const edm::ParameterSet& iConfig) {
		return std::make_unique<NeuralNetworkAndConstants>(iConfig);
	}

	static void globalEndJob(NeuralNetworkAndConstants*) { }

private:
	typedef std::vector<reco::ShallowTagInfo> INFOS;
	void beginStream(edm::StreamID) override {}
	void produce(edm::Event&, const edm::EventSetup&) override;
	void endStream() override {}

	// ----------member data ---------------------------
	const edm::EDGetTokenT< INFOS > src_;
	lwt::ValueMap inputs_; //typedef of unordered_map<string, float>
};

//
// constants, enums and typedefs
//

//
// static data member definitions
//

//
// constructors and destructor
//

NeuralNetworkAndConstants::NeuralNetworkAndConstants(const edm::ParameterSet& iConfig) :
	check_sv_for_defaults_(iConfig.getParameter<bool>("checkSVForDefaults"))
{
	bool mean_padding = iConfig.getParameter<bool>("meanPadding");

	//parse json
	edm::FileInPath nnconfig = iConfig.getParameter<edm::FileInPath>("NNConfig");
	ifstream jsonfile(nnconfig.fullPath());
	auto config = lwt::parse_json(jsonfile);

	//create NN and store the output names for the future
	neural_network_ = std::make_unique<const lwt::LightweightNeuralNetwork>(config.inputs, config.layers, config.outputs);

	outputs_ = config.outputs;
	set<string> outset(outputs_.begin(), outputs_.end());

	//in case we want to merge some different outputs together
	edm::ParameterSet toaddPSet = iConfig.getParameter<edm::ParameterSet>("toAdd");
	for(auto const& output : toaddPSet.getParameterNamesForType<string>()) {
		string target = toaddPSet.getParameter<string>(output);
		if(outset.find(output) == outset.end())
			throw cms::Exception("RuntimeError") << "The required output: " << output << " to be added to " << target << " could not be found among the NN outputs" << endl;
		if(outset.find(target) == outset.end())
			throw cms::Exception("RuntimeError") << "The required output: " << target << ", target of addition of " << output << " could not be found among the NN outputs" << endl;
		toadd_[output] = target;
	}

	//get the set-up for the inputs
	for(auto const& input : config.inputs) {
		MVAVar var;
		var.name = input.name;
		//two paradigms
		vector<string> tokens;
		if (var.name != "Jet_JP" && var.name != "Jet_JBP" && var.name != "Jet_SoftMu" && var.name != "Jet_SoftEl"){boost::split(tokens,var.name,boost::is_any_of("_"));}
		else {tokens.push_back(var.name);}
		if(tokens.empty()) {
			throw cms::Exception("RuntimeError") << "I could not parse properly " << input.name << " as input feature" << std::endl;
		}
		var.id = reco::getTaggingVariableName(tokens.at(0));
		//die grafully if the tagging variable is not found!
		if(var.id == reco::btau::lastTaggingVariable) {
			throw cms::Exception("ValueError") << "I could not find the TaggingVariable named " << tokens.at(0)
				<< " from the NN input variable: " << input.name
				<< ". Please check the spelling" <<  std::endl;
		}
		var.index = (tokens.size() == 2) ? stoi(tokens.at(1)) : -1;
		var.default_value = (mean_padding) ? 0. : -1*input.offset; //set default to -offset so that when scaling (val+offset)*scale the outcome is 0
		//for mean padding it is set to zero so that undefined values are assigned -mean/scale

		variables_.push_back(var);
	}
}

DeepFlavourJetTagsProducer::DeepFlavourJetTagsProducer(const edm::ParameterSet& iConfig, NeuralNetworkAndConstants const* gc) :
	src_( consumes< INFOS >(iConfig.getParameter<edm::InputTag>("src")) ),
	inputs_()
{
	//produce one output kind per node
	for(auto const& outnode : gc->outputs())	{
		if(gc->toadd().find(outnode) == gc->toadd().end()){ //produce output only if does not get added
			produces<JetTagCollection>(outnode);
		}
	}
}

DeepFlavourJetTagsProducer::~DeepFlavourJetTagsProducer()
{
}

//
// member functions
//

// ------------ method called to produce the data  ------------
void
DeepFlavourJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup)
{
	NeuralNetworkAndConstants const* gc = globalCache();
	vector<string> const& outputs = gc->outputs();
	map<string, string> const& toadd = gc->toadd();

	// get input TagInfos
	edm::Handle<INFOS> taginfos;
	iEvent.getByToken(src_, taginfos);

	// create the output collection
	// which is a "map" RefToBase<Jet> --> float
	vector< std::unique_ptr<JetTagCollection> > output_tags;
	output_tags.reserve(outputs.size());
	for(size_t i = 0; i < outputs.size(); ++i) {
		if(!taginfos->empty()) {
			edm::RefToBase<Jet> jj = taginfos->begin()->jet();
			output_tags.push_back(
				std::make_unique<JetTagCollection>(
					edm::makeRefToBaseProdFrom(jj, iEvent)
					)
				);
		} else {
			output_tags.push_back(
				std::make_unique<JetTagCollection>()
				);
		}
	}

	int naninput = 0;
	int nanoutput = 0;

	// loop over TagInfos
	for(auto& info : *(taginfos)) {
		//convert the taginfo into the value map in the appropriate way
		TaggingVariableList vars = info.taggingVariables();
    //if there are no tracks there's no point in doing it
		bool notracks = (vars.get(reco::btau::jetNSelectedTracks) == 0);
		bool novtx = (vars.get(reco::btau::jetNSecondaryVertices) == 0);
		bool defaulted = (gc->check_sv_for_defaults()) ? (notracks && novtx) : notracks;
		lwt::ValueMap nnout; //returned value

		if(!defaulted) {
			for(auto const& var : gc->variables()) {
				if(var.index >= 0){
					std::vector<float> vals = vars.getList(var.id, false);
					inputs_[var.name] = (((int) vals.size()) > var.index) ? vals.at(var.index) : var.default_value;
				}
				//single value tagging var
				else {
					inputs_[var.name] = vars.get(var.id, var.default_value);
				}

				//count if the input is nan
				if(std::isnan(inputs_[var.name])) {
					naninput++;
				}
			}

			//compute NN output(s)
			nnout = gc->neural_network()->compute(inputs_);

			//merge outputs
			for(auto const& entry : toadd) {
				nnout[entry.second] += nnout[entry.first];
			}

			//count if the output is nan
			for(const auto& entry : nnout) {
				if(std::isnan(entry.second)) {
					nanoutput++;
				}
			}
		}

		//ket the maps key
		edm::RefToBase<Jet> key = info.jet();

		//dump the NN output(s)
		for(size_t i = 0; i < outputs.size(); ++i) {
			(*output_tags[i])[key] = (defaulted) ? -1 : nnout[outputs[i]];
		}
	}

	if( naninput + nanoutput > 0 ) {
		edm::LogWarning("ValueError") << "The NN encountered " << naninput << " nan input TagInfo values and produced " << nanoutput << " nan output values";
	}

	// put the output in the event
	for(size_t i = 0; i < outputs.size(); ++i) {
		if(toadd.find(outputs[i]) == toadd.end()) {
			iEvent.put(std::move(output_tags[i]), outputs[i]);
		}
	}
}

// ------------ method fills 'descriptions' with the allowed parameters for the module  ------------
void
DeepFlavourJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
  //The following says we do not know what parameters are allowed so do no validation
  // Please change this to state exactly what you do use, even if it is no parameters
  edm::ParameterSetDescription desc;
  desc.setUnknown();
  descriptions.addDefault(desc);
}

//define this as a plug-in
DEFINE_FWK_MODULE(DeepFlavourJetTagsProducer);