// Copyright 2022 Haruki Uchiito
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to
// deal in the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.

// clang-format off
#include <gperftools/profiler.h>

#include <tf2/exceptions.h>
#include <tf2_ros/buffer.h>
#include <tf2_ros/transform_listener.h>

#include <ndt_2d_slam/data_types.hpp>
#include <ndt_2d_slam/scan_matcher_2d.hpp>
#include <ndt_2d_slam/conversion.hpp>
#include <ros2_utility/ros2_utility.hpp>

#include <rclcpp/rclcpp.hpp>
#include <laser_geometry/laser_geometry.hpp>

#include <std_msgs/msg/float32.hpp>
#include <std_msgs/msg/int32.hpp>
#include <geometry_msgs/msg/twist.hpp>
#include <sensor_msgs/msg/laser_scan.hpp>
#include <sensor_msgs/msg/point_cloud2.hpp>
#include <visualization_msgs/msg/marker_array.hpp>

#include <Eigen/Eigenvalues>

#include <algorithm>
#include <memory>
#include <vector>
#include <string>
#include <mutex>
#include <sstream>
// clang-format on

using std::placeholders::_1;

visualization_msgs::msg::Marker errorEllipse2D(
    const std::string& frame_id, const rclcpp::Time& timestamp, int32_t id,
    const Eigen::Vector2d& mean, const NDT2DSLAM::CovarianceMat& covariance) {
    // calc. angle of ellipse
    Eigen::SelfAdjointEigenSolver<Eigen::Matrix2d> solver(covariance);
    Eigen::Vector2d evs = solver.eigenvalues();
    std::vector<double> ev(evs.data(), evs.data() + 2);
    int idx = std::max_element(ev.begin(), ev.end()) - ev.begin();
    Eigen::Vector2d largerEigenVector = solver.eigenvectors().col(idx);
    double yaw = atan2(largerEigenVector(1), largerEigenVector(0));

    // calc. size of sphere
    const double x2 = 9.35;
    double l_len = sqrt(x2 * evs(idx));
    double s_len = sqrt(x2 * evs(static_cast<int>(!static_cast<bool>(idx))));

    return ROS2Utility::marker(
        frame_id, timestamp, id, visualization_msgs::msg::Marker::SPHERE,
        ROS2Utility::pose(mean(0), mean(1), 0.0, 0.0, 0.0, yaw),
        ROS2Utility::vector3(l_len, s_len, 0.0),
        ROS2Utility::colorrgba(1.0, 0.0, 0.0, 0.5));
}

class NDT2DSLAMNode : public rclcpp::Node {
   public:
    NDT2DSLAMNode() : Node("ndt_2d_slam") {
        _tfBuffer = std::make_unique<tf2_ros::Buffer>(this->get_clock());
        _tfListener = std::make_shared<tf2_ros::TransformListener>(*_tfBuffer);

        callback_group_subscriber1_ = this->create_callback_group(
            rclcpp::CallbackGroupType::MutuallyExclusive);
        _odomCallbackGroup = this->create_callback_group(
            rclcpp::CallbackGroupType::MutuallyExclusive);

        auto sub1_opt = rclcpp::SubscriptionOptions();
        sub1_opt.callback_group = callback_group_subscriber1_;

        _laserSub = this->create_subscription<sensor_msgs::msg::LaserScan>(
            "scan", rclcpp::QoS(10),
            std::bind(&NDT2DSLAMNode::laserCallback, this,
                      std::placeholders::_1),
            sub1_opt);

        _markerPublisher =
            this->create_publisher<visualization_msgs::msg::MarkerArray>(
                "ndtcov", 10);
        _markerPoints =
            ROS2Utility::MarkerPoints("odom", this->get_clock()->now());
        _markerPointsPublisher =
            this->create_publisher<visualization_msgs::msg::Marker>(
                "ndtTrjectory", 10);

        _scanCntPublisher =
            this->create_publisher<std_msgs::msg::Int32>("scanCnt", 10);
        _scorePublisher =
            this->create_publisher<std_msgs::msg::Float32>("score", 10);
        _mapPublisher =
            this->create_publisher<sensor_msgs::msg::PointCloud2>("ndtmap", 10);
        _lastScanPublisher =
            this->create_publisher<sensor_msgs::msg::PointCloud2>("cscan", 10);
        _currentScanPublisher =
            this->create_publisher<sensor_msgs::msg::PointCloud2>("lscan", 10);

        using namespace std::chrono_literals;
        _laserProcessTimer = this->create_wall_timer(
            100ms, std::bind(&NDT2DSLAMNode::laserProcessTimerCallback, this));

        // register odom retrieval callback
        if (false) {
            auto odom_opt = rclcpp::SubscriptionOptions();
            odom_opt.callback_group = _odomCallbackGroup;
            _twistSub = this->create_subscription<geometry_msgs::msg::Twist>(
                "cmd_vel", rclcpp::QoS(10),
                std::bind(&NDT2DSLAMNode::twistCallback, this,
                          std::placeholders::_1),
                odom_opt);
        } else {
            _odomTFTimer = this->create_wall_timer(
                50ms, std::bind(&NDT2DSLAMNode::odomTFTimerCallback, this),
                _odomCallbackGroup);
        }
    }

   private:
    std::ofstream logOut{
        "/home/xoke/ros2_ws/src/ndt_2d_slam/test/output/nodeOut.txt"};
    // NDT2DSLAM::Transform2D _estimate{0.0, 0.0, 0.0};
    ROS2Utility::MarkerPoints _markerPoints;
    rclcpp::Publisher<visualization_msgs::msg::Marker>::SharedPtr
        _markerPointsPublisher;

    std::shared_ptr<tf2_ros::TransformListener> _tfListener{nullptr};
    std::unique_ptr<tf2_ros::Buffer> _tfBuffer;

    std::mutex _slamMutex;
    NDTSLAM slam{0.3, logOut};
    rclcpp::Publisher<std_msgs::msg::Int32>::SharedPtr _scanCntPublisher;
    rclcpp::Publisher<std_msgs::msg::Float32>::SharedPtr _scorePublisher;
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr _mapPublisher;
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr
        _lastScanPublisher,
        _currentScanPublisher;
    void laserProcessTimerCallback() {
        std::stringstream frameOut;
        frameOut << "[at timer]" << std::endl;
        // ProfilerStart("profile.prof");

        if (_currentScan.size() == 0) return;
        std::cerr << "map size: " << slam.map().size() << std::endl;

        rclcpp::Time st = this->get_clock()->now();

        double score;
        {
            std::lock_guard<std::mutex> lock(_slamMutex);
            frameOut << "slam process done" << std::endl;
            score = slam.process(_currentScan);
            frameOut << "done" << std::endl;
        }
        std_msgs::msg::Float32 pscore;
        pscore.data = score;
        _scorePublisher->publish(pscore);

        std_msgs::msg::Int32 pcnt;
        pcnt.data = _scanCnt;
        _scanCntPublisher->publish(pcnt);

        std::cout << slam.estimate() << std::endl;

        double dt = (this->get_clock()->now() - st).nanoseconds();
        frameOut << "timer time: " << dt * 1e-6 << "[ms]" << std::endl;

        NDT2DSLAM::Scan2D mp = slam.map();
        if (mp.size() > 0) {
            frameOut << "got map" << std::endl;
            auto pc2 = NDT2DSLAM::scan2dToROSPointCloud2(
                mp, "odom", this->get_clock()->now());
            frameOut << "pub pc width: " << pc2.width << std::endl;
            _mapPublisher->publish(pc2);
            frameOut << "map published" << std::endl;
        } else {
            std::cerr << "slam pub skipped" << std::endl;
        }

        _markerPoints.addPoint(
            ROS2Utility::point(slam.estimate().tx, slam.estimate().ty, 0.0),
            ROS2Utility::colorrgba(0.0, 1.0, 0.0, 0.6));
        _markerPointsPublisher->publish(_markerPoints.instance);
        frameOut << "trj published" << std::endl;

        // _lastScanPublisher->publish(NDT2DSLAM::scan2dToROSPointCloud2(
        // _lastScan, "compare", this->get_clock()->now()));
        // _currentScanPublisher->publish(NDT2DSLAM::scan2dToROSPointCloud2(
        // _currentScan.transformBy(slam.estimate()), "compare",
        //  this->get_clock()->now()));

        _lastScan = _currentScan;
        _currentScan = NDT2DSLAM::Scan2D();
        // ProfilerStop();
        RCLCPP_INFO(this->get_logger(), frameOut.str().c_str());
    }

    // odometry handler objects and callbacks
    rclcpp::Subscription<geometry_msgs::msg::Twist>::SharedPtr _twistSub;
    rclcpp::Time _lastTime;
    void twistCallback(const geometry_msgs::msg::Twist::SharedPtr msg) {
        RCLCPP_INFO(this->get_logger(), "[twist callback]");

        if (_lastTime.seconds() < 1e-9) {
            _lastTime = this->get_clock()->now();
            return;
        }

        rclcpp::Time ct = this->get_clock()->now();
        double dt = (ct - _lastTime).seconds();
        RCLCPP_INFO(this->get_logger(), "dt: %f", dt);

        std::lock_guard<std::mutex> lock(_slamMutex);
        slam.motionUpdate(dt, msg->linear.x, msg->angular.z);
        RCLCPP_INFO(this->get_logger(), "pose x: %f, y: %f, th: %f",
                    slam.estimate().tx, slam.estimate().ty,
                    slam.estimate().theta);

        _lastTime = ct;
    }

    rclcpp::TimerBase::SharedPtr _odomTFTimer;
    geometry_msgs::msg::Transform _lastOdom;
    void odomTFTimerCallback() {
        std::string fromFrameRel = "base_link";
        std::string toFrameRel = "odom";
        geometry_msgs::msg::TransformStamped transformStamped;
        try {
            transformStamped = _tfBuffer->lookupTransform(
                toFrameRel, fromFrameRel, tf2::TimePointZero);
        } catch (tf2::TransformException& ex) {
            RCLCPP_INFO(this->get_logger(), "Could not transform %s to %s: %s",
                        toFrameRel.c_str(), fromFrameRel.c_str(), ex.what());
            return;
        }

        auto xyz = transformStamped.transform.translation;
        auto rpy = ROS2Utility::rpy(transformStamped.transform.rotation);
        auto lastRPY = ROS2Utility::rpy(_lastOdom.rotation);

        double dx = xyz.x - _lastOdom.translation.x;
        double dy = xyz.y - _lastOdom.translation.y;
        double dth = std::get<2>(rpy) - std::get<2>(lastRPY);

        std::lock_guard<std::mutex> lock(_slamMutex);
        slam.directMotionUpdate(dx, dy, dth);

        _lastOdom = transformStamped.transform;
    }

    // laser callback objects
    rclcpp::Subscription<sensor_msgs::msg::LaserScan>::SharedPtr _laserSub;
    laser_geometry::LaserProjection _projector;
    NDT2DSLAM::Scan2D _lastScan, _currentScan;
    int32_t _scanCnt{0};
    rclcpp::TimerBase::SharedPtr _laserProcessTimer;
    void laserCallback(const sensor_msgs::msg::LaserScan::SharedPtr msg) {
        // RCLCPP_INFO(this->get_logger(), "[laser callback]");
        sensor_msgs::msg::PointCloud2 cloud;
        try {
            _projector.transformLaserScanToPointCloud("base_link", *msg, cloud,
                                                      *_tfBuffer);

            _currentScan = NDT2DSLAM::rosPointCloud2ToScan2D(cloud);
            if (_lastScan.size() == 0) _lastScan = _currentScan;
            _scanCnt++;

            /*
                        NDT2DSLAM::NDTScanMatcher ndt(1.0, 1e-4, 1e-4,
            _lastScan); visualization_msgs::msg::MarkerArray ma; int id = 0; for
            (auto& cell : ndt._cells) { if (!cell.validity) continue; auto ell =
            errorEllipse2D("laser", this->get_clock()->now(),
                                                      ++id, cell.mean,
            cell.covariance); ma.markers.push_back(ell);
                        }
                        _markerPublisher->publish(ma);
            //            */
        } catch (tf2::TransformException& ex) {
            RCLCPP_ERROR(this->get_logger(), "got exeption");
        }
    }

    rclcpp::CallbackGroup::SharedPtr callback_group_subscriber1_;
    rclcpp::CallbackGroup::SharedPtr _odomCallbackGroup;

    rclcpp::Publisher<visualization_msgs::msg::MarkerArray>::SharedPtr
        _markerPublisher;
};

int main(int argc, char* argv[]) {
    rclcpp::init(argc, argv);

    rclcpp::executors::MultiThreadedExecutor executor;
    auto node = std::make_shared<NDT2DSLAMNode>();
    executor.add_node(node);
    executor.spin();

    rclcpp::shutdown();
    return 0;
}


#include <cstdlib>
#include <dlib/dnn.h>

namespace chrono = std::chrono;
using fus = chrono::duration<float, std::micro>;

template <typename fwd_func, typename bwd_func> auto benchmark(
    fwd_func fwd,
    bwd_func bwd,
    const std::string& name,
    int iterations = 100,
    int warmup = 10)
{
    dlib::resizable_tensor input, output, gradient;
    chrono::time_point<chrono::steady_clock> t0, t1;
    input.set_size(16, 10, 256, 256);
    output.copy_size(input);
    gradient.copy_size(input);
    t0 = chrono::steady_clock::now();
    dlib::running_stats<float> fwd_stats, bwd_stats;
    for (int i = 0; i < warmup; ++i)
    {
        fwd(output, input);
    }
    t1 = chrono::steady_clock::now();

    for (int i = 0; i < iterations; ++i)
    {
        t0 = chrono::steady_clock::now();
        fwd(output, input);
        t1 = chrono::steady_clock::now();
        fwd_stats.add(chrono::duration_cast<fus>(t1 - t0).count());
    }
    std::cout << name << " fwd: " << fwd_stats.mean() << " ± " << fwd_stats.stddev() << " µs\n";

    for (int i = 0; i < warmup; ++i)
    {
        bwd(gradient, output, input);
    }

    for (int i = 0; i < iterations; ++i)
    {
        t0 = chrono::steady_clock::now();
        bwd(gradient, output, input);
        t1 = chrono::steady_clock::now();
        bwd_stats.add(chrono::duration_cast<fus>(t1 - t0).count());
    }
    std::cout << name << " bwd: " << bwd_stats.mean() << " ± " << bwd_stats.stddev() << " µs\n";
}

auto main(/*const int argc, const char** argv*/) -> int
try
{
    setenv("CUDA_LAUNCH_BLOCKING", "1", 1);
    const int warmup = 10;
    const int iterations = 100;

    benchmark(dlib::tt::relu, dlib::tt::relu_gradient, "relu", warmup, iterations);
    benchmark(dlib::tt::mish, dlib::tt::mish_gradient, "mish", warmup, iterations);
}
catch (const std::exception& e)
{
    std::cout << e.what() << '\n';
}


//
// ********************************************************************
// * License and Disclaimer                                           *
// *                                                                  *
// * The  Geant4 software  is  copyright of the Copyright Holders  of *
// * the Geant4 Collaboration.  It is provided  under  the terms  and *
// * conditions of the Geant4 Software License,  included in the file *
// * LICENSE and available at  http://cern.ch/geant4/license .  These *
// * include a list of copyright holders.                             *
// *                                                                  *
// * Neither the authors of this software system, nor their employing *
// * institutes,nor the agencies providing financial support for this *
// * work  make  any representation or  warranty, express or implied, *
// * regarding  this  software system or assume any liability for its *
// * use.  Please see the license in the file  LICENSE  and URL above *
// * for the full disclaimer and the limitation of liability.         *
// *                                                                  *
// * This  code  implementation is the result of  the  scientific and *
// * technical work of the GEANT4 collaboration.                      *
// * By using,  copying,  modifying or  distributing the software (or *
// * any work based  on the software)  you  agree  to acknowledge its *
// * use  in  resulting  scientific  publications,  and indicate your *
// * acceptance of all terms of the Geant4 Software license.          *
// ********************************************************************
//
// $Id: PreampMedipix.cc $
//
/// \file src/PreampMedipix.cc
/// \brief Implementation of the PreampMedipix class

#include "PreampMedipix.hh"
#include "G4RunManager.hh"
#include "G4Run.hh"

#include <stdio.h>
#include <iostream>
#include <fstream>

#include <boost/concept_check.hpp>

#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/ini_parser.hpp>

using namespace std;

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
PreampMedipix::PreampMedipix(boost::property_tree::ptree pt)
{

    G4RunManager *fRM = G4RunManager::GetRunManager();
    myDet = (DetectorConstructionBase *)(fRM->GetUserDetectorConstruction());
    useCSM = myDet->GetCsmMode();

    G4int detectorType = myDet->GetDetectorType();
    // set gain modes for Medipix3RX and preamp
    if(detectorType == 0){
        if(useCSM == true){
            elecSigma = pt.get<G4double>("chip_medipix3rx.csm_enc");
        }else {
            elecSigma = pt.get<G4double>("chip_medipix3rx.spm_enc");
        }
        Cf = myDet->GetDetectorGain();
    } else if(detectorType == 1) { //Timepix
        elecSigma = pt.get<G4double>("chip_timepix.enc");
        Cf = 8e-15; // 8fF same as Medipix2MXR
    } else if(detectorType == 2) { // Dosepix
        elecSigma = pt.get<G4double>("chip_dosepix.enc");
        Cf = 8e-15;
    } else if (detectorType == 3) { // Timepix3
        elecSigma = pt.get<G4double>("chip_timepix3.enc");
        Cf = 3e-15; // 3fF
    }

    // get material and load material properties from ini file
    G4String material =  myDet->GetSensorMaterial()->GetName();
    G4String sensor;
    if (material == "G4_Si")
        sensor = "sensor_silicon";
    else if (material == "G4_CADMIUM_TELLURIDE")
        sensor = "sensor_cdte";

    pulsePrecision =    pt.get<G4double>("computation.pulsePrecision") * ns;  //must be ns!
    maxPulseTime =      pt.get<G4double>("computation.maxPulseTime") * ns;

    nPulseArrayElements = (G4int) maxPulseTime / pulsePrecision;
    ampResponseTime =   pt.get<G4double>("computation.ampResponseTime") * ns;
    nAmpResponseElements = (G4int)(ampResponseTime / pulsePrecision);

    nElectronHolePairs = pt.get<G4double>(sensor + ".nElectronHolePairs");
    minEnergy = 0.*keV;

    DEBUG = pt.get<G4bool>("testing.debug");
    writePeakToFile = pt.get<G4bool>("testing.writePeakToFile");

    preampType = pt.get<G4int>("computation.preampType");

    Ikrum = pt.get<G4double>("chip.Ikrum") * 1e-9;

    detector = MpxDetector::GetInstance();
    thresholdkeV = detector->GetTpxThreshold();
    thresholdCharge = thresholdkeV * 1000 / nElectronHolePairs ;

    //set the transfer functions
    SetTransferFunctions();

    //THL dispersion
    nPixel = myDet->GetNbPixels();
    TpxThlDisp = new G4double[nPixel * nPixel]();

    SetThresholdDispersion();
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::SetThresholdDispersion()
{
    //load .ini file with configuration data
    boost::property_tree::ptree pt;
    boost::property_tree::ini_parser::read_ini("DetectorConfig.ini", pt);

    //sensor properties
    G4String material =  myDet->GetSensorMaterial()->GetName();
    G4String sensor;
    if (material == "G4_Si")
        sensor = "sensor_silicon";
    else if (material == "G4_CADMIUM_TELLURIDE")
    {
        sensor = "sensor_cdte";
    }

    G4int detectorType = myDet->GetDetectorType();

    G4double disp = 0;

    if (detectorType == 0){ //Medipix3RX
        if(useCSM == true){
            disp = pt.get<G4double>("chip_medipix3rx.csm_threshold_dispersion");
        } else {
            disp = pt.get<G4double>("chip_medipix3rx.spm_threshold_dispersion");
        }

    }else if (detectorType == 1){ //Timepix1
        disp = pt.get<G4double>("chip_timepix.threshold_dispersion");

    }else if(detectorType == 2){ //Dosepix
        disp = pt.get<G4double>("chip_dosepix.threshold_dispersion");

    }else if (detectorType == 3) { // Timepix3
        disp = pt.get<G4double>("chip_timepix3.threshold_dispersion");
    }

    for (G4int i = 0; i < nPixel * nPixel; i++) {
        //G4cout << "DEBUG: Setting Threshold Dispersion " << disp << " i: " << i << " nPixel: " << nPixel*nPixel << G4endl;
        TpxThlDisp[i] = CLHEP::RandGauss::shoot(0, disp); //sigma of thl dispersion
    }
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::SetTransferFunctions(){

    //Preamp transfer function
    G4double Cd = 25e-15;  //Detector capacitance

    G4double Gmfb   = Ikrum/2.*20;  //Transconductance in the feedback loop (approximation from Weak Inversion)
    G4double Rf     = 2./Gmfb;       //feedback resistor

    G4double gm0=30e-6;  //transconductance parameter
    //G4double R0=20e6;    //Output resistance amplifier
    G4double C0=25e-15;  //Output capacitance amplifier
    G4double Ct=Cd*Cf+Cd*C0+Cf*C0;

    G4double rcpr=Ct/(gm0*Cf);
    G4double rcpf=Cf*Rf;

    G4double w1 = 1./rcpr;
    G4double w2 = 1./rcpf;

    //create transfer function:
    nImpResEl = nAmpResponseElements;
    impRes = new G4double[nImpResEl];
    for (G4int i = 0; i < nImpResEl; i++) {
        impRes[i] = exp(-i* pulsePrecision * 1e-9 *w2 - exp(-i* pulsePrecision * 1e-9* w1));
        //G4cout << "Debug1: " << i << " value: " << impRes[i] << G4endl;
    }
    //Shaper transfer function
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::GetPixelResponse(map<pair<G4int, G4int>, G4double *> *inducedPixelContent, MpxDigitCollection* digitCollection, G4int event)
{
    thresholdkeV = detector->GetTpxThreshold();
    thresholdCharge = thresholdkeV * 1000 / nElectronHolePairs ;

    if(preampType == 0){
        chargeIntegrationPreamp(inducedPixelContent, event, digitCollection);
    } else if(preampType == 1){
        convolutionPreamp(inducedPixelContent, event, digitCollection);
    }
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::chargeIntegrationPreamp(map<pair<G4int, G4int>, G4double *> *inducedPixelContent, G4int event, MpxDigitCollection *digitCollection)
{
    map<pair<G4int, G4int>, G4double> pixelsContent;

    map<pair<G4int, G4int>, G4double * >::iterator iPCItr;
    iPCItr = inducedPixelContent->begin();
    for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {

        G4double maxCharge = 0;
        G4double *tempArray = (*iPCItr).second;
        G4double *ampResponse = new G4double[nAmpResponseElements]();

        //preamp response
        if(preampType == 0){
            for (G4int i = 0; i < nPulseArrayElements; i++) {
                maxCharge += tempArray[i];
            }
        }

        //convert maximum in preamp response back to energy //FIXME here TOT
        G4double energy = ElectronicsNoise(maxCharge) * nElectronHolePairs * eV;

        //safe new pixel energy, not yet Digit(SPM and CSM)
        if (energy > minEnergy) {
            std::pair<G4int, G4int> newPixel;
            newPixel.first  = (*iPCItr).first.first;
            newPixel.second = (*iPCItr).first.second;
            pixelsContent[newPixel] = energy;
        }
        //write preamp response to file
        if (writePeakToFile == true) {

            G4int k = (*iPCItr).first.first * 1000 + (*iPCItr).first.second;
            char filename[64];
            std::sprintf(filename, "preampSingle%d", k);

            writeToFile(filename, ampResponse, nAmpResponseElements);
        }

    }

    //print the energy in every pixel to the console
    if (DEBUG) {
        std::map<std::pair<G4int, G4int>, G4double>::iterator testItr =  pixelsContent.begin();
        G4cout << pixelsContent.size() << " number of entries in pixel Content" << G4endl;

        for (; testItr != pixelsContent.end() ; testItr++) {
            G4cout << (*testItr).first.first << " " << (*testItr).first.second << " " << (*testItr).second / keV << G4endl;
        }
        G4cout << G4endl;
    }

    //create digits. one per pixel
    std::map<std::pair<G4int, G4int>, G4double>::iterator pCItr =  pixelsContent.begin();
    std::map<std::pair<G4int, G4int>, G4double>::iterator tempPCItr =  pixelsContent.begin();

    //charge summing mode
    if (useCSM == true) {

        map<pair<G4int, G4int>, G4double> csmSumContent;
        std::pair<G4int, G4int> tempPixelCSM;
        std::pair<G4int, G4int> iteratePixel;

        //calculate the 2x2 sums and store in csmSumContent, equivalent of the sum/pixel in the chip
        for (; pCItr != pixelsContent.end() ; pCItr++) {

            G4double sumEnergy = 0;

            //get tempPixel
            tempPixelCSM = pCItr->first;
            iteratePixel = tempPixelCSM;

            //sumEnergy of pixel
            if (pixelsContent[tempPixelCSM] > 0) sumEnergy = pixelsContent[tempPixelCSM];

            iteratePixel.first += 1;
            tempPCItr = pixelsContent.find(iteratePixel);
            if (tempPCItr != pixelsContent.end()) {
                if (tempPCItr->second > 0) sumEnergy += tempPCItr->second;
            }

            iteratePixel.second -= 1;
            tempPCItr = pixelsContent.find(iteratePixel);
            if (tempPCItr != pixelsContent.end()) {
                if (tempPCItr->second > 0) sumEnergy += tempPCItr->second;
            }

            iteratePixel.first -= 1;
            tempPCItr = pixelsContent.find(iteratePixel);
            if (tempPCItr != pixelsContent.end()) {
                if (tempPCItr->second > 0) sumEnergy += tempPCItr->second;
            }

            //store the sum with the pixel ID
            csmSumContent[tempPixelCSM] = sumEnergy;
        }

        //compare the sums
        std::map<std::pair<G4int, G4int>, G4double>::iterator csmItr =  csmSumContent.begin();
        for (; csmItr != csmSumContent.end() ; csmItr++) {

            //get tempPixel
            tempPixelCSM = csmItr->first;
            iteratePixel = tempPixelCSM;

            //neighbors
            G4double cmpEnergy1 = 0;
            G4double cmpEnergy2 = 0;
            G4double cmpEnergy3 = 0;
            G4double cmpEnergy4 = 0;

            //diagonal
            G4double cmpEnergy5 = 0;
            G4double cmpEnergy6 = 0;
            G4double cmpEnergy7 = 0;
            G4double cmpEnergy8 = 0;

            //get sumEnergy of pixel
            G4double pixelEnergy = csmSumContent[tempPixelCSM];

            iteratePixel = tempPixelCSM;
            iteratePixel.first += 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy1 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first -= 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy2 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.second += 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy3 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.second -= 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy4 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first -= 1;
            iteratePixel.second -= 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy5 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first += 1;
            iteratePixel.second -= 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy6 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first -= 1;
            iteratePixel.second += 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy7 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first += 1;
            iteratePixel.second += 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy8 += tempPCItr->second;
            }

            //create digit
            if (pixelEnergy > cmpEnergy1 && pixelEnergy > cmpEnergy2 && pixelEnergy > cmpEnergy3 && pixelEnergy > cmpEnergy4 && pixelEnergy > cmpEnergy5 && pixelEnergy > cmpEnergy6 && pixelEnergy > cmpEnergy7 && pixelEnergy > cmpEnergy8) {
                Digit *digit = new Digit;
                digit->SetColumn(tempPixelCSM.first);
                digit->SetLine(tempPixelCSM.second);
                digit->SetEnergy(pixelEnergy / keV);
                digit->SetEvent(event);
                digitCollection->insert(digit);
            }
        }

    } else {  //no CSM
        for (; pCItr != pixelsContent.end() ; pCItr++) {
            //convert all events into digits
            G4double energy = (*pCItr).second;
            if (energy > minEnergy) {
                Digit *digit = new Digit;
                digit->SetColumn((*pCItr).first.first);
                digit->SetLine((*pCItr).first.second);
                digit->SetEnergy((energy / keV));
                digit->SetEvent(event);
                digitCollection->insert(digit);
            }
        }
    }

    //free memory
    iPCItr = inducedPixelContent->begin();
    for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {
        G4double *tempArray = (*iPCItr).second;
        delete[] tempArray;
    }
    inducedPixelContent->clear();
    delete inducedPixelContent;
    pixelsContent.clear();

}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::convolutionPreamp(map<pair<G4int, G4int>, G4double *> *inducedPixelContent, G4int event, MpxDigitCollection *digitCollection)
{
    //SPM mode
    if(useCSM == false){

        map<pair<G4int, G4int>, G4double * >::iterator iPCItr;
        iPCItr = inducedPixelContent->begin();

        //iterate over induced currents in pixels
        for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {

            G4double *ampResponse = new G4double[nAmpResponseElements]();
            G4double *inducedChargeArray = (*iPCItr).second;

            //Preamp
            convolve(inducedChargeArray, ampResponse, nPulseArrayElements, impRes, nImpResEl);

            //noise
            for(G4int i=0; i<nAmpResponseElements; i++){
                ampResponse[i] = ElectronicsNoise(ampResponse[i]);
            }

            //getEnergy
            boost::tuple<G4double,G4double,G4double> totPeakToa;
            G4int pixelIndex = ((*iPCItr).first.first) * nPixel + ((*iPCItr).first.second);
            G4double thlDisp = TpxThlDisp[pixelIndex];

            totPeakToa = getToTandPeak(ampResponse, thresholdCharge, thlDisp);
            delete[] ampResponse;

            //convert to digits from peak value //FIXME here ToT
            if (boost::get<1>(totPeakToa) > minEnergy) {
                Digit *digit = new Digit;
                digit->SetColumn((*iPCItr).first.first);
                digit->SetLine((*iPCItr).first.second);
                digit->SetEnergy((boost::get<1>(totPeakToa) / keV));
                digit->SetToT(boost::get<0>(totPeakToa));
                digit->SetToA(boost::get<2>(totPeakToa));
                digit->SetEvent(event);
                digitCollection->insert(digit);
            }
        }

        //free memory
        iPCItr = inducedPixelContent->begin();
        for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {
            G4double *tempArray = (*iPCItr).second;
            delete[] tempArray;
        }
        inducedPixelContent->clear();
        delete inducedPixelContent;
    }

    //CSM mode
    if(useCSM == true){

        map<pair<G4int, G4int>, G4double * >::iterator iPCItr;
        iPCItr = inducedPixelContent->begin();

        map<pair<G4int, G4int>, G4double * >* preampResponse    = new map<pair<G4int, G4int>, G4double * >;
        map<pair<G4int, G4int>, boost::tuple<G4double,G4double,G4double> >* sumPreampResponse = new map<pair<G4int, G4int>, boost::tuple<G4double,G4double,G4double> >;
        pair<G4int,G4int> tempPixel;
        G4double* emptyDummy = new G4double[nAmpResponseElements]();

        //preampresponse of every pixel
        for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {
            G4double* ampResponse = new G4double[nAmpResponseElements]();
            G4double* inducedChargeArray = (*iPCItr).second;

            //Preamp
            convolve(inducedChargeArray, ampResponse, nPulseArrayElements, impRes, nImpResEl);

            //write single preamp response to file
            if (writePeakToFile == true) {

                G4int k = (*iPCItr).first.first * 1000 + (*iPCItr).first.second;
                char filename[64];
                std::sprintf(filename, "preampSingle%d", k);

                writeToFile(filename, ampResponse, nAmpResponseElements);
            }

            //store single preampResponse
            tempPixel.first = (*iPCItr).first.first;
            tempPixel.second = (*iPCItr).first.second;
            preampResponse->operator [](tempPixel) = ampResponse;
        }

        //sum the responses and store
        map<pair<G4int, G4int>, G4double * >::iterator pRItr;
        map<pair<G4int, G4int>, G4double * >::iterator tempPRItr;
        pRItr = preampResponse->begin();

        for(; pRItr != preampResponse->end(); pRItr++){

            pair<G4int,G4int>mainPixel = pRItr->first;
            pair<G4int,G4int>iteratePixel = mainPixel;

            //get the preamp arrays
            G4double *tempPreampReponse;
            //sum the responses
            G4double* csmSumPreamp = new G4double[nAmpResponseElements]();

            G4int sumArrLow = 0;
            G4int sumArrHigh = 1;

            for(G4int i=sumArrLow; i<=sumArrHigh; i++){
                for(G4int j=sumArrLow; j<=sumArrHigh; j++){
                    //set pixel
                    iteratePixel.first = mainPixel.first + i;
                    iteratePixel.second = mainPixel.second + j;

                    //find response pulse
                    tempPRItr = preampResponse->find(iteratePixel);
                    if(tempPRItr != preampResponse->end()){
                        tempPreampReponse = tempPRItr->second;

                        //add to sum
                        for(G4int k=0; k < nAmpResponseElements; k++){
                            csmSumPreamp[k] += tempPreampReponse[k];
                        }
                    }
                }
            }

            //add noise
            for(G4int i=0; i < nAmpResponseElements; i++){
                csmSumPreamp[i] = ElectronicsNoise(csmSumPreamp[i]);
            }

            //write CSM preamp response to file
            if (writePeakToFile == true) {

                G4int k = (*pRItr).first.first * 1000 + (*pRItr).first.second;
                char filename[64];
                std::sprintf(filename, "preampCSM%d", k);

                writeToFile(filename, csmSumPreamp, nAmpResponseElements);
            }

            //get energy
            boost::tuple<G4double,G4double,G4double> totPeakToa;
            G4int pixelIndex = ((*pRItr).first.first) * nPixel + ((*pRItr).first.second);
            G4double thlDisp = TpxThlDisp[pixelIndex];

            totPeakToa = getToTandPeak(csmSumPreamp, thresholdCharge, thlDisp);

            sumPreampResponse->operator [](mainPixel) = totPeakToa;

            //delete the array
            delete[] csmSumPreamp;
        }

        //compare the sums
        map<pair<G4int, G4int>, boost::tuple<G4double,G4double,G4double> >::iterator sPRItr = sumPreampResponse->begin();
        map<pair<G4int, G4int>, boost::tuple<G4double,G4double,G4double> >::iterator tempSPRItr;
        for (; sPRItr != sumPreampResponse->end() ; sPRItr++) {

            //get tempPixel
            pair<G4int,G4int> tempPixelCSM = sPRItr->first;
            pair<G4int,G4int> iteratePixel = tempPixelCSM;

            //first value: ToT, second: peak amplitude
            boost::tuple<G4double,G4double,G4double> cmpEnergy(0,0,0);

            //get sumEnergy of pixel
            boost::tuple<G4double,G4double,G4double> pixelEnergy = sumPreampResponse->operator [](tempPixelCSM);
            G4bool isMax = true;

            //sets the borders for the comparison
            G4int compArrLow    = -1;
            G4int compArrHigh   = 1;

            //compare neighbouring pixels with main pixel
            for(G4int i=compArrLow; i<=compArrHigh; i++){
                for(G4int j=compArrLow; j<=compArrHigh; j++){
                    if(!(i == 0 && j == 0)){
                        iteratePixel = tempPixelCSM;
                        iteratePixel.first += i;
                        iteratePixel.second += j;
                        tempSPRItr = sumPreampResponse->find(iteratePixel);
                        if (tempSPRItr != sumPreampResponse->end()) {
                            cmpEnergy = tempSPRItr->second;
                            if(boost::get<1>(cmpEnergy) > boost::get<1>(pixelEnergy)) isMax = false;
                        }
                    }
                }
            }

            //create digit
            if (isMax == true) {
                Digit *digit = new Digit;
                digit->SetColumn(tempPixelCSM.first);
                digit->SetLine(tempPixelCSM.second);
                digit->SetEnergy((boost::get<1>(pixelEnergy) / keV));
                digit->SetToT(boost::get<0>(pixelEnergy));
                digit->SetToA(boost::get<2>(pixelEnergy));
                digit->SetEvent(event);
                digitCollection->insert(digit);
            }
        }

        //free memory
        //induced pulses
        iPCItr = inducedPixelContent->begin();
        for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {
            G4double *tempArray = (*iPCItr).second;
            delete[] tempArray;
        }
        inducedPixelContent->clear();
        delete inducedPixelContent;

        //preamp response
        pRItr = preampResponse->begin();

        for(; pRItr != preampResponse->end(); pRItr++){
            G4double* tempArray = (*pRItr).second;
            delete[] tempArray;
        }
        preampResponse->clear();
        delete preampResponse;

        //preamp sum
        sumPreampResponse->clear();
        delete sumPreampResponse;

        delete[] emptyDummy;
    }
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
G4double PreampMedipix::ElectronicsNoise(G4double charge)
{
    G4double sigma = elecSigma;
    if(useCSM == true && preampType != 2) sigma = sigma * 2;
    return CLHEP::RandGauss::shoot(charge, sigma);
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
G4double PreampMedipix::convolve(G4double* in, G4double* out, G4int inLength, G4double* kernel, G4int kernel_length)
{
    G4double maxCharge = 0;

    for(G4int i=0; i<(kernel_length); i++){

        out[i] = 0.0;
        for(G4int k=0; k<kernel_length; k++){
            if((i-k) >=0 && (i-k) < inLength){
                out[i] += in[i-k] * kernel[k];
            }
        }
        if(out[i] > maxCharge) maxCharge = out[i];
    }

    return maxCharge;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
boost::tuple<G4double,G4double,G4double> PreampMedipix::getToTandPeak(G4double* preamp, G4double threshold, G4double thlDisp){
    G4int risingEdge = 0;

    boost::tuple<G4double,G4double,G4double> returnValues(0,0,0);
    // 0 = ToT, 1 = peakCharge, 2 = ToA

    for(G4int i=0; i<nAmpResponseElements; i++){

        if(risingEdge == 0 && preamp[i] >= (threshold+thlDisp)){
            risingEdge = i;
            // Set timeOfArrival
            if(boost::get<2>(returnValues) == 0) {
                boost::get<2>(returnValues) = i * pulsePrecision;
            }
        }

        if(risingEdge != 0 && preamp[i] <= (threshold+thlDisp)){
            // Set timeOverThreshold
            if((i - risingEdge) * pulsePrecision > boost::get<0>(returnValues)) {
                boost::get<0>(returnValues) = (i - risingEdge) * pulsePrecision;
            }
            risingEdge = 0;
        }

        // set the peakCharge
        if(boost::get<1>(returnValues) < (preamp[i]+thlDisp)) boost::get<1>(returnValues) = (preamp[i]+thlDisp);
    }

    boost::get<1>(returnValues) = boost::get<1>(returnValues) * nElectronHolePairs * eV;

    return returnValues;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::writeToFile(G4String filename, G4double* data, G4int nElements){
    ofstream myfile;

    myfile.open(filename);

    for (G4int it = 0; it < nElements; it++) {
        myfile << data[it] << " ";
    }
    myfile.close();
}


/**
 * @file methods/rann/ra_model.hpp
 * @author Ryan Curtin
 *
 * This is a model for rank-approximate nearest neighbor search.  It provides an
 * easy way to serialize a rank-approximate neighbor search model by abstracting
 * the types of trees and reflecting the RASearch API.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#ifndef MLPACK_METHODS_RANN_RA_MODEL_HPP
#define MLPACK_METHODS_RANN_RA_MODEL_HPP

#include <mlpack/core/tree/binary_space_tree.hpp>
#include <mlpack/core/tree/cover_tree.hpp>
#include <mlpack/core/tree/rectangle_tree.hpp>
#include <mlpack/core/tree/octree.hpp>
#include <boost/variant.hpp>
#include "ra_search.hpp"

namespace mlpack {
namespace neighbor {

/**
 * Alias template for RASearch
 */
template<typename SortPolicy,
         template<typename TreeMetricType,
                  typename TreeStatType,
                  typename TreeMatType> class TreeType>
using RAType = RASearch<SortPolicy,
                        metric::EuclideanDistance,
                        arma::mat,
                        TreeType>;

/**
 * RAMonoSearchVisitor executes a monochromatic neighbor search on the given
 * RAType. We don't make any difference for different instantiation of RAType.
 */
class RAMonoSearchVisitor : public boost::static_visitor<void>
{
 private:
  //! Number of neighbors to search for.
  const size_t k;
  //! Result matrix for neighbors.
  arma::Mat<size_t>& neighbors;
  //! Result matrix for distances.
  arma::mat& distances;

 public:
  //! Perform monochromatic nearest neighbor search.
  template<typename RAType>
  void operator()(RAType* ra) const;

  //! Construct the RAMonoSearchVisitor object with the given parameters.
  RAMonoSearchVisitor(const size_t k,
                      arma::Mat<size_t>& neighbors,
                      arma::mat& distances) :
      k(k),
      neighbors(neighbors),
      distances(distances)
  {};
};

/**
 * RABiSearchVisitor executes a bichromatic neighbor search on the given RAType.
 * We use template specialization to differentiate those tree types types that
 * accept leafSize as a parameter. In these cases, before doing neighbor search
 * a query tree with proper leafSize is built from the querySet.
 */
template<typename SortPolicy>
class RABiSearchVisitor : public boost::static_visitor<void>
{
 private:
  //! The query set for the bichromatic search.
  const arma::mat& querySet;
  //! The number of neighbors to search for.
  const size_t k;
  //! The results matrix for neighbors.
  arma::Mat<size_t>& neighbors;
  //! The result matrix for distances.
  arma::mat& distances;
  //! The number of points in a leaf (for BinarySpaceTrees).
  const size_t leafSize;

  //! Bichromatic neighbor search on the given RAType considering leafSize.
  template<typename RAType>
  void SearchLeaf(RAType* ra) const;

 public:
  //! Alias template necessary for visual c++ compiler.
  template<template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  using RATypeT = RAType<SortPolicy, TreeType>;

  //! Default Bichromatic neighbor search on the given RAType instance.
  template<template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(RATypeT<TreeType>* ra) const;

  //! Bichromatic search on the given RAType specialized for KDTrees.
  void operator()(RATypeT<tree::KDTree>* ra) const;

  //! Bichromatic search on the given RAType specialized for octrees.
  void operator()(RATypeT<tree::Octree>* ra) const;

  //! Construct the RABiSearchVisitor.
  RABiSearchVisitor(const arma::mat& querySet,
                    const size_t k,
                    arma::Mat<size_t>& neighbors,
                    arma::mat& distances,
                    const size_t leafSize);
};

/**
 * RATrainVisitor sets the reference set to a new reference set on the given
 * RAType. We use template specialization to differentiate those trees that
 * accept leafSize as a parameter. In these cases, a reference tree with proper
 * leafSize is built from the referenceSet.
 */
template<typename SortPolicy>
class RATrainVisitor : public boost::static_visitor<void>
{
 private:
  //! The reference set to use for training.
  arma::mat&& referenceSet;
  //! The leaf size, used only by BinarySpaceTree.
  size_t leafSize;

  //! Train on the given RAType considering the leafSize.
  template<typename RAType>
  void TrainLeaf(RAType* ra) const;

 public:
  //! Alias template necessary for visual c++ compiler.
  template<template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  using RATypeT = RAType<SortPolicy, TreeType>;

  //! Default Train on the given RAType instance.
  template<template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(RATypeT<TreeType>* ra) const;

  //! Train on the given RAType specialized for KDTrees.
  void operator()(RATypeT<tree::KDTree>* ra) const;

  //! Train on the given RAType specialized for Octrees.
  void operator()(RATypeT<tree::Octree>* ra) const;

  //! Construct the RATrainVisitor object with the given reference set, leafSize
  //! for BinarySpaceTrees.
  RATrainVisitor(arma::mat&& referenceSet,
                 const size_t leafSize);
};

/**
 * Exposes the SingleSampleLimit() method of the given RAType.
 */
class SingleSampleLimitVisitor : public boost::static_visitor<size_t&>
{
 public:
  template<typename RAType>
  size_t& operator()(RAType* ra) const;
};

/**
 * Exposes the FirstLeafExact() method of the given RAType.
 */
class FirstLeafExactVisitor : public boost::static_visitor<bool&>
{
 public:
  template<typename RAType>
  bool& operator()(RAType* ra) const;
};

/**
 * Exposes the SampleAtLeaves() method of the given RAType.
 */
class SampleAtLeavesVisitor : public boost::static_visitor<bool&>
{
 public:
  //! Return SampleAtLeaves (whether or not sampling is done at leaves).
  template<typename RAType>
  bool& operator()(RAType *) const;
};

/**
 * Exposes the Alpha() method of the given RAType.
 */
class AlphaVisitor : public boost::static_visitor<double&>
{
 public:
  //! Return Alpha parameter.
  template<typename RAType>
  double& operator()(RAType* ra) const;
};

/**
 * Exposes the Tau() method of the given RAType.
 */
class TauVisitor : public boost::static_visitor<double&>
{
 public:
  //! Get a reference to the Tau parameter.
  template<typename RAType>
  double& operator()(RAType* ra) const;
};

/**
 * Exposes the SingleMode() method of the given RAType.
 */
class SingleModeVisitor : public boost::static_visitor<bool&>
{
 public:
  //! Get a reference to the SingleMode parameter of the given RASearch object.
  template<typename RAType>
  bool& operator()(RAType* ra) const;
};

/**
 * Exposes the referenceSet of the given RAType.
 */
class RAReferenceSetVisitor : public boost::static_visitor<const arma::mat&>
{
 public:
  //! Return the reference set.
  template<typename RAType>
  const arma::mat& operator()(RAType* ra) const;
};

/**
 * RADeleteVisitor deletes the give RAType Instance.
 */
class RADeleteVisitor : public boost::static_visitor<void>
{
 public:
  //! Delete the RAType Object.
  template<typename RAType> void operator()(RAType* ra) const;
};

/**
 * NaiveVisitor exposes the Naive() method of the given RAType.
 */
class NaiveVisitor : public boost::static_visitor<bool&>
{
 public:
  /**
   * Get a reference to the naive parameter of the given RASearch object.
   */
  template<typename RAType>
  bool& operator()(RAType* ra) const;
};

/**
 * The RAModel class provides an abstraction for the RASearch class, abstracting
 * away the TreeType parameter and allowing it to be specified at runtime in
 * this class.  This class is written for the sake of the 'allkrann' program,
 * but is not necessarily restricted to that use.
 *
 * @param SortPolicy Sorting policy for neighbor searching (see RASearch).
 */
template<typename SortPolicy>
class RAModel
{
 public:
  /**
   * The list of tree types we can use with RASearch.  Does not include ball
   * trees; see #338.
   */
  enum TreeTypes
  {
    KD_TREE,
    COVER_TREE,
    R_TREE,
    R_STAR_TREE,
    X_TREE,
    HILBERT_R_TREE,
    R_PLUS_TREE,
    R_PLUS_PLUS_TREE,
    UB_TREE,
    OCTREE
  };

 private:
  //! The type of tree being used.
  TreeTypes treeType;
  //! The leaf size of the tree being used (useful only for the kd-tree).
  size_t leafSize;

  //! If true, randomly project into a new basis.
  bool randomBasis;
  //! The basis to project into.
  arma::mat q;

  //! The rank-approximate model.
  boost::variant<RAType<SortPolicy, tree::KDTree>*,
                 RAType<SortPolicy, tree::StandardCoverTree>*,
                 RAType<SortPolicy, tree::RTree>*,
                 RAType<SortPolicy, tree::RStarTree>*,
                 RAType<SortPolicy, tree::XTree>*,
                 RAType<SortPolicy, tree::HilbertRTree>*,
                 RAType<SortPolicy, tree::RPlusTree>*,
                 RAType<SortPolicy, tree::RPlusPlusTree>*,
                 RAType<SortPolicy, tree::UBTree>*,
                 RAType<SortPolicy, tree::Octree>*> raSearch;

 public:
  /**
   * Initialize the RAModel with the given type and whether or not a random
   * basis should be used.
   */
  RAModel(TreeTypes treeType = TreeTypes::KD_TREE, bool randomBasis = false);

  /**
   * Copy the given RAModel.
   *
   * @param other RAModel to copy.
   */
  RAModel(const RAModel& other);

  /**
   * Take ownership of the given RAModel.
   *
   * @param other RAModel to take ownership of.
   */
  RAModel(RAModel&& other);

  /**
   * Copy the given RAModel.
   *
   * @param other RAModel to copy.
   */
  RAModel& operator=(const RAModel& other);

  /**
   * Take ownership of the given RAModel.
   *
   * @param other RAModel to take ownership of.
   */
  RAModel& operator=(RAModel&& other);

  //! Clean memory, if necessary.
  ~RAModel();

  //! Serialize the model.
  template<typename Archive>
  void serialize(Archive& ar, const uint32_t /* version */);

  //! Expose the dataset.
  const arma::mat& Dataset() const;

  //! Get whether or not single-tree search is being used.
  bool SingleMode() const;
  //! Modify whether or not single-tree search is being used.
  bool& SingleMode();

  //! Get whether or not naive search is being used.
  bool Naive() const;
  //! Modify whether or not naive search is being used.
  bool& Naive();

  //! Get the rank-approximation in percentile of the data.
  double Tau() const;
  //! Modify the rank-approximation in percentile of the data.
  double& Tau();

  //! Get the desired success probability.
  double Alpha() const;
  //! Modify the desired success probability.
  double& Alpha();

  //! Get whether or not sampling is done at the leaves.
  bool SampleAtLeaves() const;
  //! Modify whether or not sampling is done at the leaves.
  bool& SampleAtLeaves();

  //! Get whether or not we traverse to the first leaf without approximation.
  bool FirstLeafExact() const;
  //! Modify whether or not we traverse to the first leaf without approximation.
  bool& FirstLeafExact();

  //! Get the limit on the size of a node that can be approximated.
  size_t SingleSampleLimit() const;
  //! Modify the limit on the size of a node that can be approximation.
  size_t& SingleSampleLimit();

  //! Get the leaf size (only relevant when the kd-tree is used).
  size_t LeafSize() const;
  //! Modify the leaf size (only relevant when the kd-tree is used).
  size_t& LeafSize();

  //! Get the type of tree being used.
  TreeTypes TreeType() const;
  //! Modify the type of tree being used.
  TreeTypes& TreeType();

  //! Get whether or not a random basis is being used.
  bool RandomBasis() const;
  //! Modify whether or not a random basis is being used.  Be sure to rebuild
  //! the model using BuildModel().
  bool& RandomBasis();

  //! Build the reference tree.
  void BuildModel(arma::mat&& referenceSet,
                  const size_t leafSize,
                  const bool naive,
                  const bool singleMode);

  //! Perform rank-approximate neighbor search, taking ownership of the query
  //! set.
  void Search(arma::mat&& querySet,
              const size_t k,
              arma::Mat<size_t>& neighbors,
              arma::mat& distances);

  /**
   * Perform rank-approximate neighbor search, using the reference set as the
   * query set.
   */
  void Search(const size_t k,
              arma::Mat<size_t>& neighbors,
              arma::mat& distances);

  //! Get the name of the tree type.
  std::string TreeName() const;
};

} // namespace neighbor
} // namespace mlpack

#include "ra_model_impl.hpp"

#endif


// export POPLAR_ENGINE_OPTIONS='{"target.workerStackSizeInBytes":1024}'

#include <iostream>
#include <vector>
#include <boost/timer/timer.hpp>

//~ poplar
#include <poplar/DeviceManager.hpp>
#include <poplar/IPUModel.hpp>
#include <poplar/Engine.hpp>
#include <poplar/Program.hpp>

//~ poplin
#include <poplin/codelets.hpp>

//~ popops
#include <popops/ElementWise.hpp>
#include <popops/ExprOp.hpp>
#include <popops/Pad.hpp>
#include <popops/codelets.hpp>

//~ poputil
#include <poputil/TileMapping.hpp>

//~ KF header
#include "KalmanFilter.h"

/* contents
~01 KalmanFilter::connectToIPU
~02 KalmanFilter::skipSwitch
~03 KalmanFilter::smoothingState
~04 KalmanFilter::appendTo
~05 KalmanFilter::inverse
~06 KalmanFilter::packHits
~07 KalmanFilter::iterate
~08 KalmanFilter::product
~09 KalmanFilter::scaledAdd
~10 KalmanFilter::filter
~11 KalmanFilter::propagateState
~12 KalmanFilter::jacobian
~13 KalmanFilter::projectEKF
~14 KalmanFilter::project
~15 KalmanFilter::packIterationTensors
~16 KalmanFilter::calcResidual
~17 KalmanFilter::calcChiSq
~18 KalmanFilter::chiSqTest
~19 KalmanFilter::smooth
*/

using namespace poplar;
using namespace poplar::program;
using namespace popops;

Device KalmanFilter::connectToIPU()
{

    DeviceManager manager = DeviceManager::createDeviceManager();
    Device dev;

    bool success = false;

    //~ looping over each IPU
    for (auto &d : manager.getDevices(poplar::TargetType::IPU, 1))
    {
        dev = std::move(d);
        std::cerr << "Trying to attach to IPU " << dev.getId();
        if ((success = dev.attach()))
        {
            std::cerr << " - attached" << std::endl;
            break;
        }
        else
        {
            std::cerr << std::endl;
        }
    }

    if (!success)
    {
        std::cerr << "Error attaching to device" << std::endl;
        exit(1);
    }

    return dev;

} //~ end KalmanFilter::connectToIPU

std::tuple<ComputeSet, Tensor, Tensor, Tensor>
KalmanFilter::skipSwitch(Graph &graph,
                         const Tensor &inputPOld,
                         const Tensor &inputCOld,
                         const Tensor &inputPNew,
                         const Tensor &inputCNew,
                         const Tensor &chiSq,
                         uint tile)
{

    ComputeSet computeSet = graph.addComputeSet("SkipSwitch" + std::to_string(tile));
    VertexRef vtx = graph.addVertex(computeSet, "SkipSwitch");

    graph.setCycleEstimate(vtx, 100);

    Tensor outP =
        graph.addVariable(FLOAT, {2, 1}, "switchOutP" + std::to_string(tile));
    Tensor outC =
        graph.addVariable(FLOAT, {2, 2}, "switchOutC" + std::to_string(tile));
    Tensor outD =
        graph.addVariable(FLOAT, {1, 1}, "switchOutD" + std::to_string(tile));

    graph.connect(vtx["inputPOld"], inputPOld);
    graph.connect(vtx["inputCOld"], inputCOld);
    graph.connect(vtx["inputPNew"], inputPNew);
    graph.connect(vtx["inputCNew"], inputCNew);
    graph.connect(vtx["chiSq"], chiSq);
    graph.connect(vtx["outP"], outP);
    graph.connect(vtx["outC"], outC);
    graph.connect(vtx["outD"], outD);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(outP, tile);
    graph.setTileMapping(outC, tile);
    graph.setTileMapping(outD, tile);

    return std::make_tuple(computeSet, outP, outC, outD);

} //~ end KalmanFilter::skipSwitch

std::tuple<ComputeSet, Tensor>
KalmanFilter::smoothingState(Graph &graph,
                             const Tensor &states,
                             const Tensor &itr,
                             uint tile)
{
    ComputeSet computeSet =
        graph.addComputeSet("smoothingState" + std::to_string(tile));

    VertexRef vtx = graph.addVertex(computeSet, "StateForSmoothing");

    graph.setCycleEstimate(vtx, 100);

    Tensor tmp =
        states.reshape({states.shape()[0], states.shape()[1] * states.shape()[2]});

    Tensor state =
        graph.addVariable(FLOAT,
                          {states.shape()[1], states.shape()[2]},
                          "state" + std::to_string(tile));

    graph.connect(vtx["state"], state.reshape({state.shape()[0] * state.shape()[1]}));
    graph.connect(vtx["states"], tmp);
    graph.connect(vtx["itr"], itr);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(tmp, tile);
    graph.setTileMapping(state, tile);

    return std::make_pair(computeSet,
                          state.reshape({states.shape()[1],
                          states.shape()[2]}));

} //~ end KalmanFilter::smoothingState

std::tuple<ComputeSet, Tensor>
KalmanFilter::appendTo(Graph &graph,
                       const Tensor &state,
                       const Tensor &itr,
                       Tensor &states,
                       uint tile)
{
    ComputeSet computeSet = graph.addComputeSet("append" + std::to_string(tile));
    VertexRef vtx = graph.addVertex(computeSet, "AppendState");
    graph.setCycleEstimate(vtx, 100);

    Tensor tmp = states.reshape(
            {states.shape()[0], states.shape()[1] * states.shape()[2]});

    Tensor newStates =
            graph.addVariable(FLOAT, tmp.shape(), "newStates" + std::to_string(tile));

    graph.connect(vtx["state"], state.reshape({state.shape()[0] * state.shape()[1]}));
    graph.connect(vtx["states"], tmp);
    graph.connect(vtx["newStates"], newStates);
    graph.connect(vtx["itr"], itr);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(tmp, tile);
    graph.setTileMapping(newStates, tile);

    return std::make_pair(computeSet,
                          newStates.reshape({states.shape()[0],
                          states.shape()[1],
                          states.shape()[2]}));

} //~ end KalmanFilter::appendTo

std::tuple<ComputeSet, Tensor>
KalmanFilter::inverse(Graph &graph,
                      const Tensor &in,
                      uint tile,
                      uint dim)
{

    ComputeSet computeSet = graph.addComputeSet("inverse" + std::to_string(tile));

    VertexRef vtx =
        graph.addVertex(computeSet, "MatrixInverse" + std::to_string(dim));

    graph.setCycleEstimate(vtx, 100);

    Tensor out =
            graph.addVariable(FLOAT, {dim, dim}, "invOut" + std::to_string(tile));

    graph.connect(vtx["in"], in);
    graph.connect(vtx["out"], out);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::inverse

std::tuple<ComputeSet, Tensor>
KalmanFilter::packHits(Graph &graph,
                       const Tensor &inputHits,
                       const Tensor &itr,
                       uint tile)
{

    ComputeSet computeSet =
            graph.addComputeSet("packHits4" + std::to_string(tile));

    VertexRef vtx = graph.addVertex(computeSet, "PackHits4");

    graph.setCycleEstimate(vtx, 10);

    // Same shape as ps
    Tensor out =
            graph.addVariable(FLOAT, {4, 1}, "hitsThisItr" + std::to_string(tile));

    graph.connect(vtx["inputHits"], inputHits);
    graph.connect(vtx["itr"], itr);
    graph.connect(vtx["out"], out);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::packHits

std::tuple<ComputeSet, Tensor>
KalmanFilter::iterate(Graph &graph,
                      const Tensor &iterator,
                      uint tile)
{

    ComputeSet computeSet = graph.addComputeSet("iterate" + std::to_string(tile));
    VertexRef vtx = graph.addVertex(computeSet, "Iterate");
    graph.setCycleEstimate(vtx, 10);

    Tensor out = graph.addVariable(INT, {1}, "itrOut" + std::to_string(tile));

    graph.connect(vtx["in"], iterator);
    graph.connect(vtx["out"], out);
    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::iterate

std::tuple<ComputeSet, Tensor>
KalmanFilter::product(Graph &graph,
                      const Tensor &first,
                      const Tensor &second,
                      uint tile)
{

    ComputeSet computeSet = graph.addComputeSet("product" + std::to_string(tile));
    VertexRef vtx = graph.addVertex(computeSet, "MatrixProduct");
    graph.setCycleEstimate(vtx, 100);

    uint outRows = first.shape()[0];
    uint outCols = second.shape()[1];

    Tensor out = graph.addVariable(FLOAT, {outRows, outCols},
                                   "prodOut" + std::to_string(tile));

    graph.connect(vtx["first"], first);
    graph.connect(vtx["second"], second);
    graph.connect(vtx["out"], out);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::product

std::tuple<ComputeSet, Tensor>
KalmanFilter::scaledAdd(Graph &graph,
                        const Tensor &first,
                        const Tensor &second,
                        uint tile,
                        float s,
                        std::string name)
{

    ComputeSet computeSet =
        graph.addComputeSet("scaledAdd" + name + std::to_string(tile));

    VertexRef vtx = graph.addVertex(computeSet, "ScaledAdd");
    graph.setCycleEstimate(vtx, 100);

    uint outRows = first.shape()[0];
    uint outCols = first.shape()[1];

    Tensor out = graph.addVariable(FLOAT, {outRows, outCols},
                                   "addOut" + name + std::to_string(tile));

    graph.connect(vtx["first"], first);
    graph.connect(vtx["second"], second);
    graph.connect(vtx["s"], s);
    graph.connect(vtx["out"], out);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::scaledAdd

std::tuple<Sequence, std::vector<Tensor>, std::vector<Tensor>>
KalmanFilter::filter(Graph &graph,
                     const std::vector<Tensor> &H,
                     const std::vector<Tensor> &G,
                     const std::vector<Tensor> &hits,
                     const std::vector<Tensor> &p,
                     const std::vector<Tensor> &C)
{

    Sequence op;

    std::vector<Tensor> outp(p.size());
    std::vector<Tensor> outc(p.size());

    for (uint i = 0; i < p.size(); i++)
    {

        outp[i] = graph.addVariable(FLOAT, {4, 1});
        graph.setTileMapping(outp[i], i);

        outc[i] = graph.addVariable(FLOAT, {4, 4});
        graph.setTileMapping(outc[i], i);

        auto [computeHG, HG] = product(graph, H[i], G[i], i);
        op.add(Execute(computeHG));

        // inv_C_proj + HG @ H
        auto [computeHGH, HGH] = product(graph, HG, H[i], i);
        op.add(Execute(computeHGH));

        auto [computeInverseC, invOutC] = inverse(graph, C[i], i, 4);
        op.add(Execute(computeInverseC));

        auto [computeAddCHGH, CHGH_HGH] = scaledAdd(graph, invOutC, HGH, i);
        op.add(Execute(computeAddCHGH));

        auto [computeInverseCHGH, invOutCHGH] = inverse(graph, CHGH_HGH, i, 4);
        op.add(Execute(computeInverseCHGH));

        auto [computeHGhits, HGhits] = product(graph, HG, hits[i], i);
        op.add(Execute(computeHGhits));

        auto [computePFilt, Cp_filt] = product(graph, invOutC, p[i], i);
        op.add(Execute(computePFilt));

        auto [computeAddCpHits, Cp_filt_hits] = scaledAdd(graph, Cp_filt, HGhits, i);
        op.add(Execute(computeAddCpHits));

        auto [computeCHGHPFilt, CHGHPFilt] = product(graph, invOutCHGH, Cp_filt_hits, i);
        op.add(Execute(computeCHGHPFilt));

        op.add(Copy(CHGHPFilt, outp[i]));
        op.add(Copy(invOutCHGH, outc[i]));

    }

    return std::make_tuple(op, outp, outc);

} //~ end KalmanFilter::filter

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::propagateState(Graph &graph,
                             std::vector<Tensor> &ps,
                             std::vector<Tensor> &d)
{

    Sequence op;

    std::vector<Tensor> outP(ps.size());

    for (uint i = 0; i < ps.size(); i++)
    {

        outP[i] = graph.addVariable(FLOAT, {2, 1});
        graph.setTileMapping(outP[i], i);

        // Need to allocate this earlier, and not during every function invocation!
        Tensor updMFlat = graph.addConstant<float>(FLOAT, {4, 1}, {0., 1., 0., 0.});
        Tensor updM = graph.addVariable(FLOAT, {2, 2});

        graph.setTileMapping(updMFlat, i);
        graph.setTileMapping(updM, i);

        updM = updMFlat.reshape({2, 2}); // Can't seem to set nD constant

        graph.setTileMapping(updM, i);

        // xTan = tan([x0, x1])
        Tensor pTan = popops::map(graph, popops::expr::UnaryOpType::TAN, ps[i], op);
        graph.setTileMapping(pTan, i);

        // x += d * [[0, 1], [0, 0]] * xTan -> x0 += d * tan(x1), x1 += 0

        auto [computeProdUpdTan, updTan] = product(graph, updM, pTan, i);
        op.add(Execute(computeProdUpdTan));

        auto [computeProdUpdD, updD] = product(graph, updTan, d[i], i);
        op.add(Execute(computeProdUpdD));

        auto [computeAddP, addP] = scaledAdd(graph, ps[i], updD, i);
        op.add(Execute(computeAddP));

        op.add(Copy(addP, outP[i]));

    }

    return std::make_tuple(op, outP);

} //~ end KalmanFilter::propagateState

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::jacobian(Graph &graph,
                       const std::vector<Tensor> &ps,
                       const std::vector<Tensor> &d)
{

    Sequence op;

    std::vector<Tensor> jacs(ps.size());

    for (uint i = 0; i < ps.size(); i++)
    {

        Tensor updMFlat = graph.addConstant<float>(FLOAT, {4, 1}, {0., 1., 0., 0.});
        Tensor updM = graph.addVariable(FLOAT, {2, 2});

        jacs[i] = graph.addVariable(FLOAT, {2, 2});

        graph.setTileMapping(updMFlat, i);
        graph.setTileMapping(updM, i);
        graph.setTileMapping(jacs[i], i);

        updM = updMFlat.reshape({2, 2});

        Tensor xCos = popops::map(graph, popops::expr::UnaryOpType::COS, ps[i], op);
        Tensor xCosSq =
            popops::map(graph, popops::expr::UnaryOpType::SQUARE, xCos, op);
        Tensor xCosSqI =
            popops::map(graph, popops::expr::UnaryOpType::INVERSE, xCosSq, op);

        Tensor xCosM = popops::pad(graph, xCosSqI,
                                   1,  // pad lower
                                   0,  // pad upper,
                                   1,  // pad dim
                                   0); // pad value

        graph.setTileMapping(xCos, i);
        graph.setTileMapping(xCosM, i);
        graph.setTileMapping(xCosSq, i);
        graph.setTileMapping(xCosSqI, i);

        auto [computeJac, jac] = product(graph, updM, xCosM, i);
        op.add(Execute(computeJac));

        Tensor updM2Flat =
            graph.addConstant<float>(FLOAT, {4, 1}, {1., 0., 0., 1.});

        Tensor updM2 = graph.addVariable(FLOAT, {2, 2});

        graph.setTileMapping(updM2Flat, i);
        graph.setTileMapping(updM2, i);

        updM2 = updM2Flat.reshape({2, 2});

        auto [computeAddJacUpd, jac_upd] = scaledAdd(graph, jac, updM2, i);
        op.add(Execute(computeAddJacUpd));

        op.add(Copy(jac_upd, jacs[i]));

    }

    return std::make_tuple(op, jacs);

} //~ end KalmanFilter::jacobian

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::projectEKF(Graph &graph,
                         const std::vector<Tensor> &jacs,
                         const std::vector<Tensor> &qs,
                         std::vector<Tensor> &covs)
{

    Sequence op;

    std::vector<Tensor> outC(covs.size());

    for (uint i = 0; i < jacs.size(); i++)
    {

        outC[i] = graph.addVariable(FLOAT, {2, 2});
        graph.setTileMapping(outC[i], i);

        auto [computeCovT1, covT1] = product(graph, jacs[i], covs[i], i);
        op.add(Execute(computeCovT1));

        Tensor jacsT = jacs[i].transpose();
        graph.setTileMapping(jacsT, i);

        auto [computeCovT2, covT2] = product(graph, covT1, jacsT, i);
        op.add(Execute(computeCovT2));

        auto [computeAddCovQ, covQ] = scaledAdd(graph, covT2, qs[i], i);
        op.add(Execute(computeAddCovQ));

        op.add(Copy(covQ, outC[i]));

    }

    return std::make_tuple(op, outC);

} //~ end KalmanFilter::projectEKF

std::tuple<Sequence, std::vector<Tensor>, std::vector<Tensor>>
KalmanFilter::project(Graph &graph,
                      const std::vector<Tensor> &ps,
                      const std::vector<Tensor> &C,
                      const std::vector<Tensor> &F,
                      const std::vector<Tensor> &Q)
{

    Sequence op;

    std::vector<Tensor> outP(ps.size());
    std::vector<Tensor> outC(ps.size());

    for (uint i = 0; i < outP.size(); i++)
    {

        auto [computePproj, p_proj] = product(graph, F[i], ps[i], i);
        op.add(Execute(computePproj));

        auto [computeFC, fc] = product(graph, F[i], C[i], i);
        op.add(Execute(computeFC));

        // Check this works wrt tile mapping
        auto [computeFCF, fcf] = product(graph, fc, F[i].transpose(), i);
        op.add(Execute(computeFCF));

        auto [computeCproj, c_proj] = scaledAdd(graph, fcf, Q[i], i);
        op.add(Execute(computeCproj));

        outC[i] = graph.addVariable(FLOAT, {4, 4});
        outP[i] = graph.addVariable(FLOAT, {4, 1});

        graph.setTileMapping(outP[i], i);
        graph.setTileMapping(outC[i], i);

        op.add(Copy(c_proj, outC[i]));
        op.add(Copy(p_proj, outP[i]));

    }

    return std::make_tuple(op, outP, outC);

} //~ end KalmanFilter::project

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::packIterationTensors(Graph &graph,
                                   const std::vector<Tensor> &loop,
                                   const std::vector<Tensor> &scatterInto,
                                   const std::vector<Tensor> &inputs)
{

    Sequence op;
    std::vector<Tensor> hits_packed(inputs.size());

    for (uint i = 0; i < inputs.size(); i++)
    {

        std::string iStr = std::to_string(i);

        hits_packed[i] = graph.addVariable(FLOAT, {4, 1}, "hits_packed" + iStr);
        graph.setTileMapping(hits_packed[i], i);

        auto [computeH, h] = packHits(graph, inputs[i], loop[i], i);
        op.add(Execute(computeH));
        op.add(Copy(h, hits_packed[i]));

    }

    return std::make_tuple(op, hits_packed);

} //~ end KalmanFilter::packIterationTensors

// Residual at step t
std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::calcResidual(Graph &graph,
                           const std::vector<Tensor> &hits,
                           const std::vector<Tensor> &p_filt,
                           const std::vector<Tensor> &H)
{

    Sequence op;
    std::vector<Tensor> res_out(p_filt.size());

    for (uint i = 0; i < res_out.size(); i++)
    {

        std::string iStr = std::to_string(i);

        res_out[i] = graph.addVariable(FLOAT, {4, 1}, "res" + iStr);
        graph.setTileMapping(res_out[i], i);

        auto [computeHp, Hp] = product(graph, H[i], p_filt[i], i);
        op.add(Execute(computeHp));

        // Scaled add with prefactor of -1
        auto [computeRes, res] = scaledAdd(graph, hits[i], Hp, i, -1.0, "res");
        op.add(Execute(computeRes));
        op.add(Copy(res, res_out[i]));

    }

    return std::make_tuple(op, res_out);

} //~ end KalmanFilter::calcResidual

// ChiSq for a plane of hits
std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::calcChiSq(Graph &graph,
                        const std::vector<Tensor> &res_filt,
                        const std::vector<Tensor> &G,
                        const std::vector<Tensor> &C_proj,
                        const std::vector<Tensor> &p_proj,
                        const std::vector<Tensor> &p_filt)
{

    Sequence op;
    std::vector<Tensor> chiSq(p_filt.size());

    for (uint i = 0; i < chiSq.size(); i++)
    {

        std::string iStr = std::to_string(i);

        Tensor rT = res_filt[i].transpose();
        graph.setTileMapping(rT, i);

        auto [computeRTG, rTG] = product(graph, rT, G[i], i);
        op.add(Execute(computeRTG));

        auto [computeResTerm, resTerm] = product(graph, rTG, res_filt[i], i);
        op.add(Execute(computeResTerm));

        auto [computeStateDiff, stateDiff] =
            scaledAdd(graph, p_filt[i], p_proj[i], i, -1.0);
        op.add(Execute(computeStateDiff));

        auto [computeInverseC, C_inv] = inverse(graph, C_proj[i], i, 4);
        op.add(Execute(computeInverseC));

        Tensor diffT = stateDiff.transpose();
        graph.setTileMapping(diffT, i);

        auto [computeStateTC, stateTC] = product(graph, diffT, C_inv, i);
        op.add(Execute(computeStateTC));

        auto [computeStateTerm, stateTerm] =
            product(graph, stateTC, stateDiff, i);
        op.add(Execute(computeStateTerm));

        auto [computeChiSqTot, chiSqTot] =
            scaledAdd(graph, resTerm, stateTerm, i, 1.0);
        op.add(Execute(computeChiSqTot));

        chiSq[i] = graph.addVariable(FLOAT, {1}, "chiSq" + iStr);
        graph.setTileMapping(chiSq[i], i);

        op.add(Copy(chiSqTot, chiSq[i]));

    }

    return std::make_tuple(op, chiSq);

} //~ end KalmanFilter::calcChiSq

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::chiSqTest(Graph &graph,
                        const std::vector<Tensor> &chiSq,
                        const std::vector<Tensor> &threshold)
{

    Sequence op;
    std::vector<Tensor> chiSqPred(chiSq.size());

    for (uint i = 0; i < chiSqPred.size(); i++)
    {

        std::string iStr = std::to_string(i);

        chiSqPred[i] = popops::map(graph,
                                   popops::expr::BinaryOpType::GREATER_THAN,
                                   chiSq[i], threshold[i], op);

        graph.setTileMapping(chiSqPred[i], i);

    }

    return std::make_tuple(op, chiSqPred);

} //~ end KalmanFilter::chiSqTest

std::tuple<Sequence, std::vector<Tensor>, std::vector<Tensor>>
KalmanFilter::smooth(Graph &graph,
                     const std::vector<Tensor> &p_smooth_prev,
                     const std::vector<Tensor> &C_smooth_prev,
                     const std::vector<Tensor> &p_filt,
                     const std::vector<Tensor> &C_filt,
                     const std::vector<Tensor> &p_proj,
                     const std::vector<Tensor> &C_proj,
                     const std::vector<Tensor> &F)
{

    Sequence op;

    std::vector<Tensor> p_smoothOut(p_filt.size());
    std::vector<Tensor> C_smoothOut(p_filt.size());

    for (uint i = 0; i < p_filt.size(); i++)
    {

        std::string iStr = std::to_string(i);

        // Compute A

        auto [computeInverseC_proj, invOutC_proj] = inverse(graph, C_proj[i], i, 4);
        op.add(Execute(computeInverseC_proj));

        auto [computeFC, fc] = product(graph, F[i].transpose(), invOutC_proj, i);
        op.add(Execute(computeFC));

        auto [computeA, a] = product(graph, C_filt[i], fc, i);
        op.add(Execute(computeA));

        auto [computePDiff, pDiff] =
            scaledAdd(graph, p_smooth_prev[i], p_proj[i], i, -1);
        op.add(Execute(computePDiff));

        auto [computeCDiff, cDiff] =
            scaledAdd(graph, C_smooth_prev[i], C_proj[i], i, -1);
        op.add(Execute(computeCDiff));

        auto [computeAPDiff, aPDiff] = product(graph, a, pDiff, i);
        op.add(Execute(computeAPDiff));

        auto [computeP_smooth, p_smooth] = scaledAdd(graph, p_filt[i], aPDiff, i);
        op.add(Execute(computeP_smooth));

        auto [computeCdiffA, cDiffA] = product(graph, cDiff, a.transpose(), i);
        op.add(Execute(computeCdiffA));

        auto [computeACdiffA, aCDiffA] = product(graph, a, cDiffA, i);
        op.add(Execute(computeACdiffA));

        auto [computeC_smooth, C_smooth] = scaledAdd(graph, C_filt[i], aCDiffA, i);
        op.add(Execute(computeC_smooth));

        p_smoothOut[i] = graph.addVariable(FLOAT, {4, 1}, "p_smooth" + iStr);
        C_smoothOut[i] = graph.addVariable(FLOAT, {4, 4}, "C_smooth" + iStr);

        graph.setTileMapping(p_smoothOut[i], i);
        graph.setTileMapping(C_smoothOut[i], i);

        op.add(Copy(p_smooth, p_smoothOut[i]));
        op.add(Copy(C_smooth, C_smoothOut[i]));

    }

    return std::make_tuple(op, p_smoothOut, C_smoothOut);

} //~ end KalmanFilter::smooth