#include "combi_darknet_openface_node.hpp"

#include <limits>
#include <string>

#include <Eigen/Dense>
#include <Eigen/Geometry>

#include <tf/transform_datatypes.h>
#include <std_msgs/Int16.h>
#include <combi_darknet_openface/GazedObjectDetectionResult.h>
#include <combi_darknet_openface/DetectionResult.h>

#include "math_util.h"

CombiDarknetOpenface::CombiDarknetOpenface()
    : nh1("~"), class_names({"none"})
{
    ros_object_sub = nh1.subscribe("darknet_ros/bounding_boxes",1, &CombiDarknetOpenface::onRecognizedObject, this);
    ros_person_pos_sub= nh1.subscribe("estimate_pos",1, &CombiDarknetOpenface::onPersonPositionEstimated, this);
    ros_face_sub = nh1.subscribe("faces",1, &CombiDarknetOpenface::onRecognizedFace, this);
    rgb_img_sub = nh1.subscribe<sensor_msgs::Image>("/camera/rgb/image_rect_color", 1, &CombiDarknetOpenface::onRgbImageUpdated, this);
    depth_img_sub = nh1.subscribe<sensor_msgs::Image>("/camera/depth_registered/image_raw", 1, &CombiDarknetOpenface::onDepthImageUpdated, this);
    camera_info_sub = nh1.subscribe<sensor_msgs::CameraInfo>("/camera/color/camera_info", 1, &CombiDarknetOpenface::onCameraInfoUpdated, this);
    measurement_pub = nh1.advertise<geometry_msgs::PoseStamped>("filter_measurement", 1);
    headpose_arrow_pub = nh1.advertise<visualization_msgs::Marker>("/visualization_headpose_arrow", 1);
    origin_marker_pub = nh1.advertise<visualization_msgs::Marker>("/visualization_origin_marker", 1);
    person_arrow_pub = nh1.advertise<visualization_msgs::Marker>("/visualization_person_arrow", 1);
    person_marker_pub = nh1.advertise<visualization_msgs::Marker>("/visualization_person_marker", 1);
    object_marker_pub = nh1.advertise<visualization_msgs::Marker>("/visualization_object_marker", 1);
    estimate_marker_pub = nh1.advertise<visualization_msgs::Marker>("/visualization_estimateperson_marker", 1);
    gazed_object_detection_pub = nh1.advertise<combi_darknet_openface::GazedObjectDetectionResult>("gazed_object_detection_result", 1);
    cv::namedWindow("RGB image", CV_WINDOW_AUTOSIZE);
}

CombiDarknetOpenface::~CombiDarknetOpenface()
{
}

void CombiDarknetOpenface::onCameraInfoUpdated(const sensor_msgs::CameraInfo::ConstPtr& msg)
{
    this->camera_matrix = cv::Mat(3, 3, CV_64FC1);
    for(std::size_t i = 0; i < msg->K.size(); ++i)
    {
        this->camera_matrix.at<double>(i / 3, i % 3) = msg->K[i];
    }
    this->dist_coeffs = cv::Mat(msg->D.size(), 1, CV_64FC1);
    for(std::size_t i = 0; i < msg->D.size(); ++i)
    {
        this->dist_coeffs.at<double>(i, 0) = msg->D[i];
    }
    this->image_size.width = msg->width;
    this->image_size.height = msg->height;
}

void CombiDarknetOpenface::onRecognizedFace(const openface_ros::Faces::ConstPtr& msg)
{
    this->nose_tip_position_ptr.reset();
    this->chin_position_ptr.reset();
    this->left_eye_position_ptr.reset();
    this->right_eye_position_ptr.reset();
    this->left_mouth_position_ptr.reset();
    this->right_mouth_position_ptr.reset();
    this->nose_end_point2D_drawtmp.reset();

    for(const auto& detected_face : msg->faces)
    {
        if(detected_face.landmarks_2d.empty())
        {
            continue;
        }

        this->nose_tip_position_ptr.reset(new cv::Point2i(std::round(detected_face.landmarks_2d[33].x), std::round(detected_face.landmarks_2d[33].y)));
        this->chin_position_ptr.reset(new cv::Point2i(std::round(detected_face.landmarks_2d[8].x), std::round(detected_face.landmarks_2d[8].y)));
        this->left_eye_position_ptr.reset(new cv::Point2i(std::round(detected_face.landmarks_2d[36].x), std::round(detected_face.landmarks_2d[36].y)));
        this->right_eye_position_ptr.reset(new cv::Point2i(std::round(detected_face.landmarks_2d[45].x), std::round(detected_face.landmarks_2d[45].y)));
        this->left_mouth_position_ptr.reset(new cv::Point2i(std::round(detected_face.landmarks_2d[48].x), std::round(detected_face.landmarks_2d[48].y)));
        this->right_mouth_position_ptr.reset(new cv::Point2i(std::round(detected_face.landmarks_2d[54].x), std::round(detected_face.landmarks_2d[54].y)));

        this->updateExtrinsicParameters(
            *this->nose_tip_position_ptr,
            *this->chin_position_ptr,
            *this->left_eye_position_ptr,
            *this->right_eye_position_ptr,
            *this->left_mouth_position_ptr,
            *this->right_mouth_position_ptr
        );

        nose_end_point2D_drawtmp.reset(new cv::Point2i(this->getProjectedPoint(cv::Point3f(0.0, 0.0, this->TMP_DISTANCE))));

        cv::Mat rotation_matrix  = cv::Mat::zeros(3, 3, CV_64FC1); // Rotation in axis-angle form
        cv::Rodrigues(this->rotation_vector, rotation_matrix);
        cv::Mat proj_matrix = (cv::Mat_<double>(3,4) <<
            rotation_matrix.at<double>(0, 0), rotation_matrix.at<double>(0, 1), rotation_matrix.at<double>(0, 2), 0,
            rotation_matrix.at<double>(1, 0), rotation_matrix.at<double>(1, 1), rotation_matrix.at<double>(1, 2), 0,
            rotation_matrix.at<double>(2, 0), rotation_matrix.at<double>(2, 1), rotation_matrix.at<double>(2, 2), 0
        );
        cv::Mat head_translation_vector = this->translation_vector.clone();
        cv::Vec3d euler_angles;
        cv::decomposeProjectionMatrix(
            proj_matrix,
            this->camera_matrix,
            rotation_matrix,
            head_translation_vector,
            cv::noArray(),
            cv::noArray(),
            cv::noArray(),
            euler_angles
        );
        EulerAngles head_orientation = {
            euler_angles[2],  // roll
            euler_angles[0],  // pitch
            euler_angles[1]  // yaw
        };
        this->applyCacheIf(
            head_orientation,
            this->rotation_vector,
            this->translation_vector,
            !this->isInImageArea(*this->nose_end_point2D_drawtmp)
        );

        if(this->estimate_position_ptr && (!head_orientation.empty()))
        {
            double head_arrow_theta = calcHeadArrowAngle(head_orientation);
            this->publishHeadPoseArrow(*this->estimate_position_ptr, head_arrow_theta);
        }
    }
}

bool CombiDarknetOpenface::isInImageArea(const cv::Point2i& point) const
{
    return (point.x > 0) && (point.x < this->image_size.width) && (point.y > 0) && (point.y < this->image_size.height);
}

void CombiDarknetOpenface::updateExtrinsicParameters(const cv::Point2i& nose_tip_position, const cv::Point2i& chin_position, const cv::Point2i& left_eye_position, const cv::Point2i& right_eye_position, const cv::Point2i& left_mouth_position, const cv::Point2i& right_mouth_position)
{
    std::vector<cv::Point2f> image_points = {
        nose_tip_position,
        chin_position,
        left_eye_position,
        right_eye_position,
        left_mouth_position,
        right_mouth_position
    };

    const float face_model_scale = 0.225;
    std::vector<cv::Point3f> model_points = {
        cv::Point3f(0.0f, 0.0f, 0.0f),  // Nose tip
        cv::Point3f(0.0f, -330.0*face_model_scale, -65.0*face_model_scale),  // Chin
        cv::Point3f(-225.0*face_model_scale, 170.0*face_model_scale, -135.0*face_model_scale),  // Left eye left corner
        cv::Point3f(225.0*face_model_scale, 170.0*face_model_scale, -135.0*face_model_scale),  // Right eye right corner
        cv::Point3f(-150.0*face_model_scale, -150.0*face_model_scale, -125.0*face_model_scale),  // Left Mouth corner
        cv::Point3f(150.0*face_model_scale, -150.0*face_model_scale, -125.0*face_model_scale)  // Right mouth corner
    };

    this->updateExtrinsicParameters(model_points, image_points);
}

void CombiDarknetOpenface::updateExtrinsicParameters(const std::vector<cv::Point3f>& model_points, const std::vector<cv::Point2f>& image_points)
{
    this->rotation_vector  = cv::Mat::zeros(3, 1, CV_64FC1); // Rotation in axis-angle form
    this->translation_vector = cv::Mat::zeros(3, 1, CV_64FC1);

    cv::solvePnP(model_points, image_points, this->camera_matrix, this->dist_coeffs, this->rotation_vector, this->translation_vector, CV_ITERATIVE);
}

void CombiDarknetOpenface::applyCacheIf(EulerAngles& head_orientation, cv::Mat& rotation_vector, cv::Mat& translation_vector, bool condition) const
{
    static bool initialized = false;
    static EulerAngles previous_head_orientation;
    static cv::Mat previous_rotation_vector;
    static cv::Mat previous_translation_vector;

    if(!initialized)
    {
        previous_head_orientation =  head_orientation;
        previous_rotation_vector = rotation_vector.clone();
        previous_translation_vector = translation_vector.clone();
        initialized = true;
    }

    if(condition)
    {
        head_orientation = previous_head_orientation;
        rotation_vector = previous_rotation_vector.clone();
        translation_vector = previous_translation_vector.clone();
    }
    else
    {
        previous_head_orientation = head_orientation;
        previous_rotation_vector = rotation_vector.clone();
        previous_translation_vector = translation_vector.clone();
    }
}

cv::Point2i CombiDarknetOpenface::getProjectedPoint(const cv::Point3f& point_3d) const
{
    std::vector<cv::Point3f> point_3d_list = {point_3d};
    std::vector<cv::Point2f> point_2d_list;
    cv::projectPoints(point_3d_list, this->rotation_vector, this->translation_vector, this->camera_matrix, this->dist_coeffs, point_2d_list);
    return cv::Point2i(std::round(point_2d_list[0].x), std::round(point_2d_list[0].y));
}

double CombiDarknetOpenface::calcHeadArrowAngle(const EulerAngles& head_orientation) const
{
    double ret_val = 0.0;
    if(0 <= head_orientation[2] && head_orientation[2] < 90)
    {
        ret_val = 180 - head_orientation[2];
    }
    else if(-90 <= head_orientation[2] && head_orientation[2]<0)
    {
        ret_val = std::abs(head_orientation[2]) + 180;
    }
    else
    {
        ret_val = head_orientation[2];
    }

    ret_val += this->robotyaw;
    if(ret_val>=360)
    {
        ret_val-=360;
    }
    return ret_val;
}

void CombiDarknetOpenface::onRecognizedObject(const darknet_ros_msgs::BoundingBoxes::ConstPtr& msg )
{
    std::vector<std::string> current_frame_class_names = {"none"};
    std::vector<cv::Rect> current_frame_object_boxes = {cv::Rect(0, 0, 0, 0)};

    this->nearest_object_index = 0;
    this->person_box.reset();

    for(const auto& detected_object :  msg->boundingBoxes)
    {
        current_frame_class_names.push_back(detected_object.Class);
        current_frame_object_boxes.push_back(cv::Rect(cv::Point(detected_object.xmin, detected_object.ymin), cv::Point(detected_object.xmax, detected_object.ymax)));
        if(detected_object.Class=="person")
        {
            this->person_box.reset(new cv::Rect(cv::Point(detected_object.xmin, detected_object.ymin), cv::Point(detected_object.xmax, detected_object.ymax)));
        }
        if(std::find(this->class_names.begin(),this->class_names.end(),detected_object.Class) == this->class_names.end())
        {
            this->class_names.push_back(detected_object.Class);
        }
    }
    this->object_boxes.resize(this->class_names.size());
    std::fill(this->object_boxes.begin(), this->object_boxes.end(), cv::Rect(0, 0, 0, 0));
    for(std::size_t i = 0; i < current_frame_class_names.size(); ++i)
    {
        auto box_index = std::distance(
            this->class_names.begin(),
            std::find(this->class_names.begin(), this->class_names.end(), current_frame_class_names[i])
        );
        this->object_boxes[box_index] = current_frame_object_boxes[i];
    }
    this->object_centers.resize(this->object_boxes.size());
    for(std::size_t i = 0; i < this->class_names.size(); ++i)
    {
        if(this->isIgnoredObjectClass(this->class_names[i]))
        {
            this->object_centers[i] = this->invalidPoint();
        }
        else
        {
            this->object_centers[i] = (this->object_boxes[i].tl() + this->object_boxes[i].br()) / 2;
        }
    }
    if(this->person_box && this->person_moving_state == PersonMovingState::Stopping)
    {
        this->calculateTimeUse();
    }
    else{
        this->calculateTimeUseOutofView();
    }
    this->publishGazedObjectDetectionResult(msg->header);
}

bool CombiDarknetOpenface::isIgnoredObjectClass(const std::string& class_name) const
{
    return std::find(this->ignore_object_list.begin(), this->ignore_object_list.end(), class_name) != this->ignore_object_list.end();
}

void CombiDarknetOpenface::calculateTimeUse()
{
    this->nearest_object_index = 0;
    this->nearest_gaze_position_ptr.reset();
    if(this->nose_tip_position_ptr && this->nose_end_point2D_drawtmp)
    {
        double nearest_object_distance;
        this->nearest_gaze_position_ptr.reset(new cv::Point2i(this->invalidPoint()));
        std::tie(this->nearest_object_index, nearest_object_distance, *this->nearest_gaze_position_ptr) = this->findGazedObject(*this->nose_tip_position_ptr, *this->nose_end_point2D_drawtmp, this->object_centers);
        if(this->nearest_gaze_position_ptr && !this->isInvalidPoint(*this->nearest_gaze_position_ptr))
        {
            if(nearest_object_distance > this->gaze_assigned_thresh)
            {
                this->nearest_object_index = 0;
            }
            if(this->object_distance_cache.value() > this->person_distance_cache.value())
            {
                this->nearest_object_index = 0;
            }
        }
    }
}

std::tuple<std::size_t, double, cv::Point2i> CombiDarknetOpenface::findGazedObject(const cv::Point2i& nose_tip_point, const cv::Point2i& nose_end_point, const std::vector<cv::Point2i>& object_centers, float angle_distance_thresh) const
{
    double nose_angle = math_util::radToDeg(math_util::atan2(nose_end_point - nose_tip_point));
    std::size_t min_object_index = 0;
    float min_object_distance = std::numeric_limits<float>::infinity();
    cv::Point2i min_gaze_position = this->invalidPoint();
    for(double distance = this->search_distance_step; distance <= this->search_distance_range; distance += this->search_distance_step)
    {
        cv::Point2i target_point = this->getProjectedPoint(cv::Point3d(0,0,distance));
        double target_point_angle = math_util::radToDeg(math_util::atan2(target_point - nose_tip_point));
        if(abs(target_point_angle - nose_angle) > angle_distance_thresh)
        {
            continue;
        }
        std::vector<float> nose_end_distances;
        for(const auto& center_position : object_centers)
        {
            if(this->isInvalidPoint(center_position))
            {
                nose_end_distances.push_back(std::numeric_limits<float>::infinity());
                continue;
            }
            nose_end_distances.push_back(cv::norm(center_position - target_point));
        }
        if(std::count(nose_end_distances.begin(), nose_end_distances.end(), std::numeric_limits<float>::infinity()) == nose_end_distances.size())
        {
            continue;
        }
        auto min_element_iter = std::min_element(std::begin(nose_end_distances), std::end(nose_end_distances));
        if(*min_element_iter < min_object_distance)
        {
            min_object_index = std::distance(nose_end_distances.begin(), min_element_iter);
            min_object_distance = *min_element_iter;
            min_gaze_position = target_point;

        }
    }
    return std::forward_as_tuple(min_object_index, min_object_distance, min_gaze_position);
}

void CombiDarknetOpenface::calculateTimeUseOutofView()
{
    this->nearest_object_index = 0;
    this->nearest_gaze_position_ptr.reset();
}

void CombiDarknetOpenface::onRgbImageUpdated(const sensor_msgs::ImageConstPtr& msg) const
{
    cv_bridge::CvImagePtr cv_ptr;
    try
    {
        cv_ptr = cv_bridge::toCvCopy(msg, msg->encoding);
    }
    catch (cv_bridge::Exception& ex)
    {
        ROS_ERROR_STREAM("cv_bridge exception: " << ex.what());
        return;
    }
    if(this->nose_tip_position_ptr && this->isInImageArea(*this->nose_tip_position_ptr))
    {
        cv::circle(cv_ptr->image, *nose_tip_position_ptr, 5, cv::Scalar(0,0,255), 5);
    }
    if(this->nose_end_point2D_drawtmp && this->isInImageArea(*this->nose_end_point2D_drawtmp))
    {
        cv::circle(cv_ptr->image, *nose_end_point2D_drawtmp, 5, cv::Scalar(0,0,255), 5);
    }
    if(this->nose_tip_position_ptr && this->nose_end_point2D_drawtmp && this->isInImageArea(*this->nose_tip_position_ptr) && this->isInImageArea(*this->nose_end_point2D_drawtmp))
    {
        cv::line(cv_ptr->image, *this->nose_tip_position_ptr, *this->nose_end_point2D_drawtmp, cv::Scalar(0, 255, 0), 2);
    }
    if(this->nearest_gaze_position_ptr)
    {
        cv::circle(cv_ptr->image, *nearest_gaze_position_ptr, 7, cv::Scalar(0,0,255), 5);
    }
    cv::imshow("RGB image", cv_ptr->image);
    cv::waitKey(10);
}

void CombiDarknetOpenface::onDepthImageUpdated(const sensor_msgs::ImageConstPtr& msg)
{
    if(!this->person_box && !this->hasGazedObject())
    {
        return;
    }
    cv_bridge::CvImagePtr cv_ptr;
    try
    {
        cv_ptr = cv_bridge::toCvCopy(msg, msg->encoding);
    }
    catch (cv_bridge::Exception& ex)
    {
        ROS_ERROR_STREAM("cv_bridge exception: " << ex.what());
        return;
    }
    if(this->person_box)
    {
        cv::Rect person_area = nose_tip_position_ptr ?
            cv::Rect(nose_tip_position_ptr->x-1, nose_tip_position_ptr->y-1, 3, 3) :
            cv::Rect((this->person_box->tl() + this->person_box->br()) / 2, cv::Size(3, 3));
        double current_person_distance = this->getDepthValue(cv_ptr->image, person_area);
        this->person_distance_cache.update(current_person_distance);;
        double person_angle = this->angle_of_view / 2 - (static_cast<double>((this->person_box->tl().x + this->person_box->br().x) / 2) / this->image_size.width) * this->angle_of_view;
        cv::Point2d person_position = this->getPositionInGlobal(
            cv::Point2d(0, 0),
            robotyaw,
            this->person_distance_cache.value(),
            person_angle
        );
        this->publishPersonMeasurement(person_position, this->estimate_position_ptr);
        this->publishPersonMarker(person_angle, person_position);
    }

    if(this->hasGazedObject())
    {
        cv::Rect object_area = cv::Rect((this->getGazedObjectBox().tl() + this->getGazedObjectBox().br()) / 2, cv::Size(3, 3));
        double current_object_distance = this->getDepthValue(cv_ptr->image, object_area);
        this->object_distance_cache.update(current_object_distance);
        double object_angle = this->angle_of_view / 2 - (static_cast<double>((this->getGazedObjectBox().tl().x + this->getGazedObjectBox().br().x) / 2) / this->image_size.width) * this->angle_of_view;
        cv::Point2d object_position = this->getPositionInGlobal(
            cv::Point2d(0, 0),
            robotyaw,
            this->object_distance_cache.value(),
            object_angle
        );
        this->publishObjectMarker(object_position);
    }
}

cv::Point2d CombiDarknetOpenface::getPositionInGlobal(const cv::Point2d& robot_trans, double robot_yaw, double distance, double angle)const
{
    Eigen::Matrix3d A;
    A << cos(math_util::degToRad(robot_yaw)), sin(math_util::degToRad(robot_yaw)), robot_trans.x,
        sin(math_util::degToRad(robot_yaw)), cos(math_util::degToRad(robot_yaw)), robot_trans.y,
        0,  0,   1;
    Eigen::Matrix3d B;
    B << cos(math_util::degToRad(angle)), sin(math_util::degToRad(angle)), distance * cos(math_util::degToRad(angle)),
        sin(math_util::degToRad(angle)), cos(math_util::degToRad(angle)), distance * sin(math_util::degToRad(angle)),
        0,  0,   1;
    Eigen::Matrix3d T = A * B;
    return cv::Point2d(T(0,2), T(1,2));
}

double CombiDarknetOpenface::getDepthValue(const cv::Mat& depth_image, const cv::Rect& area)const
{
    cv::Mat conv_image;
    depth_image.convertTo(conv_image, CV_64F);
    double ret_val = 0.0;
    for(std::size_t i = area.tl().y; i < area.br().y; ++i)
    {
        for(std::size_t j = area.tl().x; j < area.br().x; ++j)
        {
            if(conv_image.at<double>(i, j) > 0.0)
            {
                ret_val = static_cast<double>(conv_image.at<double>(i, j));
            }
        }
    }
    return ret_val;
}

void CombiDarknetOpenface::onPersonPositionEstimated(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
    static std::unique_ptr<cv::Point2d> last_estimate_position_ptr;
    if(!last_estimate_position_ptr)
    {
        last_estimate_position_ptr.reset(new cv::Point2d(msg->pose.position.x, msg->pose.position.y));
    }
    this->estimate_position_ptr.reset(new cv::Point2d(msg->pose.position.x, msg->pose.position.y));
    this->person_velocity_ptr.reset(new cv::Vec2d((*this->estimate_position_ptr) - (*last_estimate_position_ptr)));
    last_estimate_position_ptr.reset(new cv::Point2d(*this->estimate_position_ptr));
    if(!this->person_box)
    {
        this->person_moving_state = PersonMovingState::Unrecognized;
    }
    else if((abs((*this->person_velocity_ptr)[0]) > 0.04 && abs((*this->person_velocity_ptr)[1]) > 0.04) ||
            (abs((*this->person_velocity_ptr)[0]) > 0.06 || abs((*this->person_velocity_ptr)[1]) > 0.06))
    {
        person_move_cnt += 1;
    }
    else
    {
        this->person_moving_state = PersonMovingState::Stopping;
        person_move_cnt = 0;
    }
    if(person_move_cnt>4)
    {
        this->person_moving_state = PersonMovingState::Moving;
    }

    this->publishEstimatedPersonPositionMarker(*this->estimate_position_ptr);
}

void CombiDarknetOpenface::publishEstimatedPersonPositionMarker(const cv::Point2d& position) const
{
    visualization_msgs::Marker estimate_person_marker;
    estimate_person_marker.header.stamp = ros::Time::now();
    estimate_person_marker.header.frame_id = this->FIXED_FRAME;
    estimate_person_marker.pose.position.x = position.x;
    estimate_person_marker.pose.position.y = position.y;
    estimate_person_marker.lifetime = ros::Duration(0);
    estimate_person_marker.scale.x = 0.2;
    estimate_person_marker.scale.y = 0.2;
    estimate_person_marker.scale.z = 0.01;
    estimate_person_marker.type = visualization_msgs::Marker::SPHERE;
    estimate_person_marker.color.a = 0.75;
    estimate_person_marker.color.r = 1.0;
    estimate_person_marker.color.g = 0;
    estimate_person_marker.color.b = 1.0;
    this->estimate_marker_pub.publish(estimate_person_marker);
}

void CombiDarknetOpenface::publishPersonMeasurement(const cv::Point2d& position, const std::unique_ptr<cv::Point2d>& estimated_position) const
{
    static int kf_failed_cnt = 0;
    geometry_msgs::PoseStamped input_pose;
    input_pose.pose.position.x = position.x;
    input_pose.pose.position.y = position.y;
    if(!estimated_position)
    {
        input_pose.pose.position.z = 1;
        measurement_pub.publish(input_pose);
    }
    else
    {
        double dist_error = cv::norm(position - *estimated_position);
        if(dist_error<1.2)
        {
            kf_failed_cnt = 0;
            input_pose.pose.position.z = 0;
            measurement_pub.publish(input_pose);
        }
        else
        {
            if(kf_failed_cnt == 5)
            {
                input_pose.pose.position.z = 1;
                measurement_pub.publish(input_pose);
                kf_failed_cnt = 0;
            }
            else
            {
                kf_failed_cnt += 1;
            }
        }
    }
}

void CombiDarknetOpenface::publishHeadPoseArrow(const cv::Point2d& position, double head_arrow_angle_deg) const
{
    visualization_msgs::Marker head_pose_arrow;
    head_pose_arrow.header.frame_id = FIXED_FRAME;
    head_pose_arrow.header.stamp = ros::Time::now();
    head_pose_arrow.ns = "basic_shapes";
    head_pose_arrow.type = visualization_msgs::Marker::ARROW;
    head_pose_arrow.action = visualization_msgs::Marker::ADD;
    head_pose_arrow.pose.position.x = position.x;
    head_pose_arrow.pose.position.y = position.y;
    head_pose_arrow.pose.orientation = tf::createQuaternionMsgFromYaw(math_util::degToRad(head_arrow_angle_deg));
    head_pose_arrow.scale.x = 0.3;
    head_pose_arrow.scale.y = 0.1;
    head_pose_arrow.scale.z = 0.1;
    head_pose_arrow.color.r = 0.0f;
    head_pose_arrow.color.g = 0.0f;
    head_pose_arrow.color.b = 1.0f;
    head_pose_arrow.color.a = 1.0f;
    head_pose_arrow.lifetime = ros::Duration();
    headpose_arrow_pub.publish(head_pose_arrow);
}

void CombiDarknetOpenface::publishPersonMarker(double theta, const cv::Point2d& position) const
{
    visualization_msgs::Marker person_arrow;
    person_arrow.header.frame_id = FIXED_FRAME;
    person_arrow.header.stamp = ros::Time::now();
    person_arrow.ns = "basic_shapes";
    person_arrow.type = visualization_msgs::Marker::ARROW;
    person_arrow.action = visualization_msgs::Marker::ADD;
    person_arrow.pose.position.x = position.x;
    person_arrow.pose.position.y = position.y;
    person_arrow.pose.orientation = tf::createQuaternionMsgFromYaw(math_util::degToRad(theta));
    person_arrow.scale.x = 0.3;
    person_arrow.scale.y = 0.1;
    person_arrow.scale.z = 0.1;
    person_arrow.color.r = 1.0f;
    person_arrow.color.g = 0.0f;
    person_arrow.color.b = 0.0f;
    person_arrow.color.a = 1.0f;
    person_arrow.lifetime = ros::Duration();
    person_arrow_pub.publish(person_arrow);

    visualization_msgs::Marker person_marker;
    person_marker.header.stamp = ros::Time::now();
    person_marker.header.frame_id = FIXED_FRAME;
    person_marker.pose.position.x = position.x;
    person_marker.pose.position.y = position.y;
    person_marker.lifetime = ros::Duration();
    person_marker.scale.x = 0.2;
    person_marker.scale.y = 0.2;
    person_marker.scale.z = 0.01;
    person_marker.type = visualization_msgs::Marker::CUBE;
    person_marker.color.a = 0.75;
    person_marker.color.g =  1.0;
    person_marker_pub.publish(person_marker);

    visualization_msgs::Marker origin_marker;
    origin_marker.header.stamp = ros::Time::now();
    origin_marker.header.frame_id = FIXED_FRAME;
    origin_marker.pose.position.x = 0.0;
    origin_marker.pose.position.y = 0.0;
    origin_marker.lifetime = ros::Duration();
    origin_marker.scale.x = 0.1;
    origin_marker.scale.y = 0.1;
    origin_marker.scale.z = 0.1;
    origin_marker.color.a = 1.0f;
    origin_marker.color.r = 1.0;
    origin_marker.color.g = 1.0;
    origin_marker.color.b = 1.0;
    origin_marker.type = visualization_msgs::Marker::SPHERE;
    origin_marker_pub.publish(origin_marker);
}

void CombiDarknetOpenface::publishObjectMarker(const cv::Point2d& position) const
{
    visualization_msgs::Marker object_marker;
    object_marker.header.stamp = ros::Time::now();
    object_marker.header.frame_id = FIXED_FRAME;
    object_marker.pose.position.x = position.x;
    object_marker.pose.position.y = position.y;
    object_marker.lifetime = ros::Duration(1);
    object_marker.scale.x = 0.2;
    object_marker.scale.y = 0.2;
    object_marker.scale.z = 0.01;
    object_marker.type = visualization_msgs::Marker::CUBE;
    object_marker.color.a = 0.75;
    object_marker.color.b =  1.0;
    object_marker_pub.publish(object_marker);
}

void CombiDarknetOpenface::publishGazedObjectDetectionResult(const std_msgs::Header& header) const
{
    if(!this->nose_tip_position_ptr || !this->nose_end_point2D_drawtmp)
    {
        return;
    }
    combi_darknet_openface::GazedObjectDetectionResult result;
    result.header = header;
    result.face.x = this->nose_tip_position_ptr->x;
    result.face.y = this->nose_tip_position_ptr->y;
    result.face.theta = math_util::atan2(*this->nose_end_point2D_drawtmp - *this->nose_tip_position_ptr);
    result.gazedObjectIndex = this->nearest_object_index == 0 ? combi_darknet_openface::GazedObjectDetectionResult::OUT_OF_VIEW : this->nearest_object_index - 1;
    for(std::size_t i = 1; i < this->class_names.size(); ++i)
    {
        combi_darknet_openface::DetectionResult d;
        d.Class = this->class_names[i];
        d.xmin = this->object_boxes[i].tl().x;
        d.ymin = this->object_boxes[i].tl().y;
        d.xmax = this->object_boxes[i].br().x;
        d.ymax = this->object_boxes[i].br().y;
        result.detectionResults.push_back(d);
    }
    this->gazed_object_detection_pub.publish(result);
}

cv::Point2i CombiDarknetOpenface::invalidPoint()const
{
    return cv::Point2i(std::numeric_limits<int>::max(), std::numeric_limits<int>::max());
}

bool CombiDarknetOpenface::isInvalidPoint(const cv::Point2i& p)const
{
    return p.x == std::numeric_limits<int>::max() && p.y == std::numeric_limits<int>::max();
}

bool CombiDarknetOpenface::hasGazedObject() const
{
    return this->nearest_object_index != 0;
}

cv::Rect CombiDarknetOpenface::getGazedObjectBox() const
{
    return this->object_boxes.at(this->nearest_object_index);
}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "combi_darknet_openface");
    volatile CombiDarknetOpenface pt;
    ros::spin();
    return 0;
}


// Copyright (c) 2019 Graphcore Ltd. All rights reserved.
#define BOOST_TEST_MODULE NonLinearityGradSweepTest
// Test for the Non-Linearity Gradient Operations
// Used to verify the accuracy of Non-Linearity backward processing functions
// over any part of the Half-Precision range for both inputs.

#include <poplar/CSRFunctions.hpp>
#include <poplar/Engine.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <popops/Zero.hpp>

#include "poputil/VertexTemplates.hpp"

#include "../lib/popops/ExprOpUtil.hpp"
#include "popops/ElementWise.hpp"
#include <poplibs_test/NonLinearity.hpp>
#include <poplibs_test/Util.hpp>
#include <popnn/NonLinearity.hpp>
#include <popnn/NonLinearityDefUtil.hpp>
#include <popnn/codelets.hpp>
#include <popops/codelets.hpp>
#include <poputil/TileMapping.hpp>

#include <boost/program_options.hpp>

#include <exception>

using namespace poplar;
using namespace poplar::program;
using namespace poputil;
using namespace poplibs_test::util;
using namespace popops;
using namespace popnn;
using namespace poplibs_support;

const poplar::OptionFlags options{{"debug.executionProfile", "compute_sets"}};

#define TOL 0.1 // tolerance of 0.1%
#define FLOAT_ATOL 1e-20
#define HALF_ATOL 1e-7

// The number of activation sweeps executed per tile
#define MAX_ACT_SWEEP_PER_TILE 512

// The largest number that can be represented in half-precision
#define HALF_0_TO_6 (0x45ff)

//*****************************************************************************
// Execute the Non-Linearity Gradient function over every possible range of
// valid Half-Precision value, and optionally stream the function output to
// stdout.
//
// The Non-Linearity Gradient function is a function of two tensors:
//    1. Tensor of activation inputs
//    2. Tensor of Gradient Outputs
//
// One sweep of the activation inputs is performed for every Gradient output
// value. Alternatively it is also possible to decimate the range of Gradient
// outputs used. In order to speed up execution, multiple activation sweeps
// can be packed into a single Tensor, with the following caveats/restrictions:
//   - The number of gradient output points in a single tensor must be a
//     power of 2.
//   - Every execution can only process the same multiple of activation sweeps.
//     Therefore if there are not sufficient activation sweeps left over for
//     the final execution, the final execution will be abandoned.
//
// Exclude the following values:
//   o     NaN: Exponent=0b11111 && Fraction=0
//   o  +/-Inf: Exponent=0b11111 && Fraction!=0
//
// The set of valid half-precision codes with the exponent in the range
// [0b00000 - 0b11110]. The full range can be defined by sweeping across the
// following ranges:
//   o codes from 0 to 0b0111101111111111 (0x7bff) and
//   o codes from -0 to 0b1111101111111111 (-0x7bff).
//
// The constant HALF_MAX_CODE is defined as 0x7bff. The full range of inputs is
// stored in the tensor in the following order:
//
//    Index into Tensor       Half-Precision Value
//    -----------------       --------------------
//          0                  -HALF_MAX_CODE
//          1                  -(HALF_MAX_CODE-1)
//          .                        .
//          .                        .
//          .                        .
//    HALF_MAX_CODE                 -0
//    HALF_MAX_CODE+1               +0
//          .                        .
//          .                        .
//          .                        .
//   (HALF_MAX_CODE+1)*2-1      HALF_MAX_CODE
//
//
// Format of the CSV output:
//   - The elements in each row are in increasing order with unit decimation.
//   - Each row corresponds to one particular Gradient Output code.
//   - The rows are arranged in decreasing order of Gradient Outputs, beginning
//     at the upper limit and proceeding to as close to the lower limit as
//     possible.
//
//*****************************************************************************
bool doNonLinearityGradSweep(const DeviceType &deviceType,
                             const unsigned ntiles, NonLinearityType nlType,
                             const unsigned activation_max_range,
                             const unsigned gradout_lower_code,
                             const unsigned gradout_upper_code,
                             const unsigned stride_vertex,
                             const unsigned stride_delta, bool doPrintTensors,
                             bool ignoreData) {

  unsigned sweep_activation = (activation_max_range + 1) * 2;
  unsigned num_elements = sweep_activation;

  auto invertHalfCodeSign = [](const unsigned code) { return code ^ 0x8000; };

  auto halfPrecisionToIndex = [invertHalfCodeSign](const unsigned code) {
    int index = ((code & (1 << 15)) ? (-invertHalfCodeSign(code) - 1) : code);
    return index;
  };

  auto indexToHalfPrecision = [invertHalfCodeSign](const int index) {
    unsigned code = ((index >= 0) ? index : invertHalfCodeSign((-index) - 1));
    return code;
  };

  int sweep_upper = halfPrecisionToIndex(gradout_upper_code);
  int sweep_lower = halfPrecisionToIndex(gradout_lower_code);
  if (sweep_lower >= sweep_upper) {
    throw std::invalid_argument("range-hi must be greater than range-lo!");
  }

  const std::size_t xSize = num_elements * stride_vertex;

  auto device = createTestDevice(deviceType, 1, ntiles);
  auto &target = device.getTarget();
  Graph graph(target);
  popnn::addCodelets(graph);

  // layer parameters
  auto actH = graph.addVariable(HALF, {xSize}, "actH");
  auto deltaH = graph.addVariable(HALF, {xSize}, "deltaH");

  // arbitrary mappings
  mapTensorLinearly(graph, actH);
  mapTensorLinearly(graph, deltaH);

  graph.createHostWrite("inH", actH);
  graph.createHostWrite("inDeltaH", deltaH);
  graph.createHostRead("outDeltaH", deltaH);

  //
  // Sweep_Full the full range of valid half values
  //
  // Exclude the following values:
  //   o     NaN: Exponent=0b11111 && Fraction=0
  //   o  +/-Inf: Exponent=0b11111 && Fraction!=0
  //
  // The set of valid half codes is the set of all values with the
  // exponent in the range [0b00000 - 0b11110]
  //
  std::vector<uint16_t> rawHActInH(xSize);
  std::vector<uint16_t> rawHDeltaInH(xSize);
  std::vector<uint16_t> rawHDeltaOutH(xSize);

  for (unsigned block = 0; block < stride_vertex; block++) {
    for (unsigned col = 0; (col < sweep_activation / 2) && (col < num_elements);
         col++) {
      rawHActInH[(block * sweep_activation) + col] =
          (((sweep_activation / 2) - 1 - col) | (1 << 15));
    }

    for (unsigned col = 0; (col < sweep_activation / 2) &&
                           ((col + (sweep_activation / 2)) < num_elements);
         col++) {
      rawHActInH[(block * sweep_activation) + (sweep_activation / 2) + col] =
          col;
    }
  }

  // reference results calculated in harness
  const auto batchSize = 1;
  boost::multi_array<double, 2> hRefActIn(boost::extents[batchSize][xSize]),
      hRefDeltaOut(boost::extents[batchSize][xSize]),
      hDeltaOutH(boost::extents[batchSize][xSize]);

  // Save the inputs
  poplar::copyDeviceHalfToDouble(target, rawHActInH.data(), &hRefActIn[0][0],
                                 xSize);

  int stride = stride_vertex * stride_delta;
  unsigned row_code = 0;
  unsigned total_rows = [](const unsigned a, const unsigned b) {
    return (a / b) * b;
  }((sweep_upper - sweep_lower), stride);
  for (int tile_row = sweep_upper; tile_row >= (sweep_lower + stride);
       tile_row -= stride) {

    std::cerr << std::hex << "GradOut:" << row_code << std::dec << "; Progress "
              << (sweep_upper - tile_row) << "/" << total_rows << "\n";

    // Construct activations within limited ranged, one stride_vertex of column
    // sweeps at a time
    for (unsigned block = 0; block < stride_vertex; block++) {
      for (unsigned col = 0; col < num_elements; col++) {
        int row = tile_row - (block * stride_delta);
        row_code = indexToHalfPrecision(row);
        rawHDeltaInH[(block * sweep_activation) + col] = row_code;
      }
    }

    // Obtain backward gradient calculation for reference
    poplar::copyDeviceHalfToDouble(target, rawHDeltaInH.data(),
                                   &hRefDeltaOut[0][0], xSize);
    poplibs_test::bwdNonLinearity(nlType, hRefActIn, hRefDeltaOut);

    // build and run the target code
    auto bwdProg = Sequence();

    //
    // Exceptions:
    //
    //    - Invalid Instruction    Enabled
    //    - Divide by 0            Enabled
    //    - Overflow               Disabled  // Overflow Expected
    //    - Stochastic Rounding    Disabled  // For deterministic results
    //    - NANOO                  Disabled  // Overflow Expected
    //
    setFloatingPointBehaviour(graph, bwdProg, {true, true, false, false, false},
                              "Set");

    auto deltaHH =
        nonLinearityInputGradient(graph, nlType, actH, deltaH, bwdProg);
    bwdProg.add(Copy(deltaHH, deltaH));
    Engine bwdEng(graph, bwdProg);
    device.bind([&](const Device &d) {
      bwdEng.load(d);
      bwdEng.writeTensor("inH", rawHActInH.data(),
                         rawHActInH.data() + rawHActInH.size());
      bwdEng.writeTensor("inDeltaH", rawHDeltaInH.data(),
                         rawHDeltaInH.data() + rawHDeltaInH.size());
      bwdEng.run();
      bwdEng.readTensor("outDeltaH", rawHDeltaOutH.data(),
                        rawHDeltaOutH.data() + rawHDeltaOutH.size());
    });

    poplar::copyDeviceHalfToDouble(target, rawHDeltaOutH.data(),
                                   &hDeltaOutH[0][0], xSize);

    /* Print output as a matrix */
    if (doPrintTensors) {
      for (auto block = 0U; block < xSize; block += sweep_activation) {
        unsigned i;
        for (i = 0U; i < num_elements - 1; ++i) {
          std::cout << hDeltaOutH[0][block + i] << ",";
        }
        std::cout << hDeltaOutH[0][block + i] << "\n";
      }
    }

    /* Check result */
    if (!ignoreData &&
        !checkIsClose("outDeltaH", hDeltaOutH, hRefDeltaOut, TOL, HALF_ATOL)) {
      return false;
    }
  }

  return true;
}

//******************************************************************************
int main(int argc, char **argv) {
  namespace po = boost::program_options;

  DeviceType deviceType;
  Type dataType;

  bool doPrintTensors = false;
  bool ignoreData = false;
  unsigned tiles = 1;
  unsigned stride = 1;
  std::string act_max_code_string;
  unsigned act_max_code = HALF_0_TO_6;
  NonLinearityType nlType;
  unsigned gradout_upper_code = 0x3C00;
  unsigned gradout_lower_code = 0x3BFF;
  std::string range_lo;
  std::string range_hi;

  po::options_description desc("Options");

  // clang-format off
  desc.add_options()
    ("help", "Print help")
    ("print",
     po::value<bool>(&doPrintTensors)->default_value(doPrintTensors),
     "Print the tensors")
    ("ignoreCheck",
     po::value<bool>(&ignoreData)->default_value(ignoreData),
     "Ignore mismatches from the expected results")
    ("device-type",
     po::value<DeviceType>(&deviceType)->required(),
     "Device Type")
    ("tiles",
     po::value<unsigned>(&tiles)->default_value(tiles),
     "Number of tiles to use")
    ("nl-type",
     po::value<NonLinearityType>(&nlType)->required(),
     "Non-linearity type")
    ("act-max-code",
     po::value<std::string>(&act_max_code_string)
         ->value_name(act_max_code_string),
     "The limit of the sweep of activations, as the Half-precision limit")
    ("stride",
     po::value<unsigned>(&stride)->default_value(stride),
     "Sweep Interval")
    ("range-lo",
     po::value<std::string>(&range_lo)->value_name("range-lo"),
     "Lower range of the sweep")
    ("range-hi",
     po::value<std::string>(&range_hi)->value_name("range-hi"),
     "Upper range of the sweep");
  // clang-format on

  po::variables_map vm;
  try {
    po::store(po::parse_command_line(argc, argv, desc), vm);
    if (vm.count("help")) {
      std::cout << desc << "\n\n";
      return 1;
    }
    po::notify(vm);
  } catch (std::exception &e) {
    std::cerr << "error: " << e.what() << "\n";
    return 1;
  }

  if (stride & (stride - 1)) {
    throw std::invalid_argument("stride must be a power of 2!");
  }

  if (stride > MAX_ACT_SWEEP_PER_TILE) {
    throw std::invalid_argument("stride exceeds maximum!");
  }

  if (vm.count("act-max-code")) {
    act_max_code = std::stoul(act_max_code_string, nullptr, 0);
  }

  if (vm.count("range-lo")) {
    gradout_lower_code = std::stoul(range_lo, nullptr, 0);
  }

  if (vm.count("range-hi")) {
    gradout_upper_code = std::stoul(range_hi, nullptr, 0);
  }

  unsigned stride_vertex = (gradout_upper_code - gradout_lower_code) / stride;
  if (stride_vertex > MAX_ACT_SWEEP_PER_TILE) {
    stride_vertex = MAX_ACT_SWEEP_PER_TILE;
  }

  std::cerr << "Input arguments:"
            << "\n";
  std::cerr << "  Activation range-lo = 0x" << std::hex
            << [](const unsigned code) { return code ^ 0x8000; }(act_max_code)
            << std::dec << "\n";

  std::cerr << "  Activation range-hi = 0x" << std::hex << act_max_code
            << std::dec << "\n";
  std::cerr << "  GradOut range-lo = 0x" << std::hex << gradout_lower_code
            << std::dec << "\n";
  std::cerr << "  GradOut range-hi = 0x" << std::hex << gradout_upper_code
            << std::dec << "\n";
  std::cerr << "  GradOut decimation factor = " << stride << "\n";
  std::cerr << "  Activation Sweeps per Vertex = " << stride_vertex << "\n\n";

  return !doNonLinearityGradSweep(
      deviceType, tiles, nlType, act_max_code, gradout_lower_code,
      gradout_upper_code, stride_vertex, stride, doPrintTensors, ignoreData);
}


//
// This module looks for early tracker digis (before the flash blanking)
// and uses them to estimate the proton bunch time of this event WRT
// the DAQ clock.  This is used mainly as a diagnostic.
//
// Original author David Brown, 10/1/2020 LBNL
//
// framework
#include "art/Framework/Principal/Event.h"
#include "art_root_io/TFileService.h"
#include "art/Framework/Core/EDProducer.h"
#include "art/Framework/Core/ModuleMacros.h"
#include "Offline/ProditionsService/inc/ProditionsHandle.hh"
#include "Offline/RecoDataProducts/inc/StrawDigi.hh"
#include "Offline/RecoDataProducts/inc/ProtonBunchTime.hh"
#include "Offline/TrackerConditions/inc/StrawResponse.hh"
#include "TH1F.h"
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/error_of.hpp>
#include <boost/accumulators/statistics/error_of_mean.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/mean.hpp>

#include <iostream>
#include <limits>

namespace mu2e {
  using namespace TrkTypes;

  class ProtonBunchTimeFromStrawDigis : public art::EDProducer
  {
    public:
      using Name=fhicl::Name;
      using Comment=fhicl::Comment;

      struct Config {
	fhicl::Atom<int> debug{ Name("DebugLevel"), Comment("Debug level"), 0};
	fhicl::Atom<int> diag{ Name("DiagLevel"), Comment("Diag level"), 0};
	fhicl::Atom<int> print{ Name("PrintLevel"), Comment("Print level"), 0};
	fhicl::Atom<unsigned> nbins{ Name("NBins"), Comment("Number of bins in time histogram"), 30};
	fhicl::Atom<unsigned> minnhits{ Name("MinNHits"), Comment("Minimum number of hits needed to use result"), 100};
	fhicl::Atom<float> maxtime { Name("MaxDigiTime"), Comment("Maximum time to accept the digi as 'early' (nsec)"),200.0 };
	fhicl::Atom<float> meantime { Name("MeanTimeOffset"), Comment("Mean time offset of 'early' hits (nsec)"),99.5}; // this should come from a database FIXME!
	fhicl::Atom<float> phasewindow { Name("PhaseWindow"), Comment("Width of Event Marker Phase Window (nsec)"),25.0};
	fhicl::Atom<art::InputTag> strawDigiTag{ Name("StrawDigiTag"), Comment("StrawDigi producer"),"makeSD" };
      };
      using Parameters = art::EDProducer::Table<Config>;
      explicit ProtonBunchTimeFromStrawDigis(Parameters const& config);
      virtual ~ProtonBunchTimeFromStrawDigis();
      void produce( art::Event& e) override;
      void beginJob ( ) override;
    private:
      int debug_, diag_, print_;
      unsigned nbins_, minnhits_;
      float maxtime_, meantime_;
      float phasemean_, phaseerr_;
      art::InputTag ewmtag_, sdtag_;
      ProditionsHandle<StrawResponse> _strawResponse_h;
  };

  ProtonBunchTimeFromStrawDigis::ProtonBunchTimeFromStrawDigis(Parameters const& config) :
    EDProducer(config),
    debug_(config().debug()),
    diag_(config().diag()),
    print_(config().print()),
    nbins_(config().nbins()),
    minnhits_(config().minnhits()),
    maxtime_(config().maxtime()),
    meantime_(config().meantime()),
    sdtag_(config().strawDigiTag())
  {
    float phasewindow = config().phasewindow();
    phasemean_ = phasewindow/2.0;
    phaseerr_ = phasewindow/sqrt(12.0);
    consumes<StrawDigiCollection>(sdtag_);
    produces<ProtonBunchTime>();
  }

  ProtonBunchTimeFromStrawDigis::~ProtonBunchTimeFromStrawDigis() {}

  void ProtonBunchTimeFromStrawDigis::beginJob(){
 }

  void ProtonBunchTimeFromStrawDigis::produce(art::Event& event) {
    using namespace boost::accumulators;
    auto const& srep = _strawResponse_h.get(event.id());
    auto sdH = event.getValidHandle<StrawDigiCollection>(sdtag_);
    auto const& sdcol(*sdH);
    std::unique_ptr<ProtonBunchTime> pbto(new ProtonBunchTime());
// statistics accumulators
//
    accumulator_set<float, stats<tag::mean, tag::error_of<tag::mean> >> meanacc;

    TH1F* timeplot(0);
    if(debug_ > 1 && event.id().event() < 1000) {
      art::ServiceHandle<art::TFileService> tfs;
      char hname[100];
      snprintf(hname,100,"PBT%i",event.id().event());
      timeplot = tfs->make<TH1F>(hname,"time spectrum;nsec",nbins_,0.0,maxtime_);
    }

    for (size_t isd=0;isd<sdcol.size();++isd) {
      const StrawDigi& digi = sdcol[isd];
    // correct to physical time
      TDCTimes times;
      srep.calibrateTimes(digi.TDC(),times,digi.strawId());
      // take the early time
      float mintime = std::min(times[StrawEnd::hv],times[StrawEnd::cal]);
      if(mintime < maxtime_){
      // This can be improved by using TOT to approximately correct for the propagation time, and maybe peak/ped to correct
      // for time slewing TODO!
	meanacc(mintime);
// early hit: accumulate the histogram
	if(print_ > 1) std::cout << "Early hit time " << mintime << " SID " << digi.strawId() << std::endl;
	if(timeplot)timeplot->Fill(mintime);
      }
    }
    // if there aren't enought hits, set to 0
    unsigned nhits = extract_result<tag::count>(meanacc);
    if(nhits < minnhits_){
      pbto->pbtime_ = phasemean_;
      pbto->pbterr_ = phaseerr_;
    } else {
      pbto->pbtime_ = extract_result<tag::mean>(meanacc) - meantime_;
      pbto->pbterr_ = extract_result<tag::error_of<tag::mean>>(meanacc);
    }
    event.put(std::move(pbto));
  }
}

using mu2e::ProtonBunchTimeFromStrawDigis;
DEFINE_ART_MODULE(ProtonBunchTimeFromStrawDigis);


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2014-2019, Oracle and/or its affiliates.

// Contributed and/or modified by Menelaos Karavelas, on behalf of Oracle
// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

// Licensed under the Boost Software License version 1.0.
// http://www.boost.org/users/license.html

#ifndef BOOST_GEOMETRY_ALGORITHMS_DETAIL_IS_VALID_SEGMENT_HPP
#define BOOST_GEOMETRY_ALGORITHMS_DETAIL_IS_VALID_SEGMENT_HPP

#include <boost/core/ignore_unused.hpp>

#include <boost/geometry/core/point_type.hpp>
#include <boost/geometry/core/tags.hpp>

#include <boost/geometry/algorithms/assign.hpp>
#include <boost/geometry/algorithms/validity_failure_type.hpp>
#include <boost/geometry/algorithms/detail/equals/point_point.hpp>
#include <boost/geometry/algorithms/detail/is_valid/has_invalid_coordinate.hpp>
#include <boost/geometry/algorithms/dispatch/is_valid.hpp>

namespace boost { namespace geometry
{

#ifndef DOXYGEN_NO_DISPATCH
namespace dispatch
{

// A segment is a curve.
// A curve is simple if it does not pass through the same point twice,
// with the possible exception of its two endpoints
// A curve is 1-dimensional, hence we have to check is the two
// endpoints of the segment coincide, since in this case it is
// 0-dimensional.
//
// Reference: OGC 06-103r4 (6.1.6.1)
template <typename Segment>
struct is_valid<Segment, segment_tag>
{
    template <typename VisitPolicy, typename Strategy>
    static inline bool apply(Segment const& segment, VisitPolicy& visitor, Strategy const& strategy)
    {
        boost::ignore_unused(visitor);

        typename point_type<Segment>::type p[2];
        detail::assign_point_from_index<0>(segment, p[0]);
        detail::assign_point_from_index<1>(segment, p[1]);

        if (detail::is_valid::has_invalid_coordinate
                <
                    Segment
                >::apply(segment, visitor))
        {
            return false;
        }
        else if (! detail::equals::equals_point_point(p[0], p[1], strategy))
        {
            return visitor.template apply<no_failure>();
        }
        else
        {
            return
                visitor.template apply<failure_wrong_topological_dimension>();
        }
    }
};

} // namespace dispatch
#endif // DOXYGEN_NO_DISPATCH

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_ALGORITHMS_DETAIL_IS_VALID_SEGMENT_HPP


/* Copyright © 2017 Apple Inc. All rights reserved.
 *
 * Use of this source code is governed by a BSD-3-clause license that can
 * be found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
 */
const bool enable_expensive_tests = false;

#define BOOST_TEST_MODULE
#include <boost/test/unit_test.hpp>
#include <util/test_macros.hpp>
#include <vector>
#include <string>
#include <functional>

#include <random/random.hpp>

#include <sframe/sframe_iterators.hpp>
#include <unity/toolkits/recsys/models/linear_models/linear_model.hpp>
#include <unity/toolkits/recsys/models/linear_models/factorization_model.hpp>
#include <unity/toolkits/recsys/models/linear_models/matrix_factorization.hpp>
#include <unity/toolkits/util/data_generators.hpp>
#include <unity/toolkits/ml_data_2/ml_data.hpp>
#include <unity/toolkits/ml_data_2/ml_data_iterator.hpp>
#include <sframe/testing_utils.hpp>
#include <util/testing_utils.hpp>

#include <cfenv>

using namespace turi;
using namespace turi::recsys;

void run_exact_test(
    const std::vector<size_t>& n_categorical_values,
    std::map<std::string, flexible_type> opts,
    const std::string& model_type) {

  bool binary_target = false;

  if(model_type == "linear" || model_type == "fm" || model_type == "mf") {
    binary_target = false;
    opts["y_mode"] = "squared_error";
  } else if(model_type == "logistic"  || model_type == "logistic_fm" || model_type == "logistic_mf") {
    binary_target = true;
    opts["y_mode"] = "logistic";
  } else {
    DASSERT_TRUE(false);
  }

  if(model_type == "mf" || model_type == "logistic_mf")
    opts["only_2_factor_terms"] = true;

  std::feraiseexcept(FE_ALL_EXCEPT);

  size_t n_observations = opts.at("n_observations");
  opts.erase("n_observations");

  std::string target_column_name = "target";

  std::vector<std::string> column_names = {"user", "item"};

  DASSERT_NE(n_categorical_values[0], 0);
  DASSERT_NE(n_categorical_values[1], 0);

  for(size_t i = 2; i < n_categorical_values.size(); ++i) {
    column_names.push_back("C-" + std::to_string(i));
  }

  lm_data_generator lmdata(column_names, n_categorical_values, opts);

  sframe train_data = lmdata.generate(n_observations, target_column_name, 0, 0);
  sframe test_data = lmdata.generate(n_observations, target_column_name, 1, 0);

  std::map<std::string, flexible_type> options =
      { {"optimization_method", "auto"},
        {"binary_target", binary_target},
        {"target", target_column_name},
        {"regularization", 0},
        {"sgd_step_size", 0},
        {"max_iterations", binary_target ? 200 : 100},
        {"sgd_convergence_threshold", 1e-10},

        // Add in the ranking regularizer
        {"ranking_regularization", 0.1},
        {"unobserved_rating_value", 0}
      };

  opts.erase("y_mode");

  if(model_type == "mf" || model_type == "logistic_mf")
    opts.erase("only_2_factor_terms");

  options.insert(opts.begin(), opts.end());

  if(model_type == "fm" || model_type == "logistic_fm"
     || model_type == "mf" || model_type == "logistic_mf") {
    options["linear_regularization"] = 0;
  }

  // Instantiate some alternate versions to make sure the save and load work
  typedef std::shared_ptr<recsys::recsys_model_base> model_ptr;

  auto new_model = [&]() -> model_ptr {
    model_ptr ret;

    if(model_type == "linear" || model_type == "logistic") {
      ret.reset(new recsys_linear_model);
    } else if(model_type == "fm" || model_type == "logistic_fm") {
      ret.reset(new recsys_factorization_model);
    } else if(model_type == "mf" || model_type == "logistic_mf") {
      ret.reset(new recsys_matrix_factorization);
    } else {
      ASSERT_TRUE(false);
    }

    return ret;
  };

  model_ptr model = new_model();

  model->init_option_info();

  model->set_options(options);

  model->setup_and_train(train_data);

  // Instantiate some alternate versions to make sure the save and load work

  std::vector<model_ptr> all_models =
      {model,
       new_model(),
       model_ptr((recsys::recsys_model_base*)(model->clone()))};

  ////////////////////////////////////////
  {
    // Save it
    dir_archive archive_write;
    archive_write.open_directory_for_write("linear_regression_cxx_tests");

    turi::oarchive oarc(archive_write);

    oarc << *model;

    archive_write.close();

    // Load it
    dir_archive archive_read;
    archive_read.open_directory_for_read("linear_regression_cxx_tests");

    turi::iarchive iarc(archive_read);

    iarc >> (*all_models[1]);
  }

  for(model_ptr m : all_models) {

    sframe y_hat_sf = m->predict(model->create_ml_data(test_data));
    std::vector<double> y_hat = testing_extract_column<double>(y_hat_sf.select_column(0));

    {
      size_t i = 0;

      for(ml_data_iterator it(m->create_ml_data(test_data)); !it.done(); ++it, ++i) {
        double y = it.target_value();

        if(!binary_target) {
          if(y > 0)
            TS_ASSERT_LESS_THAN(y_hat[i], 1.1 * (y + 0.5));
        } else {
          if(y == 0)
            TS_ASSERT_LESS_THAN(y_hat[i], 0.75);
        }
      }
    }
  }
}

struct linear_tests  {
 public:

  ////////////////////////////////////////////////////////////////////////////////

  void test_regression_se_really_bloody_basic_2d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      100 } };

    run_exact_test({1, 1}, opts, "linear");
  }

  ////////////////////////////////////////////////////////////////////////////////

  void test_regression_se_basic_3d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      100 } };

    // run_exact_test({1,1,0}, opts, "linear");
  }

  ////////////////////////////////////////////////////////////////////////////////

  void test_regression_se_basic_5d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      250 } };

    // run_exact_test({1,1,0,0,0}, opts, "linear");
  }

  // ////////////////////////////////////////////////////////////////////////////////

  void test_regression_se_multiuser_basic_2d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      500 } };

    // run_exact_test({2,2,0,0,0}, opts, "linear");
  }

  // ////////////////////////////////////////////////////////////////////////////////

  void test_regression_se_large_no_side() {

    std::map<std::string, flexible_type> opts = {
      {"n_observations",      100000 } };

    run_exact_test({100,100}, opts, "linear");
  }

  // // ////////////////////////////////////////////////////////////////////////////////

  void test_regression_se_large_some_side() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      5*10*10 } };

    run_exact_test({10,10,0,0,0}, opts, "linear");
  }
};

struct log_linear_tests  {
 public:

  ////////////////////////////////////////////////////////////////////////////////

  void test_regression_log_really_bloody_basic_2d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      100 } };

    run_exact_test({1, 1}, opts, "logistic");
  }

  ////////////////////////////////////////////////////////////////////////////////

  void test_regression_log_basic_3d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      100 } };

    // run_exact_test({1,1,0}, opts, "logistic");
  }

  // ////////////////////////////////////////////////////////////////////////////////

  void test_regression_log_large_no_side() {

    if(enable_expensive_tests) {
      std::map<std::string, flexible_type> opts = {
        {"n_observations",     10*100*100 } };

      run_exact_test({100,100}, opts, "logistic");
    }
  }
};

struct factorization_tests  {
 public:

  ////////////////////////////////////////////////////////////////////////////////

  void test_factorization_se_really_bloody_basic_2d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      10 },
      {"n_factors",           1 } };

    run_exact_test({1, 1}, opts, "fm");
  }

  // // ////////////////////////////////////////////////////////////////////////////////

  void test_factorization_se_many_factors() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",           5 } };

    run_exact_test({1, 8}, opts, "fm");
  }

  // // ////////////////////////////////////////////////////////////////////////////////

  void test_factorization_se_many_columns() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",           1 } };

    run_exact_test({16, 1, 1, 1}, opts, "fm");
  }
};

struct log_factorization_tests  {
 public:

  ////////////////////////////////////////////////////////////////////////////////

  void test_factorization_log_really_bloody_basic_2d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      10 },
      {"n_factors",           1 } };

    run_exact_test({1, 1}, opts, "logistic_fm");
  }

  // // ////////////////////////////////////////////////////////////////////////////////

  void test_factorization_log_many_factors() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",          5 } };

    run_exact_test({1, 8}, opts, "logistic_fm");
  }

  // // ////////////////////////////////////////////////////////////////////////////////

  void test_factorization_log_many_categories() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      2000 },
      {"n_factors",           1 } };

    run_exact_test({2, 50}, opts, "logistic_fm");
  }

  void test_factorization_log_many_dimensions() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",           1 } };

    run_exact_test({16, 1, 1, 1}, opts, "logistic_fm");
  }
};

struct matrix_factorization_tests  {
 public:

  ////////////////////////////////////////////////////////////////////////////////

  void test_mf_se_really_bloody_basic_2d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      10 },
      {"n_factors",           1 } };

    run_exact_test({1, 1}, opts, "mf");
  }

  // // ////////////////////////////////////////////////////////////////////////////////

  void test_mf_se_many_factors() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",           5 } };

    run_exact_test({8, 1}, opts, "mf");
  }

  // // ////////////////////////////////////////////////////////////////////////////////

  void test_mf_se_many_categories() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",           1 } };

    run_exact_test({2, 50}, opts, "mf");
  }

  void test_mf_se_many_columns() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",           1 } };

    run_exact_test({16, 1, 1, 1}, opts, "mf");
  }
};

struct log_mf_tests  {
 public:

  ////////////////////////////////////////////////////////////////////////////////

  void test_mf_log_really_bloody_basic_2d() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      10 },
      {"n_factors",           1 } };

    run_exact_test({1, 1}, opts, "logistic_mf");
  }

  // // ////////////////////////////////////////////////////////////////////////////////

  void test_mf_log_many_factors() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",           5 } };

    run_exact_test({5, 5}, opts, "logistic_mf");
  }

  // // ////////////////////////////////////////////////////////////////////////////////

  void test_mf_log_many_categories() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",           1 } };

    run_exact_test({2, 30}, opts, "logistic_mf");
  }

  void test_mf_log_many_dimensions() {
    std::map<std::string, flexible_type> opts = {
      {"n_observations",      1000 },
      {"n_factors",           1 } };

    run_exact_test({16, 1, 1, 1}, opts, "logistic_mf");
  }
};

struct class pure_ranking  {
//  public:

//   void run_pure_ranking_test(size_t n_users, size_t n_items,
//                              size_t n_items_train, size_t n_items_test,
//                              std::map<std::string, flexible_type> opts,
//                              const std::string& model_type) {

//     std::feraiseexcept(FE_ALL_EXCEPT);

//     std::vector<std::string> column_names = {"user", "item"};

//     opts["y_mode"] = "squared_error";
//     lm_data_generator lmdata(column_names, {n_users, n_items}, opts);
//     opts.erase("y_mode");

//     sframe train_data, test_data;

//     std::tie(train_data, test_data) = lmdata.generate_for_ranking(
//         n_items_train, n_items_test, 0, 0);

//     std::shared_ptr<recsys_model_base> model;

//     std::map<std::string, flexible_type> options =
//         { {"optimization_method", "auto"},
//           {"sgd_step_size", 0},
//           {"max_iterations", 20},
//           {"sgd_convergence_threshold", 1e-6},

//           // No target column; so should be able to properly rank
//           // things.
//         };

//     options.insert(opts.begin(), opts.end());

//     // Instantiate some alternate versions to make sure the save and load work
//     typedef std::shared_ptr<recsys::recsys_model_base> model_ptr;

//     auto new_model = [&]() {
//       model_ptr ret;

//       if(model_type == "linear" || model_type == "logistic") {
//         ret.reset(new recsys_linear_model);
//       } else if(model_type == "fm" || model_type == "logistic_fm") {
//         ret.reset(new recsys_factorization_model);
//       } else if(model_type == "mf" || model_type == "logistic_mf") {
//         ret.reset(new recsys_matrix_factorization);
//       } else {
//         ASSERT_TRUE(false);
//       }

//       return ret;
//     };

//     model = new_model();

//     model->set_options(options);

//     model->setup_and_train(train_data);

//     std::vector<model_ptr> all_models =
//         {model,
//          new_model(),
//          model_ptr((recsys::recsys_model_base*)(model->clone()))};

//     ////////////////////////////////////////
//     {
//       // Save it
//       dir_archive archive_write;
//       archive_write.open_directory_for_write("linear_regression_cxx_tests");

//       turi::oarchive oarc(archive_write);

//       oarc << *model;

//       archive_write.close();

//       // Load it
//       dir_archive archive_read;
//       archive_read.open_directory_for_read("linear_regression_cxx_tests");

//       turi::iarchive iarc(archive_read);

//       iarc >> (*all_models[1]);
//     }

//     auto build_map = [&](const sframe& user_item_data) {
//       std::map<size_t, std::set<size_t> > user_map;

//       for(parallel_sframe_iterator it(user_item_data); !it.done(); ++it) {
//         size_t user = it.value(0);
//         size_t item = it.value(1);

//         user_map[user].insert(item);
//       }

//       return user_map;
//     };

//     auto true_user_map = build_map(test_data);

//     for(model_ptr m : all_models) {

//       sframe rec_out = m->recommend(nullptr, n_items_test);

//       auto user_map = build_map(rec_out);

//       for(size_t i = 0; i < n_users; ++i) {

//         // ASSERT_EQ(user_map[i], true_user_map[i]);
//       }
//     }
//   }

//   void test_linear_1() {
//     std::map<std::string, flexible_type> opts = {
//       {"_num_sampled_negative_examples", 1} };

//     run_pure_ranking_test(100, 10, 2, 1, opts, "linear");
//   }

//   void test_linear_2() {
//     std::map<std::string, flexible_type> opts;

//     run_pure_ranking_test(1000, 50, 20, 5, opts, "linear");
//   }

//   ////////////////////////////////////////////////////////////
//   // Test the different aspects of the procedure to sample negative
//   // points.  There a number of internal asserts to check that the
//   // sampling is consistent; below are the set of things which would
//   // likely hit the corner cases of the sampling algorithm.

//   void test_linear_most_items_rated_1() {
//     std::map<std::string, flexible_type> opts = {
//       {"_num_sampled_negative_examples", 1} };

//     run_pure_ranking_test(1000, 32, 31, 1, opts, "linear");
//   }

//   void test_linear_most_items_rated_2() {
//     std::map<std::string, flexible_type> opts = {
//       {"_num_sampled_negative_examples", 16} };

//     run_pure_ranking_test(10, 32, 31, 1, opts, "linear");
//   }

//     void test_linear_most_items_rated_1b() {
//     std::map<std::string, flexible_type> opts = {
//       {"_num_sampled_negative_examples", 1} };

//     run_pure_ranking_test(10, 65, 63, 1, opts, "linear");
//   }

//   void test_linear_most_items_rated_2b() {
//     std::map<std::string, flexible_type> opts = {
//       {"_num_sampled_negative_examples", 16} };

//     run_pure_ranking_test(10, 65, 64, 1, opts, "linear");
//   }

//   void test_linear_most_items_rated_3() {
//     std::map<std::string, flexible_type> opts = {
//       {"_num_sampled_negative_examples", 4} };

//     run_pure_ranking_test(10, 100, 90, 4, opts, "linear");
//   }

//   void test_linear_few_items_rated_1() {
//     std::map<std::string, flexible_type> opts = {
//       {"_num_sampled_negative_examples", 4} };

//     run_pure_ranking_test(10, 1000, 2, 1, opts, "linear");
//   }

//   ////////////////////////////////////////////////////////////
//   // Test matrix factorization

//   void test_mf_1() {
//     std::map<std::string, flexible_type> opts = {
//       {"n_factors",           1 } };

//     run_pure_ranking_test(10, 10, 4, 2, opts, "mf");
//   }

//   void test_mf_2() {
//     std::map<std::string, flexible_type> opts = {
//       {"n_factors",           4 } };

//     run_pure_ranking_test(100, 10, 4, 2, opts, "mf");
//   }

//   void test_mf_3() {
//     std::map<std::string, flexible_type> opts = {
//       {"n_factors",           1 } };

//     run_pure_ranking_test(10, 50, 20, 5, opts, "mf");
//   }

//   void test_fm_1() {
//     std::map<std::string, flexible_type> opts = {
//       {"n_factors",           1 } };

//     run_pure_ranking_test(10, 10, 4, 2, opts, "fm");
//   }

//   void test_fm_2() {
//     std::map<std::string, flexible_type> opts = {
//       {"n_factors",           4 } };

//     run_pure_ranking_test(100, 10, 4, 2, opts, "fm");
//   }

//   void test_fm_3() {
//     std::map<std::string, flexible_type> opts = {
//       {"n_factors",           1 } };

//     run_pure_ranking_test(10, 50, 20, 5, opts, "fm");
//   }

// };

BOOST_FIXTURE_TEST_SUITE(_linear_tests, linear_tests)
BOOST_AUTO_TEST_CASE(test_regression_se_really_bloody_basic_2d) {
  linear_tests::test_regression_se_really_bloody_basic_2d();
}
BOOST_AUTO_TEST_CASE(test_regression_se_basic_3d) {
  linear_tests::test_regression_se_basic_3d();
}
BOOST_AUTO_TEST_CASE(test_regression_se_basic_5d) {
  linear_tests::test_regression_se_basic_5d();
}
BOOST_AUTO_TEST_CASE(test_regression_se_multiuser_basic_2d) {
  linear_tests::test_regression_se_multiuser_basic_2d();
}
BOOST_AUTO_TEST_CASE(test_regression_se_large_no_side) {
  linear_tests::test_regression_se_large_no_side();
}
BOOST_AUTO_TEST_CASE(test_regression_se_large_some_side) {
  linear_tests::test_regression_se_large_some_side();
}
BOOST_AUTO_TEST_SUITE_END()
BOOST_FIXTURE_TEST_SUITE(_log_linear_tests, log_linear_tests)
BOOST_AUTO_TEST_CASE(test_regression_log_really_bloody_basic_2d) {
  log_linear_tests::test_regression_log_really_bloody_basic_2d();
}
BOOST_AUTO_TEST_CASE(test_regression_log_basic_3d) {
  log_linear_tests::test_regression_log_basic_3d();
}
BOOST_AUTO_TEST_CASE(test_regression_log_large_no_side) {
  log_linear_tests::test_regression_log_large_no_side();
}
BOOST_AUTO_TEST_SUITE_END()
BOOST_FIXTURE_TEST_SUITE(_factorization_tests, factorization_tests)
BOOST_AUTO_TEST_CASE(test_factorization_se_really_bloody_basic_2d) {
  factorization_tests::test_factorization_se_really_bloody_basic_2d();
}
BOOST_AUTO_TEST_CASE(test_factorization_se_many_factors) {
  factorization_tests::test_factorization_se_many_factors();
}
BOOST_AUTO_TEST_CASE(test_factorization_se_many_columns) {
  factorization_tests::test_factorization_se_many_columns();
}
BOOST_AUTO_TEST_SUITE_END()
BOOST_FIXTURE_TEST_SUITE(_log_factorization_tests, log_factorization_tests)
BOOST_AUTO_TEST_CASE(test_factorization_log_really_bloody_basic_2d) {
  log_factorization_tests::test_factorization_log_really_bloody_basic_2d();
}
BOOST_AUTO_TEST_CASE(test_factorization_log_many_factors) {
  log_factorization_tests::test_factorization_log_many_factors();
}
BOOST_AUTO_TEST_CASE(test_factorization_log_many_categories) {
  log_factorization_tests::test_factorization_log_many_categories();
}
BOOST_AUTO_TEST_CASE(test_factorization_log_many_dimensions) {
  log_factorization_tests::test_factorization_log_many_dimensions();
}
BOOST_AUTO_TEST_SUITE_END()
BOOST_FIXTURE_TEST_SUITE(_matrix_factorization_tests, matrix_factorization_tests)
BOOST_AUTO_TEST_CASE(test_mf_se_really_bloody_basic_2d) {
  matrix_factorization_tests::test_mf_se_really_bloody_basic_2d();
}
BOOST_AUTO_TEST_CASE(test_mf_se_many_factors) {
  matrix_factorization_tests::test_mf_se_many_factors();
}
BOOST_AUTO_TEST_CASE(test_mf_se_many_categories) {
  matrix_factorization_tests::test_mf_se_many_categories();
}
BOOST_AUTO_TEST_CASE(test_mf_se_many_columns) {
  matrix_factorization_tests::test_mf_se_many_columns();
}
BOOST_AUTO_TEST_SUITE_END()
BOOST_FIXTURE_TEST_SUITE(_log_mf_tests, log_mf_tests)
BOOST_AUTO_TEST_CASE(test_mf_log_really_bloody_basic_2d) {
  log_mf_tests::test_mf_log_really_bloody_basic_2d();
}
BOOST_AUTO_TEST_CASE(test_mf_log_many_factors) {
  log_mf_tests::test_mf_log_many_factors();
}
BOOST_AUTO_TEST_CASE(test_mf_log_many_categories) {
  log_mf_tests::test_mf_log_many_categories();
}
BOOST_AUTO_TEST_CASE(test_mf_log_many_dimensions) {
  log_mf_tests::test_mf_log_many_dimensions();
}
BOOST_AUTO_TEST_SUITE_END()


#include <Eigen/Dense>
#include <memory>
#include <iostream>
#include "sparsetensor.h"
#include "latentprior.h"
#include "noisemodels.h"

using namespace Eigen;

void SparseMode::init(Eigen::MatrixXi &idx, Eigen::VectorXd &vals, int m, int mode_size) {
  if (idx.rows() != vals.size()) {
    throw std::runtime_error("idx.rows() must equal vals.size()");
  }
  num_modes = idx.cols();
  row_ptr.resize(mode_size + 1);
  row_ptr.setZero();
  values.resize(vals.size());
  indices.resize(idx.rows(), idx.cols() - 1);

  mode = m;
  auto rows = idx.col(mode);
  const int nrow = mode_size;
  nnz  = idx.rows();

  // compute number of non-zero entries per each element for the mode
  for (int i = 0; i < nnz; i++) {
    if (rows(i) >= mode_size) {
      throw std::runtime_error("SparseMode: mode value larger than mode_size");
    }
    row_ptr(rows(i))++;
  }
  // cumsum counts
  for (int row = 0, cumsum = 0; row < nrow; row++) {
    int temp     = row_ptr(row);
    row_ptr(row) = cumsum;
    cumsum      += temp;
  }
  row_ptr(nrow) = nnz;

  // writing idx and vals to indices and values
  for (int i = 0; i < nnz; i++) {
    int row  = rows(i);
    int dest = row_ptr(row);
    for (int j = 0, nj = 0; j < idx.cols(); j++) {
      if (j == mode) continue;
      indices(dest, nj) = idx(i, j);
      nj++;
    }
    //A->cols[dest] = cols[i];
    values(dest) = vals(i);
    row_ptr(row)++;
  }
  // fixing row_ptr
  for (int row = 0, prev = 0; row <= nrow; row++) {
    int temp     = row_ptr(row);
    row_ptr(row) = prev;
    prev         = temp;
  }
}

Eigen::MatrixXd MatrixData::getTestData() {
  MatrixXd coords( getTestNonzeros(), 3);
#pragma omp parallel for schedule(dynamic, 2)
  for (int k = 0; k < Ytest.outerSize(); ++k) {
    int idx = Ytest.outerIndexPtr()[k];
    for (SparseMatrix<double>::InnerIterator it(Ytest,k); it; ++it) {
      coords(idx, 0) = it.row();
      coords(idx, 1) = it.col();
      coords(idx, 2) = it.value();
      idx++;
    }
  }
  return coords;
}

void MatrixData::setTrain(int* rows, int* cols, double* values, int nnz, int nrows, int ncols) {
	Y.resize(nrows, ncols);
	sparseFromIJV(Y, rows, cols, values, nnz);
	Yt = Y.transpose();
	mean_value = Y.sum() / Y.nonZeros();
	dims.resize(2);
	dims << Y.rows(), Y.cols();
}

void MatrixData::setTest(int* rows, int* cols, double* values, int nnz, int nrows, int ncols) {
	Ytest.resize(nrows, ncols);
	sparseFromIJV(Ytest, rows, cols, values, nnz);
}

void MatrixData::setTrain(int* columns, int nmodes, double* values, int nnz, int* d) {
  if (nmodes != 2) {
    throw std::runtime_error("MatrixData: tensor training input not supported.");
  }
  auto idx  = toMatrix(columns, nnz, nmodes);
  auto vals = toVector(values, nnz);

  Y.resize(d[0], d[1]);
  sparseFromIJV(Y, idx, vals);

  Yt = Y.transpose();
  mean_value = Y.sum() / Y.nonZeros();
  dims.resize(2);
  dims << Y.rows(), Y.cols();
}

void MatrixData::setTest(int* columns, int nmodes, double* values, int nnz, int* d) {
  if (nmodes != 2) {
    throw std::runtime_error("MatrixData: tensor training input not supported.");
  }
  auto idx  = toMatrix(columns, nnz, nmodes);
  auto vals = toVector(values, nnz);

  Ytest.resize(d[0], d[1]);
  sparseFromIJV(Ytest, idx, vals);
}

void MatrixData::setCensoring(int* columns, int nmodes, int* censored, int nnz, int* d)
{std::cout << "No censoring" << std::endl;}

////////  TensorData  ///////

void TensorData::setTrain(Eigen::MatrixXi &idx, Eigen::VectorXd &vals, Eigen::VectorXi &d) {
  if (idx.rows() != vals.size()) {
    throw std::runtime_error("setTrain(): idx.rows() must equal vals.size()");
  }
  dims = d;
	mean_value = vals.mean();
  for (int mode = 0; mode < N; mode++) {
    Y->push_back( std::unique_ptr<SparseMode>(new SparseMode(idx, vals, mode, dims(mode))) );
  }
}

void TensorData::setTrain(int* columns, int nmodes, double* values, int nnz, int* dims) {
  auto idx  = toMatrix(columns, nnz, nmodes);
  auto vals = toVector(values, nnz);
  auto d    = toVector(dims, nmodes);
  setTrain(idx, vals, d);
}

void TensorData::setTest(int* columns, int nmodes, double* values, int nnz, int* dims) {
  auto idx  = toMatrix(columns, nnz, nmodes);
  auto vals = toVector(values, nnz);
  auto d    = toVector(dims, nmodes);
  setTest(idx, vals, d);
}

void TensorData::setTest(Eigen::MatrixXi &idx, Eigen::VectorXd &vals, Eigen::VectorXi &d) {
  if (idx.rows() != vals.size()) {
    throw std::runtime_error("setTest(): idx.rows() must equal vals.size()");
  }
  if ((d - dims).norm() != 0) {
    throw std::runtime_error("setTest(): train and test Tensor sizes are not equal.");
  }
  Ytest = SparseMode(idx, vals, 0, d(0));
}

void TensorData::setTrain(int* rows, int* cols, double* values, int nnz, int nrows, int ncols) {
  auto idx  = toMatrix(rows, cols, nnz);
  auto vals = toVector(values, nnz);

  VectorXi d(2);
  d << nrows, ncols;

  setTrain(idx, vals, d);
}

void TensorData::setTest(int* rows, int* cols, double* values, int nnz, int nrows, int ncols) {
  auto idx  = toMatrix(rows, cols, nnz);
  auto vals = toVector(values, nnz);

  VectorXi d(2);
  d << nrows, ncols;

  setTest(idx, vals, d);
}

Eigen::MatrixXd TensorData::getTestData() {
  MatrixXd coords( getTestNonzeros(), N + 1);
#pragma omp parallel for schedule(dynamic, 2)
  for (int k = 0; k < Ytest.modeSize(); ++k) {
    for (int idx = Ytest.row_ptr[k]; idx < Ytest.row_ptr[k+1]; idx++) {
      coords(idx, 0) = (double)k;
      for (int col = 0; col < Ytest.indices.cols(); col++) {
        coords(idx, col + 1) = (double)Ytest.indices(idx, col);
      }
      coords(idx, N) = Ytest.values(idx);
    }
  }
  return coords;
}

void TensorData::setCensoring(int* columns, int nmodes, int* censored, int nnz, int* d)
{
  std::cout <<"No censoring" << std::endl;
}

//MatrixDataCensored

Eigen::MatrixXd MatrixDataCensored::getTestData() {
  MatrixXd coords( getTestNonzeros(), 3);
#pragma omp parallel for schedule(dynamic, 2)
  for (int k = 0; k < Ytest.outerSize(); ++k) {
    int idx = Ytest.outerIndexPtr()[k];
    for (SparseMatrix<double>::InnerIterator it(Ytest,k); it; ++it) {
      coords(idx, 0) = it.row();
      coords(idx, 1) = it.col();
      coords(idx, 2) = it.value();
      idx++;
    }
  }
  return coords;
}

void MatrixDataCensored::setTrain(int* rows, int* cols, double* values, int nnz, int nrows, int ncols) {
	Y.resize(nrows, ncols);
	sparseFromIJV(Y, rows, cols, values, nnz);
	Yt = Y.transpose();
	mean_value = Y.sum() / Y.nonZeros();
	dims.resize(2);
	dims << Y.rows(), Y.cols();
}

void MatrixDataCensored::setTest(int* rows, int* cols, double* values, int nnz, int nrows, int ncols) {
	Ytest.resize(nrows, ncols);
	sparseFromIJV(Ytest, rows, cols, values, nnz);
}

void MatrixDataCensored::setTrain(int* columns, int nmodes, double* values, int nnz, int* d) {
  if (nmodes != 2) {
    throw std::runtime_error("MatrixData: tensor training input not supported.");
  }
  auto idx  = toMatrix(columns, nnz, nmodes);
  auto vals = toVector(values, nnz);

  Y.resize(d[0], d[1]);
  sparseFromIJV(Y, idx, vals);

  Yt = Y.transpose();
  mean_value = Y.sum() / Y.nonZeros();
  dims.resize(2);
  dims << Y.rows(), Y.cols();
}

void MatrixDataCensored::setTest(int* columns, int nmodes, double* values, int nnz, int* d) {
  if (nmodes != 2) {
    throw std::runtime_error("MatrixDataCensored: tensor training input not supported.");
  }
  auto idx  = toMatrix(columns, nnz, nmodes);
  auto vals = toVector(values, nnz);

  Ytest.resize(d[0], d[1]);
  sparseFromIJV(Ytest, idx, vals);
}

void MatrixDataCensored::setCensoring(int* columns, int nmodes, int* censored, int nnz, int* d) {
  if (nmodes != 2) {
    throw std::runtime_error("MatrixDataCensored: tensor training input not supported.");
  }
  auto idx  = toMatrix(columns, nnz, nmodes);
  auto cens = toVector(censored, nnz);

  C.resize(d[0],d[1]);
  sparseFromIJV(C,idx,cens); //fill matrix C
  Ct=C.transpose();

  std::cout <<"Censoring Enabled" << std::endl;
}

// util functions
Eigen::MatrixXi toMatrix(int* col1, int* col2, int nrows) {
  Eigen::MatrixXi idx(nrows, 2);
  for (int row = 0; row < nrows; row++) {
    idx(row, 0) = col1[row];
    idx(row, 1) = col2[row];
  }
  return idx;
}

Eigen::MatrixXi toMatrix(int* columns, int nrows, int ncols) {
  Eigen::MatrixXi idx(nrows, ncols);
  for (int row = 0; row < nrows; row++) {
    for (int col = 0; col < ncols; col++) {
      idx(row, col) = columns[col * nrows + row];
    }
  }
  return idx;
}

Eigen::VectorXd toVector(double* vals, int size) {
  Eigen::VectorXd v(size);
  for (int i = 0; i < size; i++) {
    v(i) = vals[i];
  }
  return v;
}

Eigen::VectorXi toVector(int* vals, int size) {
  Eigen::VectorXi v(size);
  for (int i = 0; i < size; i++) {
    v(i) = vals[i];
  }
  return v;
}


#include <eos/types/Asset.hpp>
#include <boost/rational.hpp>
#include <boost/multiprecision/cpp_int.hpp>

namespace EOS {
      typedef boost::multiprecision::int128_t  int128_t;

      uint8_t Asset::decimals()const {
         auto a = (const char*)&symbol;
         return a[0];
      }
      void Asset::set_decimals(uint8_t d){
         auto a = (char*)&symbol;
         a[0] = d;
      }

      std::string Asset::symbol_name()const {
         auto a = (const char*)&symbol;
         assert( a[7] == 0 );
         return &a[1];
      }

       int64_t Asset::precision()const {

         static int64_t table[] = {
                           1, 10, 100, 1000, 10000,
                           100000, 1000000, 10000000, 100000000ll,
                           1000000000ll, 10000000000ll,
                           100000000000ll, 1000000000000ll,
                           10000000000000ll, 100000000000000ll
                         };
         return table[ decimals() ];
      }

      string Asset::toString()const {
         string result = fc::to_string(amount.value / precision());
         if( decimals() )
         {
            auto fract = amount.value % precision();
            result += "." + fc::to_string(precision() + fract).erase(0,1);
         }
         return result + " " + symbol_name();
      }

      Asset Asset::fromString( const string& from )
      {
         try
         {
            string s = fc::trim( from );
            auto space_pos = s.find( " " );
            auto dot_pos = s.find( "." );

            Asset result;
            result.symbol = uint64_t(0);
            auto sy = (char*)&result.symbol;
            *sy = (char) dot_pos; // Mask due to undefined architecture behavior

            auto intpart = s.substr( 0, dot_pos );
            result.amount = fc::to_int64(intpart);
            std::string fractpart;
            if( dot_pos != std::string::npos )
            {
               auto fractpart = "1" + s.substr( dot_pos + 1, space_pos - dot_pos - 1 );
               result.set_decimals( fractpart.size() - 1 );

               result.amount.value *= result.precision();
               result.amount.value += fc::to_int64(fractpart);
               result.amount.value -= result.precision();
            }
            auto symbol = s.substr( space_pos + 1 );

            if( symbol.size() )
               memcpy( sy+1, symbol.c_str(), std::min(symbol.size(),size_t(6)) );

            return result;
         }
         FC_CAPTURE_LOG_AND_RETHROW( (from) )
      }

      bool operator == ( const Price& a, const Price& b )
      {
         if( std::tie( a.base.symbol, a.quote.symbol ) != std::tie( b.base.symbol, b.quote.symbol ) )
             return false;

         const auto amult = uint128_t( b.quote.amount.value ) * a.base.amount.value;
         const auto bmult = uint128_t( a.quote.amount.value ) * b.base.amount.value;

         return amult == bmult;
      }

      bool operator < ( const Price& a, const Price& b )
      {
         if( a.base.symbol < b.base.symbol ) return true;
         if( a.base.symbol > b.base.symbol ) return false;
         if( a.quote.symbol < b.quote.symbol ) return true;
         if( a.quote.symbol > b.quote.symbol ) return false;

         const auto amult = uint128_t( b.quote.amount.value ) * a.base.amount.value;
         const auto bmult = uint128_t( a.quote.amount.value ) * b.base.amount.value;

         return amult < bmult;
      }

      bool operator <= ( const Price& a, const Price& b )
      {
         return (a == b) || (a < b);
      }

      bool operator != ( const Price& a, const Price& b )
      {
         return !(a == b);
      }

      bool operator > ( const Price& a, const Price& b )
      {
         return !(a <= b);
      }

      bool operator >= ( const Price& a, const Price& b )
      {
         return !(a < b);
      }

      Asset operator * ( const Asset& a, const Price& b )
      {
         if( a.symbol_name() == b.base.symbol_name() )
         {
            FC_ASSERT( b.base.amount.value > 0 );
            uint128_t result = (uint128_t(a.amount.value) * b.quote.amount.value)/b.base.amount.value;
            FC_ASSERT( result.hi == 0 );
            return Asset( result.to_uint64(), b.quote.symbol );
         }
         else if( a.symbol_name() == b.quote.symbol_name() )
         {
            FC_ASSERT( b.quote.amount.value > 0 );
            uint128_t result = (uint128_t(a.amount.value) * b.base.amount.value)/b.quote.amount.value;
            FC_ASSERT( result.hi == 0 );
            return Asset( result.to_uint64(), b.base.symbol );
         }
         FC_THROW_EXCEPTION( fc::assert_exception, "invalid Asset * Price", ("Asset",a)("Price",b) );
      }

      Price operator / ( const Asset& base, const Asset& quote )
      { try {
         FC_ASSERT( base.symbol_name() != quote.symbol_name() );
         return Price{ base, quote };
      } FC_CAPTURE_AND_RETHROW( (base)(quote) ) }

      Price Price::max( SymbolType base, SymbolType quote ) { return Asset( ShareType(EOS_MAX_SHARE_SUPPLY), base ) / Asset( ShareType(1), quote); }
      Price Price::min( SymbolType base, SymbolType quote ) { return Asset( 1, base ) / Asset( EOS_MAX_SHARE_SUPPLY, quote); }

      bool Price::is_null() const { return *this == Price(); }

      void Price::validate() const
      { try {
         FC_ASSERT( base.amount > ShareType(0) );
         FC_ASSERT( quote.amount > ShareType(0) );
         FC_ASSERT( base.symbol_name() != quote.symbol_name() );
      } FC_CAPTURE_AND_RETHROW( (base)(quote) ) }

}  // EOS


#include "self_adjoint_angular_flux.h"
#include <deal.II/dofs/dof_accessor.h>

template <int dim>
SelfAdjointAngularFlux<dim>::SelfAdjointAngularFlux(
    const std::string equation_name,
    const dealii::ParameterHandler &prm,
    std::shared_ptr<FundamentalData<dim>> &data_ptr)
    : EquationBase<dim>(equation_name, prm, data_ptr) {}

/*
 * =============================================================================
 * PUBLIC FUNCTIONS
 * =============================================================================
 */

template <int dim>
void SelfAdjointAngularFlux<dim>::AssembleLinearForms (const int &g) {
  if (have_reflective_bc_) {
    for (int dir = 0; dir < n_dir_; ++dir) {
      global_angular_flux_[this->GetCompInd(g, dir)] =
          *(mat_vec_->sys_flxes[equ_name_][this->GetCompInd(g, dir)]);
    }
  }
  EquationBase<dim>::AssembleLinearForms(g);
}

template<int dim>
void SelfAdjointAngularFlux<dim>::IntegrateBoundaryBilinearForm (
      typename dealii::DoFHandler<dim>::active_cell_iterator &,
      const int &,
      dealii::FullMatrix<double> &cell_matrix,
      const int &,
      const int &dir) {

  // Get the boundary ID and the normal vector
  const dealii::Tensor<1, dim> normal_vector = fvf_->normal_vector(0);

  double normal_dot_omega = normal_vector * omega_[dir];

  if (normal_dot_omega > 0) {
    // Integrate bilinear term
    for (int q = 0; q < n_qf_; ++q) {
      for (int i = 0; i < dofs_per_cell_; ++i) {
        for (int j = 0; j < dofs_per_cell_; ++j) {
          cell_matrix(i,j) += normal_dot_omega *
                              fvf_->shape_value(i, q) *
                              fvf_->shape_value(j, q) *
                              fvf_->JxW(q);
        }
      }
    }
  }
}

template<int dim>
void SelfAdjointAngularFlux<dim>::IntegrateBoundaryLinearForm (
      typename dealii::DoFHandler<dim>::active_cell_iterator &cell,
      const int &fn,/*face number*/
      dealii::Vector<double> &cell_rhs,
      const int &g,
      const int &dir) {

  // Get boundary id to determine if we have reflective BCs
  int boundary_id = cell->face(fn)->boundary_id();

  if (have_reflective_bc_ && is_reflective_bc_.at(boundary_id)) {

    // Get normal vector and dot product with omega
    const dealii::Tensor<1, dim> normal_vector = fvf_->normal_vector(0);
    double normal_dot_omega = normal_vector * omega_[dir];

    if (normal_dot_omega < 0) {
      // Retrieve previous iteration angular flux and get local values for the
      // reflected angle
      int reflected_dir = this->ref_dir_ind_[std::make_pair(boundary_id, dir)];

      std::vector<double> angular_flux(n_qf_);
      fvf_->get_function_values(
          global_angular_flux_[this->GetCompInd(g, reflected_dir)],
          angular_flux);

      for (int q = 0; q < n_qf_; ++q) {
        for (int i = 0; i < dofs_per_cell_; ++i) {
          cell_rhs(i) -= normal_dot_omega *
                         fvf_->shape_value(i, q) *
                         angular_flux[q] *
                         fvf_->JxW(q);
        }
      }
    }
  }
}

template<int dim>
void SelfAdjointAngularFlux<dim>::IntegrateCellBilinearForm (
      typename dealii::DoFHandler<dim>::active_cell_iterator &cell,
      dealii::FullMatrix<double> &cell_matrix,
      const int &g,
      const int &dir) {
  // Get material id for the given cell and cross-sections
  int material_id = cell->material_id();
  auto sigma_t = xsec_->sigt.at(material_id)[g];
  auto inv_sigma_t = xsec_->inv_sigt.at(material_id)[g];

  // Integrate and add both bilinear terms using precomputed values
  for (int q = 0; q < n_q_; ++q) {
    for (int i = 0; i < dofs_per_cell_; ++i) {
      for (int j = 0; j < dofs_per_cell_; ++j) {

        cell_matrix(i, j) += (pre_streaming_[{dir, q}](i, j) * inv_sigma_t
                              +
                              pre_collision_[q](i, j) * sigma_t) * fv_->JxW(q);
      }
    }
  }
}

template<int dim>
void SelfAdjointAngularFlux<dim>::IntegrateCellFixedLinearForm (
    typename dealii::DoFHandler<dim>::active_cell_iterator &cell,
    dealii::Vector<double> &cell_rhs,
    const int &g,
    const int &dir) {
  // get material id for the given cell
  int material_id = cell->material_id();

  // SAAF has two fixed source terms, the first is proportional to q/4pi, the
  // second contains an additional division by sigma_t. These two vectors hold
  // these values, respectively, for each quadrature point.
  std::vector<double> cell_q(n_q_);
  std::vector<double> cell_q_over_total(n_q_);

  if (!is_eigen_problem_) {

    // Fill the two vectors with the appropriate q/4pi and q/4pi*sigma_t values
    auto q_per_ster = xsec_->q_per_ster.at(material_id)[g];
    std::fill(cell_q.begin(), cell_q.end(),
              q_per_ster);
    std::fill(cell_q_over_total.begin(), cell_q_over_total.end(),
              q_per_ster * xsec_->inv_sigt.at(material_id)[g]);

  } else if (xsec_->is_material_fissile.at(material_id)) {

    // Fill the two vectors with the appropriate fission sources, summed
    // over all groups
    for (int group_in = 0; group_in < n_group_; ++group_in) {
      // Retrieve cell scalar flux, inverse sigma t, fission terms
      std::vector<double> group_cell_scalar_flux(n_q_);
      this->GetGroupCellScalarFlux(group_cell_scalar_flux, group_in);
      auto inv_sigma_t = xsec_->inv_sigt.at(material_id)[g];
      auto scaled_fission_transfer =
          scaled_fiss_transfer_.at(material_id)(group_in, g);

      for (int q = 0; q < n_q_; ++q) {
        cell_q[q] +=
            scaled_fission_transfer * group_cell_scalar_flux[q];
        cell_q_over_total[q] +=
            cell_q[q] * inv_sigma_t;
      }
    }
  }

  //Integrate and add both source terms
  for (int q = 0; q < n_q_; ++q) {

    cell_q[q] *= fv_->JxW(q);
    cell_q_over_total[q] *= fv_->JxW(q);

    for (int i = 0; i < dofs_per_cell_; ++i) {
      // First scattering term
      cell_rhs(i) += fv_->shape_value(i, q) * cell_q[q];
      // Second scattering term
      cell_rhs(i) +=
          omega_[dir] * fv_->shape_grad(i, q) * cell_q_over_total[q];
    }
  }
}

template<int dim>
void SelfAdjointAngularFlux<dim>::IntegrateScatteringLinearForm (
      typename dealii::DoFHandler<dim>::active_cell_iterator &cell,
      dealii::Vector<double> &cell_rhs,
      const int &g,
      const int &dir) {
  // Get material id for the given cell:
  int material_id = cell->material_id();

  // SAAF has two scattering terms, the first is proportional to scalar flux
  // time sigma_s over 4pi, the second has an additional division by sigma_t.
  // These two vectors hold those values, respectively,  at each quadrature
  // point:
  std::vector<double> cell_scatter_flux(n_q_);
  std::vector<double> cell_scatter_over_total_flux(n_q_);

  // Iterate over groups to populate cell_scatter_flux
  for (int group_in = 0; group_in < n_group_; ++group_in) {
    std::vector<double> group_cell_scalar_flux(n_q_);
    this->GetGroupCellScalarFlux(group_cell_scalar_flux, group_in);

    // Get needed cross-sections:
    auto sigma_s_per_ster = xsec_->sigs_per_ster.at(material_id)(group_in, g);
    auto inv_sigma_t = xsec_->inv_sigt.at(material_id)[g];

    // Fold group cell scalar flux into total cell scalar flux and multiply by
    // appropriate cross-sections:
    for (int q = 0; q < n_q_; ++q) {
      cell_scatter_flux[q] += sigma_s_per_ster * group_cell_scalar_flux[q];
      cell_scatter_over_total_flux[q] = cell_scatter_flux[q] * inv_sigma_t;
    }
  }

  // Integrate and add both scattering terms
  for (int q = 0; q < n_q_; ++q) {
    cell_scatter_flux[q] *= fv_->JxW(q);
    cell_scatter_over_total_flux[q] *= fv_->JxW(q);
    for (int i = 0; i < dofs_per_cell_; ++i) {
      // First scattering term
      cell_rhs(i) += fv_->shape_value(i, q) * cell_scatter_flux[q];
      // Second scattering term
      cell_rhs(i) +=
          omega_[dir] * fv_->shape_grad(i, q) * cell_scatter_over_total_flux[q];
    }
  }
}

template <int dim>
void SelfAdjointAngularFlux<dim>::PreassembleCellMatrices () {
  // Reinitialize FEM values to an arbitrary cell (the first one) in the list of
  // local cells.
  fv_->reinit(dat_ptr_->local_cells[0]);

  // For each quadrature angle, generate the Collision and Streaming matrices
  for (int q = 0; q < n_q_; ++q) {

    dealii::FullMatrix<double> temp_matrix(dofs_per_cell_, dofs_per_cell_);

    for (int i = 0; i < dofs_per_cell_; ++i) {
      for (int j = 0; j < dofs_per_cell_; ++j) {
        temp_matrix(i,j) =
            (fv_->shape_value(i,q) * fv_->shape_value(j,q));
      }
    }

    pre_collision_[q] = temp_matrix;

    temp_matrix = 0;

    // Streaming terms also depend on direction
    for (int dir = 0; dir < n_dir_; ++dir) {
      for (int i = 0; i < dofs_per_cell_; ++i) {
        for (int j = 0; j < dofs_per_cell_; ++j) {
          temp_matrix(i,j) =
              (fv_->shape_grad(i,q) * omega_[dir])
              *
              (fv_->shape_grad(j,q) * omega_[dir]);
        }
      }
      pre_streaming_[{dir, q}] = temp_matrix;
    }
  }
}

/*
 * =============================================================================
 * PROTECTED FUNCTIONS
 * =============================================================================
 */

template <int dim>
void SelfAdjointAngularFlux<dim>::GetGroupCellScalarFlux
(std::vector<double> &to_fill, int group) {

  // Get the global scalar flux for the current group
  auto & group_global_scalar_flux =
      mat_vec_->moments[equ_name_][std::make_tuple(group, 0, 0)];
  // Evaluate the global scalar flux at the quadrature points of the current
  // cell and store in return_vector (dealii function in FEValuesBase)
  fv_->get_function_values(group_global_scalar_flux, to_fill);

}

template class SelfAdjointAngularFlux<1>;
template class SelfAdjointAngularFlux<2>;
template class SelfAdjointAngularFlux<3>;