#include <iostream>
#include <iomanip>
#include <string>
#include <vector>
#include <queue>
#include <fstream>
#include <thread>
#include <atomic>
#include <functional>
#include <mutex>              // std::mutex, std::unique_lock
#include <condition_variable> // std::condition_variable
#include <boost/program_options.hpp> //Nice argparser

//#include "yolo_v2_class.hpp"        // imported functions from DLL
#include "opencv_utils.hpp"           // Split wrapper
#include "detector-wrapper.hpp"
#include "tracker.hpp"

#include "opencv_kalman.hpp"

#include "object.h"                 // municHMotorsport Object
#include <connector-1.0/client.h>   // Jussie & Alex: UDP/TCP Connector

#include <pylon/PylonIncludes.h>    // Pylon SDK

#include <opencv2/opencv.hpp>           // C++
#include "opencv2/core/version.hpp"
#include "opencv2/videoio/videoio.hpp"

namespace po = boost::program_options;

const int IMAGE_HEIGHT = 1024;
const int IMAGE_WIDTH  = 1280;

std::chrono::steady_clock::time_point steady_clara_start;

object_list_t  bbox_into_object_list( std::vector<bbox_t> boxes, Distance_Strategy strat, double distance_threshold = 20.0 )
{
    // Dirty conversion before adjusting bbox_t

    const double image_height = 1024.0;
    const double image_width  = 1280.0;

    // AOV calculated in
    // https://drive.google.com/open?id=1yfDZr2MJyqkEoziy0suvq2ejFtKxsR_Ukpk-Wdw1F5k
    // Degrees
    //const double h_AOV          = 75.3937;
    //const double v_AOV          = 63.5316;
    // RADIANS
    const double h_AOV         = 1.3158683; // 0.0229662;
    const double v_AOV         = 1.1088353; // 0.0193528;

    // All sensor and lens measures are in μm

    const double focal_length  = 4210.0;
    const double sensor_width  = 6182.4;
    const double sensor_height = 4953.6;

    // Cone measurements also in μm

    const double small_cone_width  = 228000.0;
    const double small_cone_height = 325000.0;

    const double large_cone_width  = 285000.0;
    const double large_cone_height = 505000.0;

    int index = 0;

    //  Remove detections that clip boundaries of the screen
    /*boxes.erase(std::remove_if(
            boxes.begin(),
            boxes.end(),
            [image_width, image_height](bbox_t & box)
           {
               return
                box.x          <= 0            ||
               (box.x + box.w) >= image_width  ||
                box.y          <= 0            ||
               (box.y + box.h) >= image_height;
            }), boxes.end());
    */
    std::vector<object_t> tmp_objects;

    std::chrono::steady_clock::time_point steady_clara_measure = std::chrono::steady_clock::now();
    std::chrono::duration<double> clara_delta = steady_clara_measure - steady_clara_start;

    //std::cout << "Calculating using Height" << std::endl;
    // Calculate Distances with cone height
    for( bbox_t& box : boxes )
    {
        double dimension_on_sensor;
        double small_cone_dimension;
        double large_cone_dimension;

        double snd_dimension_on_sensor;
        double snd_small_cone_dimension;
        double snd_large_cone_dimension;

        switch( strat )
        {
            case CONE_HEIGHT:
                dimension_on_sensor  = (box.h / image_height) * sensor_height;
                small_cone_dimension = small_cone_height;
                large_cone_dimension = large_cone_height;
            case CONE_WIDTH:
                dimension_on_sensor  = (box.w / image_width) * sensor_width;
                small_cone_dimension = small_cone_width;
                large_cone_dimension = large_cone_width;
            case CONE_AVERAGE:
                dimension_on_sensor  = (box.h / image_height) * sensor_height;
                snd_dimension_on_sensor  = (box.w / image_width) * sensor_width;

                small_cone_dimension = small_cone_height;
                snd_small_cone_dimension = small_cone_width;

                large_cone_dimension = large_cone_height;
                snd_large_cone_dimension = large_cone_width;
        }

        double perpendicular_distance;
        double snd_perpendicular_distance;

        double straight_distance;
        double snd_straight_distance;

        // Angles in FZModell-Koordinaten
        double angle     = (( box.x + box.w/2 ) / image_width ) * (- h_AOV) + ( h_AOV / 2);

        if(box.obj_id > 2)
        {
            perpendicular_distance = ((large_cone_dimension * focal_length) / dimension_on_sensor) / 1000000.0;
            if( strat == CONE_AVERAGE ) snd_perpendicular_distance = ((snd_large_cone_dimension * focal_length) / snd_dimension_on_sensor) / 1000000.0;
        }
        else
        {
            perpendicular_distance = ((small_cone_dimension * focal_length) / dimension_on_sensor) / 1000000.0;
            if( strat == CONE_AVERAGE ) snd_perpendicular_distance = ((snd_small_cone_dimension * focal_length) / snd_dimension_on_sensor) / 1000000.0;
        }

        //if (strat == CONE_AVERAGE) perpendicular_distance = ( perpendicular_distance + snd_perpendicular_distance ) / 2.0;

        straight_distance = perpendicular_distance / std::cos( angle );
        if (strat == CONE_AVERAGE)
        {
            snd_straight_distance = snd_perpendicular_distance / std::cos( angle );
            //std::cout << straight_distance << "\n" << snd_straight_distance << "\n";
            straight_distance = ( straight_distance + snd_straight_distance ) / 2;
        }

        if(straight_distance < distance_threshold)
        {
            object_t temp;
            temp.distance = straight_distance;
            temp.angle    = angle;
            temp.type     = box.obj_id;
            temp.x_car = 0;
            temp.y_car = 0;
            temp.angle_yaw = 0;
            temp.vx = 0;
            temp.vy = 0;
            temp.ax = 0;
            temp.ay = 0;
            temp.yaw_rate = 0;
            temp.steering_rad = 0;
            temp.time_s = clara_delta.count();

            tmp_objects.push_back(temp);
        }
        else
        {
            object_t temp;
            temp.distance = straight_distance;
            temp.angle    = angle;
            temp.type     = box.obj_id;
            temp.x_car = 0;
            temp.y_car = 0;
            temp.angle_yaw = 0;
            temp.vx = 0;
            temp.vy = 0;
            temp.ax = 0;
            temp.ay = 0;
            temp.yaw_rate = 0;
            temp.steering_rad = 0;
            temp.time_s = clara_delta.count();

            tmp_objects.push_back(temp);
        }
    }

    // Put object_t's into real list
    object_list_t cones;
    cones.size  = tmp_objects.size();
    cones.time_s = clara_delta.count();

    for(int i = 0; i < tmp_objects.size(); i++)
    {
        cones.element[i] = tmp_objects[i];
    }
    return cones;
}

// Put into opencv-utils after decoupling distance estimation and list creation
void draw_boxes(cv::Mat mat_img, std::vector<bbox_t> result_vec, std::vector<std::string> obj_names, double distance_threshold, int current_det_fps = -1, int current_cap_fps = -1)
{
    object_list_t height_objects = bbox_into_object_list( result_vec, Distance_Strategy::CONE_HEIGHT, distance_threshold);

    int index = 0;
    int const colors[6][3] = { { 1,0,1 },{ 0,0,1 },{ 0,1,1 },{ 0,1,0 },{ 1,1,0 },{ 1,0,0 } };

    for (auto &i : result_vec) {
        cv::Scalar color = obj_id_to_color(i.obj_id);
        cv::rectangle(mat_img, cv::Rect(i.x, i.y, i.w, i.h), color, 2);
        if (obj_names.size() > i.obj_id) {
            std::string obj_caption = std::to_string(height_objects.element[index].distance) + " , angle: " + std::to_string(height_objects.element[index].angle);
            if (i.track_id > 0) std::to_string(height_objects.element[index].distance) += " - " + std::to_string(i.track_id);
            cv::Size const text_size = getTextSize(obj_caption, cv::FONT_HERSHEY_SIMPLEX, 0.8, 2, 0);
            int const max_width = (text_size.width > i.w + 2) ? text_size.width : (i.w + 2);
            cv::rectangle(mat_img, cv::Point2f(std::max((int)i.x - 1, 0), std::max((int)i.y - 30, 0)),
                cv::Point2f(std::min((int)i.x + max_width, mat_img.cols-1), std::min((int)i.y, mat_img.rows-1)),
                color, CV_FILLED, 8, 0);
            // Yellow captions in black
            if (i.obj_id == 0) putText(mat_img, obj_caption, cv::Point2f(i.x, i.y - 10), cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 0, 0), 2);
            else putText(mat_img, obj_caption, cv::Point2f(i.x, i.y - 10), cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(255, 255, 255), 2);
        }
        ++index;
    }
    if (current_det_fps >= 0 && current_cap_fps >= 0) {
        std::string fps_str = "FPS detection: " + std::to_string(current_det_fps) + "   FPS capture: " + std::to_string(current_cap_fps);
        putText(mat_img, fps_str, cv::Point2f(10, 20), cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 0, 0), 2);
    }
}

void show_console_result(std::vector<bbox_t> const result_vec, std::vector<std::string> const obj_names) {
    for (auto &i : result_vec) {
        if (obj_names.size() > i.obj_id) std::cout << obj_names[i.obj_id] << " - ";
        std::cout << "obj_id = " << i.obj_id << ",  x = " << i.x << ", y = " << i.y
            << ", w = " << i.w << ", h = " << i.h
            << std::setprecision(3) << ", prob = " << i.prob << std::endl;
    }
}

void show_console_result_distances(std::vector<bbox_t> const result_vec, std::vector<std::string> const obj_names) {

        int count = 0;

        // Width distance estimation
        object_list_t  width_objects  = bbox_into_object_list(result_vec, Distance_Strategy::CONE_WIDTH);
        //Height distance estimation
        object_list_t  height_objects = bbox_into_object_list(result_vec, Distance_Strategy::CONE_HEIGHT);
        //Average distance estimation
        object_list_t  avg_objects = bbox_into_object_list(result_vec, Distance_Strategy::CONE_AVERAGE);

        for (auto &i : result_vec)
        {
            if (obj_names.size() > i.obj_id) std::cout << obj_names[i.obj_id] << " - ";
        std::cout << "obj_id = " << i.obj_id << ",  x = " << i.x << ", y = " << i.y
                  << ", w = " << i.w << ", h = " << i.h
                  << std::setprecision(3) << ", prob = " << i.prob
                  << std::setprecision(6) << ", distance_width = " << width_objects.element[count].distance
                  << ", distance_height = " << height_objects.element[count].distance
                  << ", avg_height = " << avg_objects.element[count].distance

                  << ", angle = " << height_objects.element[count].angle << '\n';
            count++;
    }
}

void show_console_result_CLARA_test(std::vector<bbox_t> const result_vec, std::vector<std::string> const obj_names)
{
    int count = 0;

    //Height distance estimation
    object_list_t height_objects = bbox_into_object_list(result_vec, Distance_Strategy::CONE_HEIGHT);

    for (auto &i : result_vec)
    {
        // Distance
        std::cout << height_objects.element[count].distance << std::setprecision(6)
        << "," <<
        // Id
        i.obj_id
        << "," <<
        // Angle
        height_objects.element[count].angle << std::endl;
        ++count;
    }
}

void send_objects_tcp(std::vector<bbox_t> const result_vec, connector::client< connector::TCP > & sender, Distance_Strategy strat, double distance_threshold) {

    object_list_t  objects = bbox_into_object_list( result_vec, strat, distance_threshold );

    sender.send_tcp< uint32_t >( objects.size );
    sender.send_tcp< object_t >( objects.element[0], objects.size * sizeof( object_t ));
    sender.send_tcp< double   >( objects.time_s );
}

void send_objects_udp(std::vector<bbox_t> const result_vec, connector::client< connector::UDP > & sender, Distance_Strategy strat, double distance_threshold) {

    object_list_t  objects = bbox_into_object_list( result_vec, strat, distance_threshold );

    sender.send_udp< uint32_t >( objects.size );
    sender.send_udp< object_t >( objects.element[0], objects.size * sizeof( object_t ));
    sender.send_udp< double   >( objects.time_s );
}

std::vector<std::string> objects_names_from_file(std::string const filename) {
    std::ifstream file(filename);
    std::vector<std::string> file_lines;
    if (!file.is_open()) return file_lines;
    for(std::string line; getline(file, line);) file_lines.push_back(line);
    std::cout << "object names loaded \n";
    return file_lines;
}

int main(int argc, char *argv[])
{
    std::string       names_file;
    std::string       cfg_file;
    std::string       weights_file;
    std::string       filename;
    std::string       pfs_file_name;
    std::string       ip;
    int               pylon;
    int               record_stream;
    int               live_demo;
    int               valid_test;
    int               port;
    int               udp_test;
    int               tcp_test;
    float             thresh;
    int               undistort;
    int               tracking;
    int               strategy_index;

    double            distance_threshold;
    long              frame_counter = 0;

    Distance_Strategy distance_strategy;

    po::options_description desc("Allowed options");
    desc.add_options()
        ("help", "produce help message")
        ("classes", po::value<std::string>(&names_file)->default_value("/home/nvidia/Documents/github-repos/darknet-cvAlexey/data/cones.names"), ".txt or .list file with one class name per line")
        ("config", po::value<std::string>(&cfg_file)->default_value("/home/nvidia/Documents/github-repos/darknet-cvAlexey/cfg/yolov3-hires.cfg"),  "Darknet .cfg file")
        ("weights", po::value<std::string>(&weights_file)->default_value("/home/nvidia/Documents/github-repos/darknet-cvAlexey/weights/yolov3-hires_20000.weights"), "Darknet .weights file")
        ("image_file", po::value<std::string>(&filename), "Single Image file to run detection on")
        ("list_file", po::value<std::string>(&filename), "List file of image paths to run detections on")
        ("video_file", po::value<std::string>(&filename), "Single Video file to run detection on")
        ("basler", po::value<int>(&pylon)->default_value(1), "Run demo with Basler Cam if set to 1")
        ("pfs_file", po::value<std::string>(&pfs_file_name)->default_value("../cam_config_120fps_auto-balance.pfs"), "Camera Config File to load.")
        ("record", po::value<int>(&record_stream)->default_value(0), "Record openCV stream to .avi file")
        ("live_demo", po::value<int>(&live_demo)->default_value(0), "Show openCV stream")
        ("thresh", po::value<float>(&thresh)->default_value(0.20), "Set probability threshold for detection")
        ("undistort", po::value<int>(&undistort)->default_value(0), "Set undistortion flag")
        ("tracking", po::value<int>(&tracking)->default_value(1), "Set tracking flag")
        ("valid_test", po::value<int>(&valid_test)->default_value(0), "Set to write detections from -list_file to image files in cwd")
        ("port", po::value<int>(&port)->default_value(4401), "Set port to send objects to")
        ("ip", po::value<std::string>(&ip)->default_value("127.0.0.1"), "Set ip to send objects to, default is localhost via FSD::Connector")
        ("udp_test", po::value<int>(&udp_test)->default_value(1), "If true, sends objects to specified ip:port")
        ("tcp_test", po::value<int>(&tcp_test)->default_value(0), "If true, sends objects to specified ip:port")
        ("distance_strategy", po::value<int>(&strategy_index)->default_value(0), "Sets the distance estimation strategy to be used.\n0 := Height\n1 := Width.\n2 := Average")
        ("distance_threshold", po::value<double>(&distance_threshold)->default_value(20.0), "Sets the distance threshold value over which no detections are forwarded")

    ;

    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);

    if (vm.count("help"))
    {
        std::cout << desc << "\n";
        return 1;
    }

    // Set Distance Estimation
    if      ( strategy_index == 0 ) distance_strategy = Distance_Strategy::CONE_HEIGHT;
    else if ( strategy_index == 1 ) distance_strategy = Distance_Strategy::CONE_WIDTH;
    else if ( strategy_index == 2 ) distance_strategy = Distance_Strategy::CONE_AVERAGE;

    std::cout << "Using Distance Strat: " << distance_strategy << std::endl;

    // create a 3x3 double-precision camera matrix
    cv::Mat cam_mtx  = (cv::Mat_<double>(3,3)
                                        << 862.47564274, 0.0         , 692.5289645
                                          , 0.0        , 851.02364963, 528.42008433
                                          , 0.0        , 0.0         , 1.0);
    cv::Mat dist_mtx = (cv::Mat_<double>(1,5) << -0.25596272,  0.14915701,  0.00074994, -0.001857  , -0.05665815);
    cv::Mat rvec     = (cv::Mat_<double>(3,1) << -0.05461275,  1.12613537,  0.04233885);
    cv::Mat tvec     = (cv::Mat_<double>(3,1) << -1.65275017, -2.34888592,  13.89459771);

    cv::Mat map1, map2, identity_mtx;
    cv::setIdentity(identity_mtx);

    cv::Size imageSize( IMAGE_HEIGHT, IMAGE_WIDTH );
    // Get New Camera Matrix
    // (Old camera matrix, distortion coefficients, img_size, what to do with empty pixels, new_img_size, valid_pixels_roi, keep_center_principal_point)
    cv::Mat new_cam_mtx = cv::getOptimalNewCameraMatrix(cam_mtx, dist_mtx, imageSize, 1, imageSize, 0, true);
    cv::initUndistortRectifyMap(cam_mtx, dist_mtx, identity_mtx, cam_mtx, imageSize, CV_16SC2, map1, map2);

    std::cout << "Trying to initialize Pylon" << std::endl;
    // Initialize Pylon
    Pylon::PylonAutoInitTerm autoInitTerm;
    std::cout << "Pylon successfully initialized" << std::endl;

    Pylon::CInstantCamera camera;
    Pylon::CImageFormatConverter formatConverter;
    formatConverter.OutputPixelFormat = Pylon::PixelType_BGR8packed;
    Pylon::CPylonImage pylonImage;

    Pylon::CGrabResultPtr ptrGrabResult;

    std::cout << "Pylon Camera Instance and GrabPointer Created" << std::endl;

    if ( pylon )
    {
        try
        {
          std::cout << "Trying to Attach Cam" << std::endl;
          camera.Attach(Pylon::CTlFactory::GetInstance().CreateFirstDevice());
          std::cout << "Camera Attached.\n" << "Using Device: " << camera.GetDeviceInfo().GetModelName() << std::endl;

          camera.StartGrabbing(Pylon::EGrabStrategy::GrabStrategy_LatestImageOnly, Pylon::EGrabLoop::GrabLoop_ProvidedByUser);
          std::cout << "Grabbing Started" << std::endl;

          if ( pfs_file_name != "" )
          {
            Pylon::CFeaturePersistence::Load(pfs_file_name.c_str(), &camera.GetNodeMap(), true);
            std::cout << "Camera Config loaded: " << pfs_file_name << std::endl;
          }
        }
        catch ( Pylon::GenericException &e) { std::cerr << "An exception occurred." << std::endl << e.GetDescription() << std::endl; return 1; }

    }

    Detector detector(cfg_file, weights_file);

    steady_clara_start = std::chrono::steady_clock::now();

    auto obj_names = objects_names_from_file(names_file);

    std::string out_videofile = "/home/fsd/Videos/darknet-recording.avi";
    //std::string out_videofile = "/home/nvidia/Documents/github-repos/modi-scheduler/darknet-recording.avi";

    connector::client< connector::TCP > tcp_sender( port, ip );
    connector::client< connector::UDP > udp_sender( port, ip );

    if ( udp_test ) udp_sender.init();
    else if ( tcp_test) tcp_sender.init();

    Tracker_optflow tracker_flow;

    if ( tracking ) detector.wait_stream = true;

    while (true)
    {

        if(!pylon)
        {

        std::cout << "input image or video filename: ";
        if(filename.size() == 0) std::cin >> filename; std::cout << "File Name: " << filename << "\n";
        if (filename.size() == 0) break;

        try {
            extrapolate_coords_t extrapolate_coords;
            bool extrapolate_flag = false;
            float cur_time_extrapolate = 0, old_time_extrapolate = 0;
            preview_boxes_t large_preview(100, 150, false), small_preview(50, 50, true);
            bool show_small_boxes = false;

            std::string const file_ext = filename.substr(filename.find_last_of(".") + 1);
            std::string const protocol = filename.substr(0, 7);
            if (file_ext == "avi" || file_ext == "mp4" || file_ext == "mjpg" || file_ext == "mov" ||    // video file
                protocol == "rtmp://" || protocol == "rtsp://" || protocol == "http://" || protocol == "https:/")   // video network stream
            {
                cv::Mat cap_frame, cur_frame, det_frame, write_frame;
                std::queue<cv::Mat> track_optflow_queue;
                int passed_flow_frames = 0;
                std::shared_ptr<image_t> det_image;
                std::vector<bbox_t> result_vec, thread_result_vec;
                detector.nms = 0.02;    // comment it - if track_id is not required
                std::atomic<bool> consumed, videowrite_ready;
                bool exit_flag = false;
                consumed = true;
                videowrite_ready = true;
                std::atomic<int> fps_det_counter, fps_cap_counter, tracking_counter;
                fps_det_counter  = 0;
                fps_cap_counter  = 0;
                tracking_counter = 0;
                int current_det_fps = 0, current_cap_fps = 0, current_tracking_fps = 0;
                std::thread t_detect, t_cap, t_videowrite;
                std::mutex mtx;
                std::condition_variable cv_detected, cv_pre_tracked;
                std::chrono::steady_clock::time_point steady_start, steady_end;
                cv::VideoCapture cap(filename); cap >> cur_frame;
                int const video_fps = cap.get(CV_CAP_PROP_FPS);
                cv::Size const frame_size = cur_frame.size();
                cv::VideoWriter output_video;
                // Stream Recording
                if ( record_stream ) output_video.open(out_videofile, CV_FOURCC('M', 'J', 'P', 'G'), std::max(35, video_fps), frame_size, true);

                while (!cur_frame.empty())
                {
                    // always sync
                    if (t_cap.joinable())
                    {
                        t_cap.join();
                        ++fps_cap_counter;
                        cur_frame = cap_frame.clone();
                    }
                    t_cap = std::thread([&]()
                        {
                          cv::Mat temp;
                          cv::Mat undistort_img;

                          cap >> temp;

                          if( undistort )
                          {
                              //cv::undistort(temp, undistort_img, cam_mtx, dist_mtx, cam_mtx);
                              cv::remap(temp, undistort_img, map1, map2, cv::INTER_LINEAR);
                              cap_frame = temp;
                          }
                          else
                          {
                              cap >> cap_frame;
                          }
                        });
                    ++cur_time_extrapolate;

                    // swap result bouned-boxes and input-frame
                    if(consumed)
                    {
                        std::unique_lock<std::mutex> lock(mtx);
                        det_image = detector.mat_to_image_resize(cur_frame);
                        auto old_result_vec = detector.tracking_id(result_vec);
                        auto detected_result_vec = thread_result_vec;
                        result_vec = detected_result_vec;

                        if(tracking)
                        {
                          // track optical flow
                          if (track_optflow_queue.size() > 0)
                          {
                            cv::Mat first_frame = track_optflow_queue.front();
                            tracker_flow.update_tracking_flow(track_optflow_queue.front(), result_vec);

                            while (track_optflow_queue.size() > 1)
                            {
                                track_optflow_queue.pop();
                                result_vec = tracker_flow.tracking_flow(track_optflow_queue.front(), true);
                            }
                            track_optflow_queue.pop();
                            passed_flow_frames = 0;

                            result_vec = detector.tracking_id(result_vec);
                            auto tmp_result_vec = detector.tracking_id(detected_result_vec, false);
                            small_preview.set(first_frame, tmp_result_vec);

                            extrapolate_coords.new_result(tmp_result_vec, old_time_extrapolate);
                            old_time_extrapolate = cur_time_extrapolate;
                            extrapolate_coords.update_result(result_vec, cur_time_extrapolate - 1);
                          }
                        }
                        else
                        {
                          result_vec = detector.tracking_id(result_vec);  // comment it - if track_id is not required
                          extrapolate_coords.new_result(result_vec, cur_time_extrapolate - 1);
                        }
                        // add old tracked objects
                        for (auto &i : old_result_vec)
                        {
                            auto it = std::find_if(result_vec.begin(), result_vec.end(),
                                [&i](bbox_t const& b) { return b.track_id == i.track_id && b.obj_id == i.obj_id; });
                            bool track_id_absent = (it == result_vec.end());
                            if (track_id_absent)
                            {
                                if (i.frames_counter-- > 1)
                                    result_vec.push_back(i);
                            }
                            else
                            {
                                it->frames_counter = std::min((unsigned)3, i.frames_counter + 1);
                            }
                        }
                        if(tracking)
                        {
                          tracker_flow.update_cur_bbox_vec(result_vec);
                          result_vec = tracker_flow.tracking_flow(cur_frame, true);   // track optical flow
                        }
                        consumed = false;
                        cv_pre_tracked.notify_all();
                    }
                    // launch thread once - Detection
                    if (!t_detect.joinable())
                    {
                        t_detect = std::thread([&]()
                        {
                            auto current_image = det_image;
                            consumed = true;
                            while (current_image.use_count() > 0 && !exit_flag)
                            {
                                auto result = detector.detect_resized(*current_image, frame_size.width, frame_size.height,
                                    thresh, false); // true
                                ++fps_det_counter;
                                std::unique_lock<std::mutex> lock(mtx);
                                thread_result_vec = result;
                                consumed = true;
                                cv_detected.notify_all();
                                if (detector.wait_stream)
                                {
                                    while (consumed && !exit_flag) cv_pre_tracked.wait(lock);
                                }
                                current_image = det_image;
                            }
                        });
                    }
                    //while (!consumed);    // sync detection
                    if (!cur_frame.empty())
                    {
                        steady_end = std::chrono::steady_clock::now();
                        if (std::chrono::duration<double>(steady_end - steady_start).count() >= 1)
                        {
                            current_det_fps      = fps_det_counter;
                            current_cap_fps      = fps_cap_counter;
                            current_tracking_fps = tracking_counter;
                            steady_start = steady_end;
                            fps_det_counter  = 0;
                            fps_cap_counter  = 0;
                            tracking_counter = 0;
                        }

                        large_preview.set(cur_frame, result_vec);

                        if(tracking)
                        {
                           ++passed_flow_frames;
                          track_optflow_queue.push(cur_frame.clone());
                          result_vec = tracker_flow.tracking_flow(cur_frame); // track optical flow
                          ++tracking_counter;
                          extrapolate_coords.update_result(result_vec, cur_time_extrapolate);
                          small_preview.draw(cur_frame, show_small_boxes);
                        }

                        auto result_vec_draw = result_vec;
                        if (extrapolate_flag) {
                            result_vec_draw = extrapolate_coords.predict(cur_time_extrapolate);
                            cv::putText(cur_frame, "extrapolate", cv::Point2f(10, 40), cv::FONT_HERSHEY_COMPLEX_SMALL, 1.0, cv::Scalar(50, 50, 0), 2);
                        }

                        if      ( udp_test ) send_objects_udp( result_vec, udp_sender, distance_strategy, distance_threshold );
                        else if ( tcp_test ) send_objects_tcp( result_vec, tcp_sender, distance_strategy, distance_threshold );

                        std::cout << "Detection FPS: " << current_det_fps << "\n";
                        std::cout << "Capture   FPS: " << current_cap_fps << "\n";
                        std::cout << "Tracking FPS: " << current_tracking_fps << "\n";
                        show_console_result(result_vec, obj_names);
                        // Make Results always be on top of Console
                        std::cout << "\033[2J";
                        std::cout << "\033[1;1H";

                        if( live_demo || record_stream)  draw_boxes(cur_frame, result_vec_draw, obj_names, current_det_fps, current_cap_fps);

                        if( live_demo )
                        {
                            large_preview.draw(cur_frame);

                            cv::imshow("window name", cur_frame);
                            int key = cv::waitKey(1);   // 3 or 16ms

                            if (key == 'f') show_small_boxes = !show_small_boxes;
                            if (key == 'p') while (true) if(cv::waitKey(100) == 'p') break;
                            if (key == 't') tracking = !tracking;
                            if (key == 'e') extrapolate_flag = !extrapolate_flag;
                            if (key == 27) { exit_flag = true; break; }
                        }

                        if (output_video.isOpened() && videowrite_ready && record_stream)
                        {
                          if (t_videowrite.joinable()) t_videowrite.join();
                          write_frame = cur_frame.clone();
                          videowrite_ready = false;
                          t_videowrite = std::thread([&]()
                          {
                            output_video << write_frame; videowrite_ready = true;
                          });
                        }
                     }

                    if( record_stream && !tracking )
                    {
                      // wait detection result for video-file only (not for net-cam)
                      if (protocol != "rtsp://" && protocol != "http://" && protocol != "https:/")
                      {
                        std::unique_lock<std::mutex> lock(mtx);
                        while (!consumed) cv_detected.wait(lock);
                      }
                    }
                }
                exit_flag = true;
                if (t_cap.joinable()) t_cap.join();
                if (t_detect.joinable()) t_detect.join();
                if (t_videowrite.joinable()) t_videowrite.join();
                std::cout << "Video ended \n";
                break;
                return 0;
            }
            else if (file_ext == "txt")
                {   // list of image files
                std::cout << "Starting List File calculations" << std::endl;
                std::ifstream file(filename);
                if (!file.is_open()) std::cout << "File not found! \n";
                else
                    {
                        for (std::string line; std::getline(file, line);)  //file >> line;)
                        {

                        // Easier splitting for Alex
                        std::cout << "#" << std::endl;
                        std::cout << line << std::endl;
                        cv::Mat mat_img = cv::imread(line);
                        cv::Mat undistort_img;

                        //cv::undistort(mat_img, undistort_img, cam_mtx, dist_mtx, cam_mtx);
                        std::vector<bbox_t> result_vec;

                        if(undistort)
                        {
                          cv::Mat tmp(mat_img);
                          cv::remap(tmp, undistort_img, map1, map2, cv::INTER_LINEAR);
                          result_vec = detector.detect(undistort_img);
                        }
                        else result_vec = detector.detect(mat_img);

                        if      ( udp_test ) send_objects_udp( result_vec, udp_sender, distance_strategy, distance_threshold );
                        else if ( tcp_test ) send_objects_tcp( result_vec, tcp_sender, distance_strategy, distance_threshold );

                        //show_console_result_distances( result_vec, obj_names );
                        show_console_result_CLARA_test( result_vec, obj_names );

                        if( valid_test )
                        {
                          frame_counter++;
                          std::ostringstream stringStream;
                          stringStream << std::setfill('0');
                          stringStream << "valid_frames/jo6_";
                          stringStream << std::setw(8) << std::to_string(frame_counter);
                          stringStream << ".jpg";
                          std::string filename = stringStream.str();
                          draw_boxes(mat_img, result_vec, obj_names, distance_threshold);
                          cv::imwrite( filename , mat_img);
                        }
                    }
                    break;
                }
            }
            else
            {  // image file
                std::cout << "Single image file" << std::endl;

                cv::Mat undistort_img;
                cv::Mat mat_img = cv::imread(filename);

                cv::undistort(mat_img, undistort_img, cam_mtx, dist_mtx, cam_mtx);
                std::vector<bbox_t> result_vec = detector.detect(undistort_img);
                //result_vec = detector.tracking_id(result_vec);    // comment it - if track_id is not required

                if(live_demo)
                {
                    draw_boxes(undistort_img, result_vec, obj_names, distance_threshold);
                    cv::imshow("window name", undistort_img);
                }

                if      ( udp_test ) send_objects_udp( result_vec, udp_sender, distance_strategy, distance_threshold );
                else if ( tcp_test ) send_objects_tcp( result_vec, tcp_sender, distance_strategy, distance_threshold );

                show_console_result_distances(result_vec, obj_names);

                if(live_demo) cv::waitKey(0);
            }
        }
            catch (std::exception &e) { std::cerr << "exception: " << e.what() << "\n"; getchar(); }
            catch (...) { std::cerr << "unknown exception \n"; getchar(); }
            filename.clear();
        }
        else
        {
            try
            {
                std::cout << "Pylon Loop" << std::endl;

                extrapolate_coords_t extrapolate_coords;
                bool extrapolate_flag = false;
                float cur_time_extrapolate = 0, old_time_extrapolate = 0;
                preview_boxes_t large_preview(100, 150, false), small_preview(50, 50, true);
                bool show_small_boxes = false;

                cv::Mat cap_frame, cur_frame, det_frame, write_frame;
                std::queue<cv::Mat> track_optflow_queue;
                int passed_flow_frames = 0;
                std::shared_ptr<image_t> det_image;
                std::vector<bbox_t> result_vec, thread_result_vec;
                detector.nms = 0.02;    // comment it - if track_id is not required
                std::atomic<bool> consumed, videowrite_ready;
                consumed = true;
                videowrite_ready = true;
                std::atomic<int> fps_det_counter, fps_cap_counter, tracking_counter;
                fps_det_counter  = 0;
                fps_cap_counter  = 0;
                tracking_counter = 0;
                int current_det_fps = 0, current_cap_fps = 0, cur_tracking_fps = 0;
                std::thread t_detect, t_cap, t_videowrite;
                std::mutex mtx;
                std::condition_variable cv_detected, cv_pre_tracked;
                std::chrono::steady_clock::time_point steady_start, steady_end;
                if(camera.IsGrabbing())
                {
                    camera.RetrieveResult( 5000, ptrGrabResult, Pylon::ETimeoutHandling::TimeoutHandling_ThrowException);
                    if( ptrGrabResult->GrabSucceeded())
                    {
                        formatConverter.Convert(pylonImage, ptrGrabResult);
                        cv::Mat temp(ptrGrabResult->GetHeight(), ptrGrabResult->GetWidth(), CV_8UC3, (uint8_t *) pylonImage.GetBuffer());
                        cv::Mat undistort_img;

                        if( undistort )
                        {
                          //cv::undistort(temp, undistort_img, cam_mtx, dist_mtx, cam_mtx);
                          cv::remap(temp, undistort_img, map1, map2, cv::INTER_LINEAR);
                          cur_frame = undistort_img;
                        }
                        else
                        {
                          cur_frame = temp;
                        }
                        std::cout << "Cap Dims: " << ptrGrabResult->GetHeight() << " " << ptrGrabResult->GetWidth() << std::endl;
                    }
                }

                //cv::VideoCapture cap(filename); cap >> cur_frame;
                int const video_fps = 60; //cap.get(CV_CAP_PROP_FPS);

                int cam_video_fps;
                if (GenApi::IsAvailable(camera.GetNodeMap().GetNode("ResultingFrameRateAbs")))
                {
                    cam_video_fps = (int) GenApi::CFloatPtr(camera.GetNodeMap().GetNode("ResultingFrameRateAbs"))->GetValue();
                }
                else cam_video_fps = (int) GenApi::CFloatPtr(camera.GetNodeMap().GetNode("ResultingFrameRate"))->GetValue(); // BCON and USB use SFNC3 names

                std::cout << "Cam FPS: " << cam_video_fps << std::endl;

                cv::Size const frame_size = cur_frame.size();
                std::cout << frame_size.width << " " << frame_size.height << std::endl;
                // Record Stream
                cv::VideoWriter output_video;
                if ( record_stream ) output_video.open(out_videofile, CV_FOURCC('M', 'J', 'P', 'G'), /*std::max(35, cam_video_fps)*/ 30, frame_size, true);

                while (!cur_frame.empty())
                {
                    // always sync
                    if (t_cap.joinable())
                    {
                        t_cap.join();
                        ++fps_cap_counter;
                        cur_frame = cap_frame.clone();
                    }
                    t_cap = std::thread([&]()
                    {
                        if(camera.IsGrabbing())
                        {
                            camera.RetrieveResult( 5000, ptrGrabResult, Pylon::ETimeoutHandling::TimeoutHandling_ThrowException);
                            if( ptrGrabResult->GrabSucceeded())
                            {
                                formatConverter.Convert(pylonImage, ptrGrabResult);
                                cv::Mat temp(ptrGrabResult->GetHeight(), ptrGrabResult->GetWidth(), CV_8UC3, (uint8_t *) pylonImage.GetBuffer());

                                cv::Mat undistort_img;
                                if( undistort )
                                {
                                  //cv::undistort(temp, undistort_img, cam_mtx, dist_mtx, cam_mtx);
                                  cv::remap(temp, undistort_img, map1, map2, cv::INTER_LINEAR);
                                  cap_frame = undistort_img;
                                }
                                else
                                {
                                  cap_frame = temp;
                                }
                            }
                        }
                    });
                    ++cur_time_extrapolate;

                    // swap result bounded-boxes and input-frame
                    if(consumed)
                    {
                        std::unique_lock<std::mutex> lock(mtx);
                        det_image = detector.mat_to_image_resize(cur_frame);
                        auto old_result_vec = detector.tracking_id(result_vec);
                        auto detected_result_vec = thread_result_vec;
                        result_vec = detected_result_vec;

                        if(tracking)
                        {
                          // track optical flow
                          if (track_optflow_queue.size() > 0)
                          {
                            //std::cout << "\n !!!! all = " << track_optflow_queue.size() << ", cur = " << passed_flow_frames << std::endl;
                            cv::Mat first_frame = track_optflow_queue.front();
                            tracker_flow.update_tracking_flow(track_optflow_queue.front(), result_vec);

                            while (track_optflow_queue.size() > 1) {
                                track_optflow_queue.pop();
                                result_vec = tracker_flow.tracking_flow(track_optflow_queue.front(), true);
                            }
                            track_optflow_queue.pop();
                            passed_flow_frames = 0;

                            result_vec = detector.tracking_id(result_vec);
                            auto tmp_result_vec = detector.tracking_id(detected_result_vec, false);
                            small_preview.set(first_frame, tmp_result_vec);

                            extrapolate_coords.new_result(tmp_result_vec, old_time_extrapolate);
                            old_time_extrapolate = cur_time_extrapolate;
                            extrapolate_coords.update_result(result_vec, cur_time_extrapolate - 1);
                          }
                        }
                        // Detections only
                        else
                        {
                          result_vec = detector.tracking_id(result_vec);  // comment it - if track_id is not required
                          extrapolate_coords.new_result(result_vec, cur_time_extrapolate - 1);
                        // Connector Stuff
                        if ( udp_test ) send_objects_udp( result_vec, udp_sender, distance_strategy, distance_threshold );
                        else if ( tcp_test ) send_objects_tcp( result_vec, tcp_sender, distance_strategy, distance_threshold );
                        }
                        // add old tracked objects
                        for (auto &i : old_result_vec) {
                            auto it = std::find_if(result_vec.begin(), result_vec.end(),
                                [&i](bbox_t const& b) { return b.track_id == i.track_id && b.obj_id == i.obj_id; });
                            bool track_id_absent = (it == result_vec.end());
                            if (track_id_absent)
                            {
                                if (i.frames_counter-- > 1)
                                    result_vec.push_back(i);
                            }
                            else
                            {
                                it->frames_counter = std::min((unsigned)3, i.frames_counter + 1);
                            }
                        }

                        if(tracking)
                        {
                          tracker_flow.update_cur_bbox_vec(result_vec);
                          result_vec = tracker_flow.tracking_flow(cur_frame, true);   // track optical flow
                        }

                        consumed = false;
                        cv_pre_tracked.notify_all();
                    }
                    // launch thread once - Detection
                    if (!t_detect.joinable())
                    {
                        t_detect = std::thread([&]()
                        {
                            auto current_image = det_image;
                            consumed = true;
                            while (current_image.use_count() > 0) {
                                auto result = detector.detect_resized(*current_image, frame_size.width, frame_size.height, thresh, false);  // true
                                ++fps_det_counter;
                                std::unique_lock<std::mutex> lock(mtx);
                                thread_result_vec = result;
                                consumed = true;
                                cv_detected.notify_all();
                                if (detector.wait_stream) {
                                    while (consumed) cv_pre_tracked.wait(lock);
                                }
                                current_image = det_image;
                            }
                        });
                    }
                    //while (!consumed);    // sync detection

                    if (!cur_frame.empty())
                    {
                        steady_end = std::chrono::steady_clock::now();
                        if (std::chrono::duration<double>(steady_end - steady_start).count() >= 1)
                        {
                            current_det_fps  = fps_det_counter;
                            current_cap_fps  = fps_cap_counter;
                            cur_tracking_fps = tracking_counter;
                            steady_start = steady_end;
                            fps_det_counter  = 0;
                            fps_cap_counter  = 0;
                            tracking_counter = 0;
                        }

                        large_preview.set(cur_frame, result_vec);
                        if (tracking)
                        {
                          ++passed_flow_frames;
                          track_optflow_queue.push(cur_frame.clone());

                          result_vec = tracker_flow.tracking_flow(cur_frame); // track optical flow
                          ++tracking_counter;
                        }

                        extrapolate_coords.update_result(result_vec, cur_time_extrapolate);
                        small_preview.draw(cur_frame, show_small_boxes);

                        auto result_vec_draw = result_vec;
                        if (extrapolate_flag)
                        {
                            result_vec_draw = extrapolate_coords.predict(cur_time_extrapolate);
                            cv::putText(cur_frame, "extrapolate", cv::Point2f(10, 40), cv::FONT_HERSHEY_COMPLEX_SMALL, 1.0, cv::Scalar(50, 50, 0), 2);
                        }
                        if ( tracking )
                        {
                          if      ( udp_test ) send_objects_udp( result_vec, udp_sender, distance_strategy, distance_threshold );
                          else if ( tcp_test ) send_objects_tcp( result_vec, tcp_sender, distance_strategy, distance_threshold );
                        }
                        //std::cout << "Detection FPS: " << current_det_fps  << "\n";
                        //std::cout << "Capture   FPS: " << current_cap_fps  << "\n";
                        //std::cout << "Tracking  FPS: " << cur_tracking_fps << "\n";
                        //show_console_result_distances( result_vec, obj_names );

                        // Make Results always be on top of Console
                        //std::cout << "\033[2J";
                        //std::cout << "\033[1;1H";

                        if ( live_demo || record_stream ) draw_boxes(cur_frame, result_vec_draw, obj_names, current_det_fps, current_cap_fps);

                        if (  live_demo )
                        {
                            //cv::namedWindow( "OpenCV Display Window", CV_WINDOW_NORMAL);

                            large_preview.draw(cur_frame);

                            cv::imshow("OpenCV Display Window", cur_frame);
                            int key = cv::waitKey(1);   // 3 or 16ms
                            if (key == 'f') show_small_boxes = !show_small_boxes;
                            if (key == 't') tracking = !tracking;
                            if (key == 'p') while (true) if(cv::waitKey(100) == 'p') break;
                            if (key == 'e') extrapolate_flag = !extrapolate_flag;
                        }

                        if (output_video.isOpened() && videowrite_ready && record_stream)
                        {
                            if (t_videowrite.joinable()) t_videowrite.join();
                            write_frame = cur_frame.clone();
                            videowrite_ready = false;
                            t_videowrite = std::thread([&]()
                            {
                                 output_video << write_frame; videowrite_ready = true;
                            });
                        }
                    }

                    // wait detection result for video-file only (not for net-cam)
                    /*if ( record_stream )
                    {
                        std::unique_lock<std::mutex> lock(mtx);
                        while (!consumed) cv_detected.wait(lock);
                    }*/
                }
                if (t_cap.joinable()) t_cap.join();
                if (t_detect.joinable()) t_detect.join();
                if (t_videowrite.joinable()) t_videowrite.join();
                std::cout << "Video ended \n";

            }
        catch ( Pylon::GenericException &e) { std::cerr << "An exception occurred." << std::endl << e.GetDescription() << std::endl; return 1; }
        catch (std::exception &e) { std::cerr << "exception: " << e.what() << "\n"; getchar(); }
        catch (...) { std::cerr << "unknown exception \n"; getchar(); }
        filename.clear();
        }
    }
    return 0;
}


//Copyright (c) 2013 Singapore-MIT Alliance for Research and Technology
//Licensed under the terms of the MIT License, as described in the file:
//   license.txt   (http://opensource.org/licenses/MIT)

 /* PassengerDistribution.hpp
  *
  * \author Meenu James
  */

#pragma once

#include <boost/random.hpp>
#include <boost/math/distributions.hpp>

namespace sim_mob {

///Simple base class for passenger time distributions
class PassengerDist {
public:
    virtual ~PassengerDist() {}
    virtual double getnopassengers() = 0;
};

///Pre-specialized sub-classes for passenger time distributions
class NormalPassengerDist : public PassengerDist {
private:
    boost::mt19937 gen;
    boost::normal_distribution<double> dist;
    boost::variate_generator<boost::mt19937, boost::normal_distribution<double> > varGen;

public:
    NormalPassengerDist(double mean, double stdev)
     : gen(), dist(mean, stdev), varGen(gen, dist)
    {}
    virtual ~NormalPassengerDist() {}

    virtual double getnopassengers() { return varGen(); }
};

///Pre-specialized sub-classes for reaction time distributions
class LognormalPassengerDist : public PassengerDist {
private:
    boost::mt19937 gen;
    boost::lognormal_distribution<double> dist;
    boost::variate_generator<boost::mt19937, boost::lognormal_distribution<double> > varGen;

public:
    LognormalPassengerDist(double mean, double stdev)
     : gen(), dist(mean, stdev), varGen(gen, dist)
    {}
    virtual ~LognormalPassengerDist() {}

    virtual double getnopassengers() { return varGen(); }
};
class UniformPassengerDist : public PassengerDist {
private:
    boost::mt19937 gen;
    boost::uniform_int<int> dist;
    boost::variate_generator<boost::mt19937, boost::uniform_int<int> > varGen;

public:
    UniformPassengerDist(int min, int max)
     : gen(), dist(min ,max), varGen(gen, dist)
    {}
    virtual ~UniformPassengerDist() {}

    virtual double getnopassengers() { return varGen(); }
};

}


/** \file   cmr_spirit_recon.cpp
    \brief  Implement some functionalities commonly used in cmr applications for spirit recon
    \author Hui Xue
*/

#include "cmr_spirit_recon.h"

#include "mri_core_utility.h"
#include "mri_core_spirit.h"
#include "hoNDArray_reductions.h"
#include "hoNDFFT.h"
#include "hoSPIRIT2DOperator.h"
#include "hoLsqrSolver.h"
#include "hoSPIRIT2DTDataFidelityOperator.h"
#include "hoWavelet2DTOperator.h"
#include "hoGdSolver.h"
#include <boost/make_shared.hpp>

namespace Gadgetron {

    template <typename T>
    void perform_spirit_recon_linear_2DT(const Gadgetron::hoNDArray<T>& kspace, size_t startE1, size_t endE1, const Gadgetron::hoNDArray<T>& kerIm,
                                const Gadgetron::hoNDArray<T>& kspaceInitial, Gadgetron::hoNDArray<T>& res, size_t iter_max, double iter_thres, bool print_iter)
    {
        try
        {
            size_t RO = kspace.get_size(0);
            size_t E1 = kspace.get_size(1);
            size_t CHA = kspace.get_size(2);
            size_t N = kspace.get_size(3);
            size_t S = kspace.get_size(4);

            size_t ref_N = kerIm.get_size(4);
            size_t ref_S = kerIm.get_size(5);

            if(startE1>=E1) startE1 = 0;
            if(endE1>=E1) startE1 = 0;
            if(startE1>endE1)
            {
                startE1 = 0;
                endE1 = E1 - 1;
            }

            long long num = N*S;

            hoNDArray<T> ker_Shifted(kerIm);
            Gadgetron::hoNDFFT<typename realType<T>::Type>::instance()->ifftshift2D(kerIm, ker_Shifted);

            hoNDArray<T> kspace_Shifted;
            kspace_Shifted = kspace;
            Gadgetron::hoNDFFT<typename realType<T>::Type>::instance()->ifftshift2D(kspace, kspace_Shifted);

            hoNDArray<T> kspace_initial_Shifted;
            bool hasInitial = false;
            if ( kspaceInitial.dimensions_equal(&kspace) )
            {
                kspace_initial_Shifted = kspaceInitial;
                Gadgetron::hoNDFFT<typename realType<T>::Type>::instance()->ifftshift2D(kspaceInitial, kspace_initial_Shifted);
                hasInitial = true;
            }

            res = kspace;

#ifdef USE_OMP
            int numThreads = (int)num;
            if (numThreads > omp_get_num_procs()) numThreads = omp_get_num_procs();
            GDEBUG_CONDITION_STREAM(print_iter, "numThreads : " << numThreads);
#endif // USE_OMP

            long long ii;

            std::vector<size_t> dim(3, 1);
            dim[0] = RO;
            dim[1] = E1;
            dim[2] = CHA;

#pragma omp parallel default(none) private(ii) shared(num, N, S, RO, E1, CHA, dim, startE1, endE1, ref_N, ref_S, kspace, res, kspace_Shifted, ker_Shifted, kspace_initial_Shifted, hasInitial, iter_max, iter_thres, print_iter) num_threads(numThreads) if(num>1)
            {
                boost::shared_ptr< hoSPIRIT2DOperator< T > > oper(new hoSPIRIT2DOperator< T >(&dim));
                hoSPIRIT2DOperator< T >& spirit = *oper;
                spirit.use_non_centered_fft_ = true;
                spirit.no_null_space_ = false;

                if (ref_N == 1 && ref_S == 1)
                {
                    boost::shared_ptr<hoNDArray<T> > ker(new hoNDArray< T >(RO, E1, CHA, CHA, ker_Shifted.begin()));
                    spirit.set_forward_kernel(*ker, false);
                }

                hoLsqrSolver< T > cgSolver;
                cgSolver.set_tc_tolerance((float)iter_thres);
                cgSolver.set_max_iterations( (unsigned int)iter_max);
                cgSolver.set_output_mode(print_iter ? hoLsqrSolver< T >::OUTPUT_VERBOSE : hoLsqrSolver< T >::OUTPUT_SILENT);
                cgSolver.set_encoding_operator(oper);

                hoNDArray< T > b(RO, E1, CHA);
                hoNDArray< T > unwarppedKSpace(RO, E1, CHA);

#pragma omp for
                for (ii = 0; ii < num; ii++)
                {
                    size_t s = ii / N;
                    size_t n = ii - s*N;

                    // check whether the kspace is undersampled
                    bool undersampled = false;
                    for (size_t e1 = startE1; e1 <= endE1; e1++)
                    {
                        if ((std::abs(kspace(RO / 2, e1, CHA - 1, n, s)) == 0)
                            && (std::abs(kspace(RO / 2, e1, 0, n, s)) == 0))
                        {
                            undersampled = true;
                            break;
                        }
                    }

                    T* pKpaceShifted = &(kspace_Shifted(0, 0, 0, n, s));
                    T* pRes = &(res(0, 0, 0, n, s));

                    if (!undersampled)
                    {
                        memcpy(pRes, pKpaceShifted, sizeof(T)*RO*E1*CHA);
                        continue;
                    }

                    long long kernelN = n;
                    if (kernelN >= (long long)ref_N) kernelN = (long long)ref_N - 1;

                    long long kernelS = s;
                    if (kernelS >= (long long)ref_S) kernelS = (long long)ref_S - 1;

                    boost::shared_ptr< hoNDArray< T > > acq(new hoNDArray< T >(RO, E1, CHA, pKpaceShifted));
                    spirit.set_acquired_points(*acq);

                    boost::shared_ptr< hoNDArray<T> > initialAcq;
                    if ( hasInitial )
                    {
                        initialAcq = boost::shared_ptr< hoNDArray<T> >(new hoNDArray<T>(RO, E1, CHA, &kspace_initial_Shifted(0, 0, 0, n, s)));
                        cgSolver.set_x0(initialAcq);
                    }
                    else
                    {
                        cgSolver.set_x0(acq);
                    }

                    if (ref_N == 1 && ref_S == 1)
                    {
                        spirit.compute_righ_hand_side(*acq, b);
                        cgSolver.solve(&unwarppedKSpace, &b);
                    }
                    else
                    {
                        T* pKer = &(ker_Shifted(0, 0, 0, kernelN, kernelS));
                        boost::shared_ptr<hoNDArray< T > > ker(new hoNDArray< T >(RO, E1, CHA, CHA, pKer));
                        spirit.set_forward_kernel(*ker, false);

                        spirit.compute_righ_hand_side(*acq, b);
                        cgSolver.solve(&unwarppedKSpace, &b);
                    }

                    // restore the acquired points
                    spirit.restore_acquired_kspace(*acq, unwarppedKSpace);
                    memcpy(pRes, unwarppedKSpace.begin(), unwarppedKSpace.get_number_of_bytes());
                }
            }

            Gadgetron::hoNDFFT<typename realType<T>::Type>::instance()->fftshift2D(res, kspace_Shifted);
            res = kspace_Shifted;
        }
        catch(...)
        {
            GADGET_THROW("Exceptions happened in perform_spirit_recon_linear_2DT(...) ... ");
        }
    }

    template EXPORTCMR void perform_spirit_recon_linear_2DT(const Gadgetron::hoNDArray< std::complex<float> >& kspace, size_t startE1, size_t endE1, const Gadgetron::hoNDArray< std::complex<float> >& kerIm, const Gadgetron::hoNDArray< std::complex<float> >& kspaceInitial, Gadgetron::hoNDArray< std::complex<float> >& res, size_t iter_max, double iter_thres, bool print_iter);
    template EXPORTCMR void perform_spirit_recon_linear_2DT(const Gadgetron::hoNDArray< std::complex<double> >& kspace, size_t startE1, size_t endE1, const Gadgetron::hoNDArray< std::complex<double> >& kerIm, const Gadgetron::hoNDArray< std::complex<double> >& kspaceInitial, Gadgetron::hoNDArray< std::complex<double> >& res, size_t iter_max, double iter_thres, bool print_iter);

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

    template <typename T>
    class sCB : public hoGdSolverCallBack< hoNDArray< T >, hoWavelet2DTOperator< T > >
    {
    public:
        typedef hoGdSolverCallBack< hoNDArray<T>, hoWavelet2DTOperator<T> > BaseClass;

        sCB() : BaseClass() {}
        virtual ~sCB() {}

        void execute(const hoNDArray<T>& b, hoNDArray<T>& x)
        {
            typedef hoSPIRIT2DTDataFidelityOperator<T> SpiritOperType;
            SpiritOperType* pOper = dynamic_cast<SpiritOperType*> (this->solver_->oper_system_);
            pOper->restore_acquired_kspace(x);
        }
    };

    template <typename T>
    void perform_spirit_recon_non_linear_2DT(const Gadgetron::hoNDArray<T>& kspace, const Gadgetron::hoNDArray<T>& kerIm,
                                            const Gadgetron::hoNDArray<T>& coil_map, const Gadgetron::hoNDArray<T>& kspaceLinear, Gadgetron::hoNDArray<T>& res,
                                            size_t iter_max, double iter_thres, double data_fidelity_lamda, double image_reg_lamda, double reg_N_weighting_ratio,
                                            bool reg_use_coil_sen_map, bool reg_with_approx_coeff, const std::string& wav_name, bool print_iter)
    {
        size_t RO = kspace.get_size(0);
        size_t E1 = kspace.get_size(1);
        size_t CHA = kspace.get_size(2);
        size_t N = kspace.get_size(3);
        size_t S = kspace.get_size(4);

        size_t ref_N = kerIm.get_size(4);
        size_t ref_S = kerIm.get_size(5);

        res = kspace;

        size_t s;

        for (s=0; s<S; s++)
        {
            boost::shared_ptr< hoNDArray< T > > coilMap;

            bool hasCoilMap = false;
            if (coil_map.get_size(0) == RO && coil_map.get_size(1) == E1 && coil_map.get_size(2)==CHA)
            {
                if (ref_N < N)
                {
                    coilMap = boost::shared_ptr< hoNDArray< T > >(new hoNDArray< T >(RO, E1, CHA, const_cast<T*>(coil_map.begin()) ));
                }
                else
                {
                    coilMap = boost::shared_ptr< hoNDArray< T > >(new hoNDArray< T >(RO, E1, CHA, ref_N, const_cast<T*>(coil_map.begin()) ));
                }

                hasCoilMap = true;
            }

            size_t s_used = s;
            if(s_used>=ref_S) s_used = ref_S;

            boost::shared_ptr<hoNDArray< T > > ker(new hoNDArray< T >(RO, E1, CHA, CHA, ref_N, const_cast<T*>(kerIm.begin())+s_used*RO*E1*CHA*CHA*ref_N) );
            boost::shared_ptr<hoNDArray< T > > acq(new hoNDArray< T >(RO, E1, CHA, N, const_cast<T*>(kspace.begin())+s*RO*E1*CHA*N) );
            hoNDArray< T > kspaceInitial(RO, E1, CHA, N, const_cast<T*>(kspaceLinear.begin())+s*RO*E1*CHA*N );
            hoNDArray< T > res2DT(RO, E1, CHA, N, res.begin()+s*RO*E1*CHA*N );

            if (data_fidelity_lamda > 0)
            {
                GDEBUG_STREAM("Start the NonLinear SPIRIT data fidelity iteration - regularization strength : "
                    << image_reg_lamda
                    << " - number of iteration : "                      << iter_max
                    << " - proximity across cha : "                     << false
                    << " - redundant dimension weighting ratio : "      << reg_N_weighting_ratio
                    << " - using coil sen map : "                       << reg_use_coil_sen_map
                    << " - with approx coeff : "                        << reg_with_approx_coeff
                    << " - iter thres : "                               << iter_thres
                    << " - wavelet name : "                             << wav_name
                    );

                typedef hoGdSolver< hoNDArray< T >, hoWavelet2DTOperator< T > > SolverType;
                SolverType solver;
                solver.iterations_ = iter_max;
                solver.set_output_mode(print_iter ? SolverType::OUTPUT_VERBOSE : SolverType::OUTPUT_SILENT);
                solver.grad_thres_ = iter_thres;
                solver.proximal_strength_ratio_ = image_reg_lamda;

                boost::shared_ptr< hoNDArray< T > > x0 = boost::make_shared< hoNDArray< T > >(kspaceInitial);
                solver.set_x0(x0);

                // parallel imaging term
                std::vector<size_t> dims;
                acq->get_dimensions(dims);
                hoSPIRIT2DTDataFidelityOperator< T > spirit(&dims);
                spirit.set_forward_kernel(*ker, false);
                spirit.set_acquired_points(*acq);

                // image reg term
                hoWavelet2DTOperator< T > wav3DOperator(&dims);
                wav3DOperator.set_acquired_points(*acq);
                wav3DOperator.scale_factor_first_dimension_ = 1;
                wav3DOperator.scale_factor_second_dimension_ = 1;
                wav3DOperator.scale_factor_third_dimension_ = reg_N_weighting_ratio;
                wav3DOperator.with_approx_coeff_ = reg_with_approx_coeff;
                wav3DOperator.change_coeffcients_third_dimension_boundary_ = true;
                wav3DOperator.proximity_across_cha_ = false;
                wav3DOperator.no_null_space_ = true;
                wav3DOperator.input_in_kspace_ = true;
                wav3DOperator.select_wavelet(wav_name);

                if (reg_use_coil_sen_map && hasCoilMap)
                {
                    wav3DOperator.coil_map_ = *coilMap;
                }

                // set operators

                solver.oper_system_ = &spirit;
                solver.oper_reg_ = &wav3DOperator;

                solver.solve(*acq, res2DT);
            }
            else
            {
                GDEBUG_STREAM("Start the NonLinear SPIRIT iteration with regularization strength : "
                    << image_reg_lamda
                    << " - number of iteration : " << iter_max
                    << " - proximity across cha : " << false
                    << " - redundant dimension weighting ratio : " << reg_N_weighting_ratio
                    << " - using coil sen map : " << reg_use_coil_sen_map
                    << " - with approx coeff : "  << reg_with_approx_coeff
                    << " - iter thres : " << iter_thres
                    << " - wavelet name : " << wav_name
                    );

                typedef hoGdSolver< hoNDArray< T >, hoWavelet2DTOperator< T > > SolverType;
                SolverType solver;
                solver.iterations_ = iter_max;
                solver.set_output_mode(print_iter ? SolverType::OUTPUT_VERBOSE : SolverType::OUTPUT_SILENT);
                solver.grad_thres_ = iter_thres;
                solver.proximal_strength_ratio_ = image_reg_lamda;

                boost::shared_ptr< hoNDArray< T > > x0 = boost::make_shared< hoNDArray< T > >(kspaceInitial);
                solver.set_x0(x0);

                // parallel imaging term
                std::vector<size_t> dims;
                acq->get_dimensions(dims);

                hoSPIRIT2DTOperator< T > spirit(&dims);
                spirit.set_forward_kernel(*ker, false);
                spirit.set_acquired_points(*acq);
                spirit.no_null_space_ = true;
                spirit.use_non_centered_fft_ = false;

                // image reg term
                std::vector<size_t> dim;
                acq->get_dimensions(dim);

                hoWavelet2DTOperator< T > wav3DOperator(&dim);
                wav3DOperator.set_acquired_points(*acq);
                wav3DOperator.scale_factor_first_dimension_ = 1;
                wav3DOperator.scale_factor_second_dimension_ = 1;
                wav3DOperator.scale_factor_third_dimension_ = reg_N_weighting_ratio;
                wav3DOperator.with_approx_coeff_ = reg_with_approx_coeff;
                wav3DOperator.change_coeffcients_third_dimension_boundary_ = true;
                wav3DOperator.proximity_across_cha_ = false;
                wav3DOperator.no_null_space_ = true;
                wav3DOperator.input_in_kspace_ = true;
                wav3DOperator.select_wavelet(wav_name);

                if (reg_use_coil_sen_map && hasCoilMap)
                {
                    wav3DOperator.coil_map_ = *coilMap;
                }

                // set operators
                solver.oper_system_ = &spirit;
                solver.oper_reg_ = &wav3DOperator;

                // set call back
                sCB<T> cb;
                cb.solver_ = &solver;
                solver.call_back_ = &cb;

                hoNDArray< T > b(kspaceInitial);
                Gadgetron::clear(b);

                solver.solve(b, res2DT);
                // if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(res2DT, debug_folder_full_path_ + "spirit_nl_2DT_res");

                spirit.restore_acquired_kspace(*acq, res2DT);

                // if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(res2DT, debug_folder_full_path_ + "spirit_nl_2DT_res_restored");
            }
        }
    }

    template EXPORTCMR void perform_spirit_recon_non_linear_2DT(const Gadgetron::hoNDArray< std::complex<float> >& kspace, const Gadgetron::hoNDArray< std::complex<float> >& kerIm,
        const Gadgetron::hoNDArray< std::complex<float> >& coilMap, const Gadgetron::hoNDArray< std::complex<float> >& kspaceInitial, Gadgetron::hoNDArray< std::complex<float> >& res,
        size_t iter_max, double iter_thres, double data_fidelity_lamda, double image_reg_lamda, double reg_N_weighting_ratio,
        bool reg_use_coil_sen_map, bool reg_with_approx_coeff, const std::string& wav_name, bool print_iter);

    template EXPORTCMR void perform_spirit_recon_non_linear_2DT(const Gadgetron::hoNDArray< std::complex<double> >& kspace, const Gadgetron::hoNDArray< std::complex<double> >& kerIm,
        const Gadgetron::hoNDArray< std::complex<double> >& coilMap, const Gadgetron::hoNDArray< std::complex<double> >& kspaceInitial, Gadgetron::hoNDArray< std::complex<double> >& res,
        size_t iter_max, double iter_thres, double data_fidelity_lamda, double image_reg_lamda, double reg_N_weighting_ratio,
        bool reg_use_coil_sen_map, bool reg_with_approx_coeff, const std::string& wav_name, bool print_iter);
}