#include <iostream>
#include <fstream>
#include <regex>

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/statistics/variance.hpp>

#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/serialization/vector.hpp>
#include <boost/filesystem.hpp>
#include <boost/filesystem/fstream.hpp>

#include <boost/spirit/include/qi.hpp>
#include <boost/spirit/include/phoenix.hpp>
#include <boost/fusion/include/std_pair.hpp>
#include <boost/iterator/iterator_adaptor.hpp>

#include <tbb/tbb.h>

#include <Eigen/StdVector>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/highgui/highgui_c.h>

#include <pupiltracker/pupiltracker.h>

#include <singleeyefitter/singleeyefitter.h>
#include <singleeyefitter/fun.h>
#include <singleeyefitter/projection.h>

#include "clipper.hpp"

using namespace singleeyefitter;

struct fitEyeModel_ret {
    Sphere<double> est_eye, est_eye_lm, est_eye_contrast;
    std::vector<Circle3D<double>> est_pupils, est_pupils2, est_pupils_lm, est_pupils_contrast;
};

namespace detail {
    using namespace boost::spirit;

    template<class T> struct as_qi;
    template<> struct as_qi<int> { static const int_type& get() {return int_;} };
    template<> struct as_qi<double> { static const double_type& get() {return double_;} };

    template<class To>
    struct parse_helper {
        template<class From>
        static To parse(const From& src) {
            using std::begin;
            using std::end;
            To ret;
            bool success = qi::parse(begin(src), end(src), as_qi<To>::get(), ret);
            if (!success) {
                throw std::runtime_error("Parse int failed");
            }
            return ret;
        }

        static To parse(char* src) {
            To ret;
            bool success = qi::parse(src, src + strlen(src), as_qi<To>::get(), ret);
            if (!success) {
                throw std::runtime_error("Parse int failed");
            }
            return ret;
        }

        static To parse(const char* src) {
            To ret;
            bool success = qi::parse(src, src + strlen(src), as_qi<To>::get(), ret);
            if (!success) {
                throw std::runtime_error("Parse int failed");
            }
            return ret;
        }
    };
}

template<typename T>
int parse_int(T&& src) {
    return ::detail::parse_helper<int>::parse(std::forward<T>(src));
}
template<typename T>
double parse_double(T&& src) {
    return ::detail::parse_helper<double>::parse(std::forward<T>(src));
}

cv::Point2f toImgCoord(const cv::Point2f& point, const cv::Mat& m, double scale = 1, int shift = 0) {
    return cv::Point2f(static_cast<float>((m.cols/2 + scale*point.x) * (1<<shift)),
        static_cast<float>((m.rows/2 + scale*point.y) * (1<<shift)));
}
cv::Point toImgCoord(const cv::Point& point, const cv::Mat& m, double scale = 1, int shift = 0) {
    return cv::Point(static_cast<int>((m.cols/2 + scale*point.x) * (1<<shift)),
        static_cast<int>((m.rows/2 + scale*point.y) * (1<<shift)));
}
cv::RotatedRect toImgCoord(const cv::RotatedRect& rect, const cv::Mat& m, float scale = 1) {
    return cv::RotatedRect(toImgCoord(rect.center,m,scale),
        cv::Size2f(scale*rect.size.width,
        scale*rect.size.height),
        rect.angle);
}

namespace boost {
    namespace serialization {

        template<class Archive, class T>
        inline void serialize(Archive & ar, Ellipse2D<T> & g, const unsigned int version)
        {
            ar & g.centre;
            ar & g.major_radius;
            ar & g.minor_radius;
            ar & g.angle;
        }
        template<class Archive, typename _Scalar, int _Rows, int _Cols, int _Options, int _MaxRows, int _MaxCols>
        inline void serialize(Archive & ar,
                              Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows, _MaxCols> & t,
                              const unsigned int file_version)
        {
            int rows = t.rows(), cols = t.cols();
            ar & rows;
            ar & cols;
            if( rows * cols != t.size() )
                t.resize( rows, cols );

            for(int i=0; i<t.size(); i++)
                ar & t.data()[i];
        }
        template<class Archive, class T>
        inline void serialize(Archive & ar, cv::Point_<T> & g, const unsigned int version)
        {
            ar & g.x;
            ar & g.y;
        }

    } // namespace serializationstd::vector<Eigen::Vector2d>
} // namespace boost

template <typename T>
std::vector<size_t> sort_indexes_impl(const T &v, std::random_access_iterator_tag)
{
    static_assert(std::is_same<typename std::iterator_traits<typename T::iterator>::iterator_category, std::random_access_iterator_tag>::value, "Type matches tag");

    // initialize original index locations
    auto idx = fun::range_<std::vector<size_t>>(v.size());

    // sort indexes based on comparing values in v
    sort(std::begin(idx), std::end(idx), [&](size_t i1, size_t i2){ return v[i1] < v[i2]; });

    return idx;
}

template <typename T>
std::vector<size_t> sort_indexes(const T &v)
{
    return sort_indexes_impl(v, typename std::iterator_traits<typename T::iterator>::iterator_category());
}

template< typename T >
void reorder(std::vector<T>& vals, const std::vector<size_t>& idxs)  {
    vals = fun::map([&](size_t i){ return vals[i]; }, idxs);
}

struct PupilGroundTruth {
    Eigen::Vector3d gaze_vector;
    std::vector<Eigen::Vector2d> outline;
    PupilGroundTruth() {}
    PupilGroundTruth(
        Eigen::Vector3d gaze_vector,
        std::vector<Eigen::Vector2d> outline) :
            gaze_vector(std::move(gaze_vector)),
            outline(std::move(outline)) {}
};

double calcEllipseTruthOverlap(const Ellipse2D<double>& ellipse, const std::vector<Eigen::Vector2d>& truth_poly) {
    ClipperLib::Path truth_path;
    ClipperLib::Path ellipse_path;

    int N = truth_poly.size();
    for(int i = 0; i < N; ++i) {
        auto p = truth_poly[i];

        double theta = (double)i / N * 2 * PI;
        auto p2 = pointAlongEllipse(ellipse, -theta);

        truth_path.emplace_back(static_cast<int>(std::floor(p.x()*100+0.5)),
                                static_cast<int>(std::floor(p.y()*100+0.5)));
        ellipse_path.emplace_back(static_cast<int>(std::floor(p2.x()*100+0.5)),
                                  static_cast<int>(std::floor(p2.y()*100+0.5)));
    };

    ClipperLib::Clipper clpr;
    clpr.AddPolygon(truth_path, ClipperLib::ptSubject);
    clpr.AddPolygon(ellipse_path, ClipperLib::ptClip);

    ClipperLib::Paths intersection_paths, union_paths;
    clpr.Execute(ClipperLib::ctIntersection, intersection_paths);
    clpr.Execute(ClipperLib::ctUnion, union_paths);

    double intersection_area = fun::sum([](const ClipperLib::Path& path){ return ClipperLib::Area(path); }, intersection_paths);
    double union_area = fun::sum([](const ClipperLib::Path& path){ return ClipperLib::Area(path); }, union_paths);

    return intersection_area / union_area;
}

void writeResults(const std::string& filename,
                  const std::vector<int> ids,
                  const std::vector<EyeModelFitter::Pupil>& null_pupils,
                  const std::vector<EyeModelFitter::Pupil>& simple_pupils,
                  const std::vector<EyeModelFitter::Pupil>& contrast_pupils,
                  const std::vector<EyeModelFitter::Pupil>& edge_pupils,
                  const std::map<int, PupilGroundTruth>& true_pupils,
                  double focal_length) {
    std::ofstream of(filename);

    for (int i = 0; i < ids.size(); ++i) {
        const auto& pupil = null_pupils[i].circle;
        const auto& pupil2 = simple_pupils[i].circle;
        const auto& pupil_dlib = contrast_pupils[i].circle;
        const auto& pupil_lm = edge_pupils[i].circle;
        const auto& id = ids[i];
        auto true_pupil_it = true_pupils.find(id);
        if (true_pupil_it == true_pupils.end())
            continue;

        const auto& true_pupil = true_pupil_it->second;

        double anglediff = acos(pupil.normal.dot(true_pupil.gaze_vector))*180/PI;
        double anglediff2 = acos(pupil2.normal.dot(true_pupil.gaze_vector))*180/PI;
        double anglediff_dlib = acos(pupil_dlib.normal.dot(true_pupil.gaze_vector))*180/PI;
        double anglediff_lm = acos(pupil_lm.normal.dot(true_pupil.gaze_vector))*180/PI;

        auto true_pupil_outline = true_pupil.outline;

        double ellipdist = calcEllipseTruthOverlap(Ellipse2D<double>(project(pupil, focal_length)), true_pupil_outline);
        double ellipdist2 = calcEllipseTruthOverlap(Ellipse2D<double>(project(pupil2, focal_length)), true_pupil_outline);
        double ellipdist_dlib = calcEllipseTruthOverlap(Ellipse2D<double>(project(pupil_dlib, focal_length)), true_pupil_outline);
        double ellipdist_lm = calcEllipseTruthOverlap(Ellipse2D<double>(project(pupil_lm, focal_length)), true_pupil_outline);

        of << id;
        of << "\t" << anglediff << "\t" << anglediff2 << "\t" << anglediff_dlib << "\t" << anglediff_lm;
        of << "\t" << ellipdist << "\t" << ellipdist2 << "\t" << ellipdist_dlib << "\t" << ellipdist_lm;
        of << std::endl;
    }
}

std::string regex_escape(const std::string& pattern) {
    return regex_replace(pattern,
        std::regex("[\\^\\.\\$\\|\\(\\)\\[\\]\\*\\+\\?\\/\\\\]"),
        std::string("\\$&"));
}

struct Key {
    wchar_t code;
    bool shift;
    bool ctrl;
    bool meta;
};

#ifdef WIN32

Key cvxWaitKey(int delay) {
    int code = cv::waitKey(delay);
    Key key;
    key.code = code;
    key.shift = (GetKeyState(VK_SHIFT) & 0x8000) != 0;
    key.ctrl = (GetKeyState(VK_CONTROL) & 0x8000) != 0;
    key.meta = (GetKeyState(VK_MENU) & 0x8000) != 0;
    return key;
}

#else

Key cvxWaitKey(int delay) {
    int code = cv::waitKey(delay);
    Key key;
    if (code >= 'A' && code <= 'Z') {
        key.code = code + 'a' - 'A';
        key.shift = true;
    }
    else {
        key.code = code;
        key.shift = false;
    }
    key.ctrl = false;
    key.meta = false;
    return key;
}

#endif

int main(int argc, char* argv[])
{
    using boost::math::sign;

    if (argc < 4) {
        std::cerr << "USAGE: " << argv[0] << " <folder> <filepattern> <fov> [<truthfile> [<init_num_added>]]" << std::endl;
        return 1;
    }

    //try
    {
        namespace fs = boost::filesystem;

        std::vector<int> ids;
        std::vector<cv::Mat> obs_eye_images;
        std::vector<Ellipse2D<double>> obs_pupil_ellipses;
        std::vector<std::vector<cv::Point2f>> obs_pupil_inliers;

        fs::path imagedir(argv[1]);
        std::string filepattern(argv[2]);

        std::smatch filepatternisregex_match;
        if (regex_match(filepattern, filepatternisregex_match, std::regex("/(.*)/"))) {
            filepattern = filepatternisregex_match[1].str();
        } else {
            filepattern = regex_escape(filepattern);
            // Converts #### into (\d\d\d\d\d*)
            filepattern = regex_replace(filepattern, std::regex("#+"), std::string("($&\\d*)"));
            filepattern = regex_replace(filepattern, std::regex("#"), std::string("\\d"));
        }
        std::regex file_regex(filepattern);

        tbb::mutex push_lock;

        tbb::parallel_for_each(fs::directory_iterator(imagedir), fs::directory_iterator(), [&] (fs::path path) {
            std::smatch path_match;
            std::string path_str = path.filename().string();

            if (is_regular_file(path) && regex_match(path_str, path_match, file_regex)) {
                int id = parse_int(path_match[1].str());
                cv::Mat eye = cv::imread(path.string(), CV_LOAD_IMAGE_GRAYSCALE);
                Ellipse2D<double> el;
                std::vector<cv::Point2f> inlier_pts;

                bool cache_valid = false;
                const int CACHE_VERSION = 8;

                fs::path cache_path = path.parent_path() / fs::path(path.stem().string() + ".cache");

                if (exists(cache_path)) {
                    std::ifstream ifs(cache_path.string());
                    boost::archive::text_iarchive ia(ifs);

                    int ia_cache_version;
                    ia >> ia_cache_version;
                    if (CACHE_VERSION == ia_cache_version) {
                        std::cout << "Loading " << cache_path.string() << std::endl;

                        ia >> el;
                        ia >> inlier_pts;

                        cache_valid = true;
                    }
                }

                if (!cache_valid)
                {
                    std::cout << "Pupil tracking " << path.string() << std::endl;

                    pupiltracker::tracker_log log;

                    pupiltracker::TrackerParams pupil_tracker_params;
                    pupil_tracker_params.Radius_Min = 20;
                    pupil_tracker_params.Radius_Max = 70;
                    pupil_tracker_params.CannyThreshold1 = 20;
                    pupil_tracker_params.CannyThreshold2 = 40;
                    pupil_tracker_params.CannyBlur = 1.6;
                    pupil_tracker_params.EarlyRejection = true;
                    pupil_tracker_params.EarlyTerminationPercentage = 95;
                    pupil_tracker_params.PercentageInliers = 20;
                    pupil_tracker_params.InlierIterations = 2;
                    pupil_tracker_params.ImageAwareSupport = true;
                    pupil_tracker_params.StarburstPoints = 0;
                    //pupil_tracker_params.Seed = 0;

                    pupiltracker::findPupilEllipse_out pupil_tracker_out;
                    bool found = pupiltracker::findPupilEllipse(pupil_tracker_params, eye, pupil_tracker_out, log);

                    if (found) {
                        el = toEllipse<double>(pupil_tracker_out.elPupil);
                        el.centre -= Eigen::Vector2d(eye.cols, eye.rows)/2;

                        for (auto&& inlier : pupil_tracker_out.inliers) {
                            inlier_pts.push_back(cv::Point2f(pupil_tracker_out.roiPupil.x + inlier.x - eye.cols/2, pupil_tracker_out.roiPupil.y + inlier.y - eye.rows/2));
                        }
                    } else {
                        el = Ellipse2D<double>::Null;
                    }

                    std::ofstream ofs(cache_path.string());
                    boost::archive::text_oarchive oa(ofs);

                    oa << CACHE_VERSION;
                    oa << el;
                    oa << inlier_pts;
                }

                {
                    tbb::mutex::scoped_lock push_guard(push_lock);

                    if (el) {
                        ids.push_back(id);
                        obs_eye_images.push_back(eye);
                        obs_pupil_ellipses.push_back(el);
                        obs_pupil_inliers.push_back(std::move(inlier_pts));
                    }
                }
            }
        });

        auto&& idx_sort = sort_indexes(ids);
        reorder(ids, idx_sort);
        reorder(obs_eye_images, idx_sort);
        reorder(obs_pupil_ellipses, idx_sort);
        reorder(obs_pupil_inliers, idx_sort);

        double fov = parse_double(argv[3]);
        fov *= PI / 180;
        double focal_length = (obs_eye_images[0].cols / 2) / std::tan(fov/2);

        bool animate = false;

        // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
        // This is where the stuff actually happens
        // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

        EyeModelFitter null_fitter(focal_length, 5, 0.5);
        EyeModelFitter simple_fitter(focal_length, 5, 0.5);
        EyeModelFitter contrast_fitter(focal_length, 5, 0.5);
        EyeModelFitter edge_fitter(focal_length, 5, 0.5);

        auto N = obs_eye_images.size();
        if (argc > 5) {
            int n = parse_int(argv[5]);
            if (n >= 0 && n < N) {
                N = n;
            }
        }

        for (int i = 0; i < N; ++i) {
            null_fitter.add_observation(obs_eye_images[i], obs_pupil_ellipses[i], obs_pupil_inliers[i]);
            simple_fitter.add_observation(obs_eye_images[i], obs_pupil_ellipses[i], obs_pupil_inliers[i]);
            contrast_fitter.add_observation(obs_eye_images[i], obs_pupil_ellipses[i], obs_pupil_inliers[i]);
            edge_fitter.add_observation(obs_eye_images[i], obs_pupil_ellipses[i], obs_pupil_inliers[i]);
        }

        null_fitter.unproject_observations();

        simple_fitter.unproject_observations();
        simple_fitter.initialise_model();

        contrast_fitter.unproject_observations();
        contrast_fitter.initialise_model();

        edge_fitter.unproject_observations();
        edge_fitter.initialise_model();

        Sphere<double> true_eye;
        std::map<int, PupilGroundTruth> true_pupils;

        if (argc > 4) {
            using namespace boost::spirit;
            using boost::phoenix::at_c;

            std::ifstream ground_truth_file((imagedir / argv[4]).string());
            if (ground_truth_file.is_open()) {
                std::string str;
                std::getline(ground_truth_file, str);
                while(ground_truth_file.good()) {
                    int id;
                    PupilGroundTruth true_pupil;

                    /* One day I will understand boost qi enough to do something like this:

                    qi::rule<decltype(begin(str)), Eigen::Vector2d(), qi::ascii::space_type> vec2d =
                        qi::eps[_val = Eigen::Vector2d()] >>
                        qi::double_[at_c<0>(_val) = _1] >> ',' >>
                        qi::double_[at_c<1>(_val) = _1];
                    qi::rule<decltype(begin(str)), Eigen::Vector3d(), qi::ascii::space_type> vec3d =
                        qi::eps[qi::_val = Eigen::Vector3d()] >>
                        qi::double_[at_c<0>(_val) = _1] >> ',' >>
                        qi::double_[at_c<1>(_val) = _1] >> ',' >>
                        qi::double_[at_c<2>(_val) = _1];

                    But it is not that day */

                    std::vector<std::pair<double,double>> outlinepoints;

                    int parse_success = qi::phrase_parse(begin(str), end(str),
                        int_ >> '|' >>
                        double_ >> ',' >> double_ >> ',' >> double_ >> '|' >>
                        double_ >> '|' >>
                        double_ >> ',' >> double_ >> ',' >> double_ >> '|' >>
                        (double_ >> ',' >> double_) % ',',
                        ascii::space,
                        id,
                        true_eye.centre[0],true_eye.centre[1],true_eye.centre[2],
                        true_eye.radius,
                        true_pupil.gaze_vector[0],true_pupil.gaze_vector[1],true_pupil.gaze_vector[2],
                        outlinepoints
                    );
                    if (parse_success) {
                        true_eye.centre[1] = -true_eye.centre[1];
                        true_eye.centre[2] = -true_eye.centre[2];
                        true_pupil.gaze_vector[1] = -true_pupil.gaze_vector[1];
                        true_pupil.gaze_vector[2] = -true_pupil.gaze_vector[2];
                        if (std::find(begin(ids), end(ids), id) != ids.end()) {
                            for (const auto& p : outlinepoints) {
                                true_pupil.outline.emplace_back(p.first, -p.second);
                            }
                            true_pupils[id] = true_pupil;
                        }

                    } else {
                        std::cout << "Failed to parse ground truth: " << str << std::endl;
                    }

                    std::getline(ground_truth_file, str);
                }
            }
        }

        double displayscale = 1;
        bool do_display = true;
        int curr_i = 0;

        namespace acc=boost::accumulators;
        acc::accumulator_set<double, acc::stats<acc::tag::mean, acc::tag::variance>> acc_ellipdist, acc_ellipdist2, acc_ellipdist_dlib, acc_ellipdist_lm;
        bool recalc_ellipdist = true;

        bool display_obs_pupil_ellipses = true;
        bool display_est_pupil_ellipses = true;
        bool display_est_pupil_ellipses_contrast = true;
        bool display_est_pupil_ellipses_lm = true;
        bool display_true_pupil_ellipses = true;
        bool display_intersection_ellip = false;
        bool display_intersection_lines = false;

        bool ransac = true;
        bool use_smoothness = true;

        auto animate_func =
            [&] (const Sphere<double>& eye, const std::vector<Circle3D<double>>& pupils) {
                cv::Mat curr = obs_eye_images[curr_i];
                cv::Mat curr_disp = cvx::resize(cvx::cvtColor(cvx::convert(curr, CV_8U, 255),
                                                              cv::COLOR_GRAY2BGR),
                                                displayscale, 0, cv::INTER_CUBIC);

                cv::ellipse(curr_disp, toImgCoord(toRotatedRect(project(eye, focal_length)), curr_disp, displayscale), cvx::rgb(60,0,0), 1, CV_AA);
                for (const auto& pupil : pupils) {
                    if (pupil)
                        cv::ellipse(curr_disp, toImgCoord(toRotatedRect(Ellipse2D<double>(project(pupil, focal_length))), curr_disp, displayscale), cvx::rgb(60,60,0), 1, CV_AA);
                }
                if (curr_i < pupils.size() && pupils[curr_i]) {
                    cv::ellipse(curr_disp, toImgCoord(toRotatedRect(Ellipse2D<double>(project(pupils[curr_i], focal_length))), curr_disp, displayscale), cvx::rgb(60,0,0), 1, CV_AA);
                }
                cv::imshow("Current Eye", curr_disp);
                cv::waitKey(10);
            };

        //cv::namedWindow("Eye");
        //cv::createButton("Display observed ellipses", [](int state, void* userdata){*(bool*)userdata = state;}, &display_obs_pupil_ellipses, cv::QT_CHECKBOX, display_obs_pupil_ellipses);
        //cv::createButton("Display initial pupil model", [](int state, void* userdata){*(bool*)userdata = state;}, &display_est_pupil_ellipses, cv::QT_CHECKBOX, display_est_pupil_ellipses);
        //cv::createButton("Display contrast-optimized pupils", [](int state, void* userdata){*(bool*)userdata = state;}, &display_est_pupil_ellipses_contrast, cv::QT_CHECKBOX, display_est_pupil_ellipses_contrast);

        cv::VideoWriter recording_writer;
        while(do_display) {
            cv::Mat disp = cv::Mat::zeros(displayscale*obs_eye_images[0].rows, displayscale*obs_eye_images[0].cols, CV_8UC3);

            cv::Mat curr = obs_eye_images[curr_i];
            cv::Mat curr_disp = cvx::resize(cvx::cvtColor(curr,
                                                          cv::COLOR_GRAY2BGR),
                                            displayscale, 0, cv::INTER_CUBIC);

            cv::Mat curr_edge_disp = cv::Mat::zeros(curr_disp.rows, curr_disp.cols, CV_8UC3);

            for (const auto& inlier : obs_pupil_inliers[curr_i]) {
                int x = std::floor(curr_edge_disp.cols/2 + inlier.x * displayscale - 0.5);
                int y = std::floor(curr_edge_disp.rows/2 + inlier.y * displayscale - 0.5);
                if (x < 0 || x >= curr_edge_disp.cols || y < 0 || y >= curr_edge_disp.rows)
                    continue;

                curr_edge_disp.at<cv::Vec3b>(y,x) = cv::Vec3b::all(255);
            }

            if (display_true_pupil_ellipses) {
                if (true_eye) {
                    cv::ellipse(disp, toImgCoord(toRotatedRect(project(true_eye, focal_length)), disp, displayscale), cvx::rgb(60,60,60), -1, CV_AA);
                }
            }

            if (display_est_pupil_ellipses) {
                if (simple_fitter.eye) {
                    cv::ellipse(disp, toImgCoord(toRotatedRect(project(simple_fitter.eye, focal_length)), disp, displayscale), cvx::rgb(60,0,60), 1, CV_AA);
                    cv::ellipse(curr_disp, toImgCoord(toRotatedRect(project(simple_fitter.eye, focal_length)), curr_disp, displayscale), cvx::rgb(60,0,60), 1, CV_AA);
                }
            }
            if (display_est_pupil_ellipses_contrast) {
                if (contrast_fitter.eye) {
                    cv::ellipse(disp, toImgCoord(toRotatedRect(project(contrast_fitter.eye, focal_length)), disp, displayscale), cvx::rgb(0,60,0), 1, CV_AA);
                    cv::ellipse(curr_disp, toImgCoord(toRotatedRect(project(contrast_fitter.eye, focal_length)), curr_disp, displayscale), cvx::rgb(0,60,0), 1, CV_AA);
                }
            }
            if (display_est_pupil_ellipses_lm) {
                if (edge_fitter.eye) {
                    cv::ellipse(disp, toImgCoord(toRotatedRect(project(edge_fitter.eye, focal_length)), disp, displayscale), cvx::rgb(0,60,60), 1, CV_AA);
                    cv::ellipse(curr_disp, toImgCoord(toRotatedRect(project(edge_fitter.eye, focal_length)), curr_disp, displayscale), cvx::rgb(0,60,60), 1, CV_AA);
                }
            }
            if (display_true_pupil_ellipses) {
                if (true_eye) {
                    cv::ellipse(curr_disp, toImgCoord(toRotatedRect(project(true_eye, focal_length)), curr_disp, displayscale), cvx::rgb(60,0,0), 1, CV_AA);
                }
            }

            auto display_from = [&] (const std::vector<EyeModelFitter::Pupil>& pupils, double r, double g, double b) {
                for (const auto& pupil : pupils) {
                    if (pupil.circle)
                        cv::ellipse(disp, toImgCoord(toRotatedRect(Ellipse2D<double>(project(pupil.circle, focal_length))), disp, displayscale), cvx::rgb(r,g,b), 1, CV_AA);
                }
                if (curr_i < pupils.size() && pupils[curr_i].circle) {
                    cv::ellipse(curr_disp, toImgCoord(toRotatedRect(Ellipse2D<double>(project(pupils[curr_i].circle, focal_length))), curr_disp, displayscale), cvx::rgb(r,g,b), 1, CV_AA);
                    cv::ellipse(curr_edge_disp, toImgCoord(toRotatedRect(Ellipse2D<double>(project(pupils[curr_i].circle, focal_length))), curr_edge_disp, displayscale), cvx::rgb(r,g,b), 1, CV_AA);
                }
            };

            if (display_obs_pupil_ellipses) {
                display_from(null_fitter.pupils, 255,0,255);
            }
            if (display_est_pupil_ellipses) {
                display_from(simple_fitter.pupils, 255,255,0);
            }
            if (display_est_pupil_ellipses_contrast) {
                display_from(contrast_fitter.pupils, 0,255,0);
            }
            if (display_est_pupil_ellipses_lm) {
                display_from(edge_fitter.pupils, 0,255,255);
            }
            if (display_true_pupil_ellipses) {
                std::vector<std::vector<cv::Point>> pts;
                for (auto&& id : ids) {
                    const auto& true_pupils_id_it = true_pupils.find(id);
                    if (true_pupils_id_it != true_pupils.end()) {
                        const auto& true_pupil = true_pupils_id_it->second;

                        auto this_pts = fun::map([&](const Eigen::Vector2d& pt){
                            auto imgcoord = toImgCoord(cv::Point2f(pt.x(), pt.y()), disp, displayscale, 5);
                            return cv::Point(imgcoord.x, imgcoord.y);
                        }, true_pupil.outline);
                        pts.emplace_back(std::move(this_pts));
                    }
                }
                cv::polylines(disp, pts, true, cvx::rgb(255,0,0), 1, CV_AA, 5);

                const auto& true_pupils_curri_it = true_pupils.find(ids[curr_i]);
                if (true_pupils_curri_it != true_pupils.end()) {
                    const auto& true_pupil = true_pupils_curri_it->second;

                    auto this_pts = fun::map([&](const Eigen::Vector2d& pt){
                        auto imgcoord = toImgCoord(cv::Point2f(pt.x(), pt.y()), curr_disp, displayscale, 5);
                        return cv::Point(imgcoord.x, imgcoord.y);
                    }, true_pupil.outline);
                    cv::polylines(curr_disp, this_pts, true, cvx::rgb(255,0,0), 1, CV_AA, 5);
                    cv::polylines(curr_edge_disp, this_pts, true, cvx::rgb(255,0,0), 1, CV_AA, 5);
                }
            }

            double anglediff_i, anglediff2_i, anglediff_dlib_i, anglediff_lm_i;
            double ellipdist_i, ellipdist2_i, ellipdist_dlib_i, ellipdist_lm_i;

            acc::accumulator_set<double, acc::stats<acc::tag::mean, acc::tag::variance>> acc_anglediff, acc_anglediff2, acc_anglediff_dlib, acc_anglediff_lm;
            if (recalc_ellipdist) {
                acc_ellipdist = acc::accumulator_set<double, acc::stats<acc::tag::mean, acc::tag::variance>>();
                acc_ellipdist2 = acc::accumulator_set<double, acc::stats<acc::tag::mean, acc::tag::variance>>();
                acc_ellipdist_dlib = acc::accumulator_set<double, acc::stats<acc::tag::mean, acc::tag::variance>>();
                acc_ellipdist_lm = acc::accumulator_set<double, acc::stats<acc::tag::mean, acc::tag::variance>>();
            }
            for (int i = 0; i < N; ++i) {
                const EyeModelFitter::Circle& pupil = null_fitter.pupils[i].circle;
                const EyeModelFitter::Circle& pupil2 = simple_fitter.pupils[i].circle;
                const EyeModelFitter::Circle& pupil_dlib = contrast_fitter.pupils[i].circle;
                const EyeModelFitter::Circle& pupil_lm = edge_fitter.pupils[i].circle;
                const auto& id = ids[i];
                auto true_pupil_it = true_pupils.find(id);
                if (true_pupil_it == true_pupils.end())
                    continue;

                const auto& true_pupil = true_pupil_it->second;

                double anglediff = acos(pupil.normal.dot(true_pupil.gaze_vector))*180/PI;
                double anglediff2 = acos(pupil2.normal.dot(true_pupil.gaze_vector))*180/PI;
                double anglediff_dlib = acos(pupil_dlib.normal.dot(true_pupil.gaze_vector))*180/PI;
                double anglediff_lm = acos(pupil_lm.normal.dot(true_pupil.gaze_vector))*180/PI;

                if (i == curr_i) {
                    anglediff_i = anglediff;
                    anglediff2_i = anglediff2;
                    anglediff_dlib_i = anglediff_dlib;
                    anglediff_lm_i = anglediff_lm;
                }

                acc_anglediff(anglediff);
                acc_anglediff2(anglediff2);
                acc_anglediff_dlib(anglediff_dlib);
                acc_anglediff_lm(anglediff_lm);

                auto true_pupil_outline = true_pupil.outline;

                if (recalc_ellipdist || i == curr_i) {
                    double ellipdist = calcEllipseTruthOverlap(Ellipse2D<double>(project(pupil, focal_length)), true_pupil_outline);
                    double ellipdist2 = calcEllipseTruthOverlap(Ellipse2D<double>(project(pupil2, focal_length)), true_pupil_outline);
                    double ellipdist_dlib = calcEllipseTruthOverlap(Ellipse2D<double>(project(pupil_dlib, focal_length)), true_pupil_outline);
                    double ellipdist_lm = calcEllipseTruthOverlap(Ellipse2D<double>(project(pupil_lm, focal_length)), true_pupil_outline);

                    if (i == curr_i) {
                        ellipdist_i = ellipdist;
                        ellipdist2_i = ellipdist2;
                        ellipdist_dlib_i = ellipdist_dlib;
                        ellipdist_lm_i = ellipdist_lm;
                    }
                    if (recalc_ellipdist) {
                        acc_ellipdist(ellipdist);
                        acc_ellipdist2(ellipdist2);
                        acc_ellipdist_dlib(ellipdist_dlib);
                        acc_ellipdist_lm(ellipdist_lm);
                    }
                }
            }
            recalc_ellipdist = false;

            std::stringstream curr_i_ss;
            curr_i_ss << std::fixed << std::setprecision(4);
            curr_i_ss << "obs " << anglediff_i << ", ";
            curr_i_ss << "est " << anglediff2_i << ", ";
            curr_i_ss << "grd " << anglediff_dlib_i << ", ";
            curr_i_ss << "lm " << anglediff_lm_i << " | ";
            curr_i_ss << "obs " << ellipdist_i << ", ";
            curr_i_ss << "est " << ellipdist2_i << ", ";
            curr_i_ss << "grd " << ellipdist_dlib_i << ", ";
            curr_i_ss << "lm " << ellipdist_lm_i << " | ";
            curr_i_ss << "Frame " << curr_i << " (total " << obs_eye_images.size() << ")";

            std::stringstream ss;
            ss << std::fixed << std::setprecision(4);
            ss << "obs " << acc::mean(acc_anglediff) << " (" << acc::variance(acc_anglediff) << "), ";
            ss << "est " << acc::mean(acc_anglediff2) << " (" << acc::variance(acc_anglediff2) << "), ";
            ss << "grd " << acc::mean(acc_anglediff_dlib) << " (" << acc::variance(acc_anglediff_dlib) << "), ";
            ss << "lm " << acc::mean(acc_anglediff_lm) << " (" << acc::variance(acc_anglediff_lm) << ") | ";
            ss << "obs " << acc::mean(acc_ellipdist) << " (" << acc::variance(acc_ellipdist) << "), ";
            ss << "est " << acc::mean(acc_ellipdist2) << " (" << acc::variance(acc_ellipdist2) << "), ";
            ss << "grd " << acc::mean(acc_ellipdist_dlib) << " (" << acc::variance(acc_ellipdist_dlib) << "), ";
            ss << "lm " << acc::mean(acc_ellipdist_lm) << " (" << acc::variance(acc_ellipdist_lm) << ")";

            cv::imshow("Eye", disp);
            cv::imshow("Current Eye", curr_disp);
            cv::imshow("Current Eye Edges", curr_edge_disp);

            cv::displayStatusBar("Eye", ss.str());
            cv::displayStatusBar("Current Eye", curr_i_ss.str());

            if (recording_writer.isOpened()) {
                recording_writer << curr_disp;
            }

            if (display_intersection_ellip || display_intersection_lines) {
                cv::Mat disp_invalid_lines = cv::Mat::zeros(displayscale*obs_eye_images[0].rows, displayscale*obs_eye_images[0].cols, CV_32FC1);
                cv::Mat disp_valid_lines = cv::Mat::zeros(displayscale*obs_eye_images[0].rows, displayscale*obs_eye_images[0].cols, CV_32FC1);
                cv::Mat disp_ellipses = cv::Mat::zeros(displayscale*obs_eye_images[0].rows, displayscale*obs_eye_images[0].cols, CV_32FC1);
                for (const auto& pupil : null_fitter.pupils) {
                    Eigen::Matrix<double,2,1> c_proj = project(pupil.circle.centre, focal_length);
                    Eigen::Matrix<double,2,1> v_proj = project(pupil.circle.centre+pupil.circle.normal, focal_length) - c_proj;
                    Eigen::Matrix<double,2,1> start = c_proj - 500*v_proj;
                    Eigen::Matrix<double,2,1> end = c_proj + 500*v_proj;

                    if (display_intersection_lines) {

                        cv::Mat disp_line = cv::Mat::zeros(displayscale*obs_eye_images[0].rows, displayscale*obs_eye_images[0].cols, CV_8UC1);
                        cvx::line(disp_line,
                            cv::Point2f(disp_line.cols/2 + displayscale*start[0], disp_line.rows/2 + displayscale*start[1]),
                            cv::Point2f(disp_line.cols/2 + displayscale*end[0], disp_line.rows/2 + displayscale*end[1]),
                            cv::Scalar(25));

                        cv::Mat disp_line_f = cvx::convert(disp_line, CV_32FC1, 1./255, 0);

                        if (pupil.init_valid)
                            cv::accumulate((disp_line_f).mul(1 - disp_valid_lines), disp_valid_lines);
                        else
                            cv::accumulate((disp_line_f).mul(1 - disp_invalid_lines), disp_invalid_lines);
                    }

                    if (display_intersection_ellip) {
                        cv::Mat disp_ellipse = cv::Mat::zeros(displayscale*obs_eye_images[0].rows, displayscale*obs_eye_images[0].cols, CV_8UC1);
                        cv::ellipse(disp_ellipse, toImgCoord(toRotatedRect(Ellipse2D<double>(project(pupil.circle, focal_length))), disp, displayscale), cv::Scalar(50), 1, CV_AA);

                        cv::Mat disp_ellipse_f = cvx::convert(disp_ellipse, CV_32FC1, 1./255, 0);
                        cv::accumulate((disp_ellipse_f).mul(1 - disp_ellipses), disp_ellipses);
                    }
                }

                cv::Mat disp_intersection = cv::Mat(displayscale*obs_eye_images[0].rows, displayscale*obs_eye_images[0].cols, CV_32FC3, cvx::rgb(1,1,1));
                cv::accumulate(cvx::cvtColor(disp_invalid_lines, cv::COLOR_GRAY2BGR).mul(cvx::rgb(1,0,0) - disp_intersection), disp_intersection);
                cv::accumulate(cvx::cvtColor(disp_valid_lines, cv::COLOR_GRAY2BGR).mul(cvx::rgb(0,0,1) - disp_intersection), disp_intersection);
                cv::accumulate(cvx::cvtColor(disp_ellipses, cv::COLOR_GRAY2BGR).mul(cvx::rgb(1,0,1) - disp_intersection), disp_intersection);

                if (display_intersection_lines) {
                    auto crosspt = toImgCoord(toPoint2f(project(simple_fitter.eye.centre, focal_length)), disp, displayscale);
                    for (int i = -1; i <= 1; ++i)
                        for (int j = -1; j <= 1; ++j)
                            cvx::cross(disp_intersection, crosspt + cv::Point2f(i,j), 10, cvx::rgb(0,0,0), 3);
                    cvx::cross(disp_intersection, crosspt, 10, cvx::rgb(0,200,0), 3);
                }

                cv::imshow("Intersection", disp_intersection);
            }

            Key key = cvxWaitKey(-1);

            switch(key.code) {
            case 16: // Shift
                break;

            case '1':
                display_obs_pupil_ellipses = !display_obs_pupil_ellipses;
                break;
            case '2':
                display_est_pupil_ellipses = !display_est_pupil_ellipses;
                break;
            case '3':
                display_est_pupil_ellipses_contrast = !display_est_pupil_ellipses_contrast;
                break;
            case '4':
                display_est_pupil_ellipses_lm = !display_est_pupil_ellipses_lm;
                break;
            case '5':
                display_true_pupil_ellipses = !display_true_pupil_ellipses;
                break;

            case 27: // ESC
            case'q':
                do_display = false;
                break;
            case 'w':
                writeResults((imagedir / "result_analysis.txt").string(),
                    ids,
                    null_fitter.pupils,
                    simple_fitter.pupils,
                    contrast_fitter.pupils,
                    edge_fitter.pupils,
                    true_pupils,
                    focal_length);
                break;

            case '+':
            case '=':
                displayscale *= 1.1;
                break;
            case '-':
            case '_':
                displayscale /= 1.1;
                break;
            case 'i':
                display_intersection_ellip = !display_intersection_ellip;
                break;
            case 'o':
                display_intersection_lines = !display_intersection_lines;
                break;

            case 'a': {
                if (key.shift) {

                    animate = !animate;
                    if (animate) {
                        std::cout << "Animation enabled" << std::endl;
                    } else {
                        std::cout << "Animation disabled" << std::endl;
                    }

                } else {
                    int id;
                    if (N < obs_eye_images.size() - 1) {
                        id = null_fitter.add_observation(obs_eye_images[N], obs_pupil_ellipses[N], obs_pupil_inliers[N]);
                        null_fitter.unproject_single_observation(id);

                        id = simple_fitter.add_observation(obs_eye_images[N], obs_pupil_ellipses[N], obs_pupil_inliers[N]);
                        simple_fitter.unproject_single_observation(id);
                        simple_fitter.initialise_single_observation(id);

                        id = contrast_fitter.add_observation(obs_eye_images[N], obs_pupil_ellipses[N], obs_pupil_inliers[N]);
                        contrast_fitter.unproject_single_observation(id);
                        contrast_fitter.initialise_single_observation(id);
                        contrast_fitter.refine_single_with_contrast(id);

                        id = edge_fitter.add_observation(obs_eye_images[N], obs_pupil_ellipses[N], obs_pupil_inliers[N]);
                        edge_fitter.unproject_single_observation(id);
                        edge_fitter.initialise_single_observation(id);
                        //edge_fitter.refine_single_with_inliers(id);

                        curr_i = N;
                        N++;
                    }
                }
                break;
            }

            case 's':
                if (key.shift) {
                    for (int i = 0; i < N; ++i) {
                        std::cout << i << std::endl;
                        contrast_fitter.refine_single_with_contrast(i);
                    }
                    std::cout << "Done" << std::endl;
                }
                else {
                    use_smoothness = !use_smoothness;
                    if (use_smoothness) {
                        std::cout << "Smoothness enabled" << std::endl;
                    } else {
                        std::cout << "Smoothness disabled" << std::endl;
                    }
                }

                break;

            case 'd':
                if (key.shift) {
                    contrast_fitter.refine_single_with_contrast(curr_i);
                    contrast_fitter.print_single_contrast_metric(curr_i);
                }
                else {
                    if (animate)
                        contrast_fitter.refine_with_region_contrast(animate_func);
                    else
                        contrast_fitter.refine_with_region_contrast();
                    std::cout << "Done" << std::endl;
                }
                recalc_ellipdist = true;
                display_est_pupil_ellipses_contrast = true;

                break;

            case 'g':
                if (key.shift)
                {
                    std::string out_file = (imagedir / "contrast_sweep.txt").string();
                    std::ofstream out_stream(out_file);

                    auto start_params = contrast_fitter.pupils[curr_i].params;

                    for (auto&& theta : fun::linspace(0.0, PI, 100)) {
                        for (auto&& psi : fun::linspace(0.0, 2*PI, 100)) {
                            for (auto&& radius : fun::range_<std::vector<double>>(1, 10, 0.1)) {
                                EyeModelFitter::PupilParams params(theta, psi, radius);
                                contrast_fitter.pupils[curr_i].params = params;
                                std::cout << theta << " " << psi << " " << radius;
                                auto val = contrast_fitter.single_contrast_metric(curr_i);
                                std::cout << " -> " << val;

                                out_stream << theta << " " << psi << " " << radius << " " << val;
                            }
                        }
                    }

                    contrast_fitter.pupils[curr_i].params = start_params;

                }
                else
                {
                    auto start_params = contrast_fitter.pupils[curr_i].params;

                    auto best_params = start_params;
                    double best_contrast = contrast_fitter.single_contrast_metric(curr_i);

                    for (auto&& theta : fun::linspace(start_params.theta - 0.1, start_params.theta + 0.1, 6)) {
                        for (auto&& psi : fun::linspace(start_params.psi - 0.1, start_params.psi + 0.1, 6)) {
                            for (auto&& radius : fun::linspace(start_params.radius / 1.5, start_params.radius * 1.5, 6)) {
                                EyeModelFitter::PupilParams params(theta, psi, radius);
                                contrast_fitter.pupils[curr_i].params = params;
                                std::cout << theta-start_params.theta << " " << psi-start_params.psi << " " << radius-start_params.radius;
                                contrast_fitter.refine_single_with_contrast(curr_i);
                                auto val = contrast_fitter.single_contrast_metric(curr_i);
                                std::cout << " -> " << val;
                                if (best_contrast >= val)
                                {
                                    best_params = contrast_fitter.pupils[curr_i].params;
                                    best_contrast = val;
                                    std::cout << " *";
                                }
                                std::cout << std::endl;
                            }
                        }
                    }

                    contrast_fitter.pupils[curr_i].params = best_params;
                    contrast_fitter.pupils[curr_i].circle = contrast_fitter.circleFromParams(best_params);

                    recalc_ellipdist = true;
                    display_est_pupil_ellipses_contrast = true;
                }

                break;

            case '[':
                contrast_fitter.region_band_width--;
                std::cout << "Band width = " << contrast_fitter.region_band_width << std::endl;
                break;

            case ']':
                contrast_fitter.region_band_width++;
                std::cout << "Band width = " << contrast_fitter.region_band_width << std::endl;
                break;

            case ',':
                contrast_fitter.region_scale/=1.125;
                std::cout << "Region scale = " << contrast_fitter.region_scale << std::endl;
                break;

            case '.':
                contrast_fitter.region_scale*=1.125;
                std::cout << "Region scale = " << contrast_fitter.region_scale << std::endl;
                break;

            case 'n':
                null_fitter.unproject_observations(1, 20, ransac);

                simple_fitter.unproject_observations(1, 20, ransac);
                simple_fitter.initialise_model();

                contrast_fitter.unproject_observations(1, 20, ransac);
                contrast_fitter.initialise_model();

                edge_fitter.unproject_observations(1, 20, ransac);
                edge_fitter.initialise_model();

                std::cout << "Model reset" << std::endl;

                break;

            case 'h':
                if (key.shift) {
                    contrast_fitter.pupils[curr_i].params.psi -= 0.01;
                    contrast_fitter.pupils[curr_i].circle = contrast_fitter.circleFromParams(contrast_fitter.pupils[curr_i].params);
                    std::cout << contrast_fitter.pupils[curr_i].params.theta << "," << contrast_fitter.pupils[curr_i].params.psi << "," <<  contrast_fitter.pupils[curr_i].params.radius << std::endl;
                    contrast_fitter.print_single_contrast_metric(curr_i);
                } else {
                    curr_i = curr_i == 0 ? obs_eye_images.size() - 1 : curr_i - 1;
                }
                break;
            case 'j':
                if (key.shift) {
                    contrast_fitter.pupils[curr_i].params.theta -= 0.01;
                    contrast_fitter.pupils[curr_i].circle = contrast_fitter.circleFromParams(contrast_fitter.pupils[curr_i].params);
                    std::cout << contrast_fitter.pupils[curr_i].params.theta << "," << contrast_fitter.pupils[curr_i].params.psi << "," <<  contrast_fitter.pupils[curr_i].params.radius << std::endl;
                    contrast_fitter.print_single_contrast_metric(curr_i);
                }
                break;
            case 'k':
                if (key.shift) {
                    contrast_fitter.pupils[curr_i].params.theta += 0.01;
                    contrast_fitter.pupils[curr_i].circle = contrast_fitter.circleFromParams(contrast_fitter.pupils[curr_i].params);
                    std::cout << contrast_fitter.pupils[curr_i].params.theta << "," << contrast_fitter.pupils[curr_i].params.psi << "," <<  contrast_fitter.pupils[curr_i].params.radius << std::endl;
                    contrast_fitter.print_single_contrast_metric(curr_i);
                }
                break;
            case 'l':
                if (key.shift) {
                    contrast_fitter.pupils[curr_i].params.psi += 0.01;
                    contrast_fitter.pupils[curr_i].circle = contrast_fitter.circleFromParams(contrast_fitter.pupils[curr_i].params);
                    std::cout << contrast_fitter.pupils[curr_i].params.theta << "," << contrast_fitter.pupils[curr_i].params.psi << "," <<  contrast_fitter.pupils[curr_i].params.radius << std::endl;
                    contrast_fitter.print_single_contrast_metric(curr_i);
                } else {
                    curr_i = (curr_i + 1) % obs_eye_images.size();
                }
                break;
            case ';':
                if (key.shift) {
                    contrast_fitter.pupils[curr_i].params.radius /= 1.1;
                    contrast_fitter.pupils[curr_i].circle = contrast_fitter.circleFromParams(contrast_fitter.pupils[curr_i].params);
                    std::cout << contrast_fitter.pupils[curr_i].params.theta << "," << contrast_fitter.pupils[curr_i].params.psi << "," <<  contrast_fitter.pupils[curr_i].params.radius << std::endl;
                    contrast_fitter.print_single_contrast_metric(curr_i);
                }
                break;
            case '\'':
                if (key.shift) {
                    contrast_fitter.pupils[curr_i].params.radius *= 1.1;
                    contrast_fitter.pupils[curr_i].circle = contrast_fitter.circleFromParams(contrast_fitter.pupils[curr_i].params);
                    std::cout << contrast_fitter.pupils[curr_i].params.theta << "," << contrast_fitter.pupils[curr_i].params.psi << "," <<  contrast_fitter.pupils[curr_i].params.radius << std::endl;
                    contrast_fitter.print_single_contrast_metric(curr_i);
                }
                break;

            case 'e':
                if (animate)
                    edge_fitter.refine_with_inliers(animate_func);
                else
                    edge_fitter.refine_with_inliers();
                recalc_ellipdist = true;
                display_est_pupil_ellipses_lm = true;

                break;

            case 'r':
                if (key.shift) {
                    if (!recording_writer.isOpened()) {
                        int fourcc = CV_FOURCC('M','J','P','G');
                        std::cout << "Start recording" << std::endl;
                        std::string out_file = (imagedir / "video.avi").string();
                        int done = recording_writer.open(
                            out_file,
                            fourcc,
                            25,
                            curr_disp.size());
                        if (!done) {
                            std::cerr << "Start failed!" << std::endl;
                        }
                    } else {
                        std::cout << "Stop recording" << std::endl;
                        recording_writer.release();
                    }
                } else {
                    ransac = !ransac;
                    if (ransac) {
                        std::cout << "RANSAC enabled" << std::endl;
                    } else {
                        std::cout << "RANSAC disabled" << std::endl;
                    }
                }
                break;
            case 'p':
                std::cout << curr_i << std::endl;
                std::cout << contrast_fitter.pupils[curr_i].params.theta << "," << contrast_fitter.pupils[curr_i].params.psi << "," <<  contrast_fitter.pupils[curr_i].params.radius << std::endl;
                contrast_fitter.print_single_contrast_metric(curr_i);
                if (curr_i < N) {
                    /*std::cout << EllipseGoodnessFunction<double>()(
                        contrast_fitter.eye,
                        contrast_fitter.pupils[curr_i].params.theta, contrast_fitter.pupils[curr_i].params.psi, contrast_fitter.pupils[curr_i].params.radius,
                        contrast_fitter.focal_length,
                        contrast_fitter.pupils[curr_i].observation.image) << std::endl;

                    std::cout << exp(-sq(contrast_fitter.pupils[curr_i].params.radius - 2.5)/sq(1.0)) << std::endl;

                    if (curr_i > 0) {
                        std::cout << angleDiffGoodness(
                            contrast_fitter.pupils[curr_i-1].params.theta, contrast_fitter.pupils[curr_i-1].params.psi,
                            contrast_fitter.pupils[curr_i].params.theta, contrast_fitter.pupils[curr_i].params.psi,
                            1) << std::endl;

                        std::cout << exp(-sq(contrast_fitter.pupils[curr_i].params.radius-contrast_fitter.pupils[curr_i-1].params.radius)/sq(1.0)) << std::endl;
                    }*/
                }
                break;

            default:
                std::cout << "Unknown key code: " << (int)key.code << " (ascii " << std::string(1, key.code) << ")" << std::endl;
                break;
            }
        }
    }

    //catch (std::exception& e) {
    //    std::cerr << e.what() << std::endl;
    //    while (cv::waitKey(10) == -1) {}
    //    return 1;
    //}

    return 0;
}


/* -*- c++ -*- */
/*
 * Copyright 2015 Free Software Foundation, Inc.
 *
 * This file is part of GNU Radio
 *
 * GNU Radio is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3, or (at your option)
 * any later version.
 *
 * GNU Radio is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with GNU Radio; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#include <gnuradio/fec/ldpc_decoder.h>
#include <math.h>
#include <boost/assign/list_of.hpp>
#include <volk/volk.h>
#include <sstream>
#include <stdio.h>
#include <vector>
#include <gnuradio/fec/decoder.h>
#include <algorithm>            // for std::reverse
#include <string.h>             // for memcpy
#include <gnuradio/fec/maxstar.h>

namespace gr {
 namespace fec {

   generic_decoder::sptr
   ldpc_decoder::make(std::string alist_file, float sigma,
                      int max_iterations)
   {
     return generic_decoder::sptr
       (new ldpc_decoder(alist_file, sigma, max_iterations));
   }

   ldpc_decoder::ldpc_decoder(std::string alist_file, float sigma,
                              int max_iterations)
     : generic_decoder("ldpc_decoder")
   {
     if(!boost::filesystem::exists( alist_file ))
       throw std::runtime_error("Bad AList file name!");

     d_list.read(alist_file.c_str());
     d_code.set_alist(d_list);
     d_spa.set_alist_sigma(d_list, sigma);

     d_rate = static_cast<double>(d_code.dimension())/static_cast<double>(d_code.get_N());
     set_frame_size(d_code.dimension());

     d_spa.set_K(d_output_size);
     d_spa.set_max_iterations(max_iterations);
   }

   int
   ldpc_decoder::get_output_size()
   {
     return d_output_size;
   }

   int
   ldpc_decoder::get_input_size()
   {
     return d_input_size;
   }

   double
   ldpc_decoder::rate()
   {
     return d_rate;
   }

   bool
   ldpc_decoder::set_frame_size(unsigned int frame_size)
   {
     if(frame_size % d_code.dimension() != 0) {
       GR_LOG_ERROR(d_logger, boost::format("Frame size (%1% bits) must be a "
                                            "multiple of the information word "
                                            "size of the LDPC matrix, %2%") \
                    % frame_size % (d_code.dimension()));
       throw std::runtime_error("ldpc_decoder: cannot use frame size.");
     }

     d_output_size = frame_size;
     d_input_size = static_cast<int>(round(frame_size / d_rate));

     return true;
   }

   void
   ldpc_decoder::generic_work(void *inBuffer, void *outBuffer)
   {
     const float *in = (const float *) inBuffer;
     unsigned char *out = (unsigned char *) outBuffer;

     int j = 0;
     std::vector<float> rx(d_code.get_N());
     for(int i = 0; i < d_input_size; i+=d_code.get_N()) {
       for(int k = 0; k < d_code.get_N(); k++) {
         rx[k] = in[i+k] * (-1);
       }

       int n_iterations = 0;
       std::vector<char> estimate( d_spa.decode(rx, &n_iterations) );
       std::vector<char> data( d_code.get_systematic_bits(estimate) );
       memcpy(&out[j], &data[0], d_code.dimension());
       d_iterations = n_iterations;

       j += d_code.dimension();
     }
   }

   int
   ldpc_decoder::get_input_item_size()
   {
     return sizeof(INPUT_DATATYPE);
   }

   int
   ldpc_decoder::get_output_item_size()
   {
     return sizeof(OUTPUT_DATATYPE);
   }

   int
   ldpc_decoder::get_history()
   {
     return 0;
   }

   float
   ldpc_decoder::get_shift()
   {
     return 0.0;
   }

   const char*
   ldpc_decoder::get_conversion()
   {
     return "none";
   }

   ldpc_decoder::~ldpc_decoder()
   {
   }

 } // namespace gr
} // namespace fec


/*******************************************************************************
*
*  Filename    : KeysCompMgr.cc
*  Description : Definition of how objects are stored and data is loaded
*  Author      : Yi-Mu "Enoch" Chen [ ensc@hep1.phys.ntu.edu.tw ]
*
*******************************************************************************/
#include "TstarAnalysis/MassRecoCompare/interface/KeysCompMgr.hpp"

#include "TstarAnalysis/Common/interface/GetEventWeight.hpp"
#include "TstarAnalysis/MassRecoCompare/interface/Common.hpp"

#include "RooAddPdf.h"
#include "RooBifurGauss.h"
#include "RooConstVar.h"
#include "RooDataSet.h"
#include "RooFFTConvPdf.h"
#include "RooGaussian.h"
#include "RooKeysPdf.h"
#include "RooRealVar.h"
#include "TFile.h"

#include <boost/format.hpp>
#include <iostream>

using namespace std;
using namespace mgr;

/*******************************************************************************
*   Declaring global RooRealVars
*******************************************************************************/
void
KeysCompMgr::InitStaticVars()
{
  StaticNewVar( "x", "m_{t+jet}", "GeV/c^{2}", 400, 2500 );
  StaticNewVar( "w", "w",         "",        -1000, 1000 );
  StaticNewVar( "n", "n",         "",           -1,    7 );
}

/******************************************************************************/

RooRealVar&
KeysCompMgr::x(){ return *StaticVar( "x" ); }

/******************************************************************************/

RooRealVar&
KeysCompMgr::w(){ return *StaticVar( "w" ); }

/******************************************************************************/

RooRealVar&
KeysCompMgr::n(){ return *StaticVar( "n" );}

/*******************************************************************************
*   Constructor and data set definition
*******************************************************************************/
KeysCompMgr::KeysCompMgr( const std::string& name, const std::string& latexname ) :
  Named( name ),
  RooFitMgr( name )
{
  SetLatexName( latexname );
  RooDataSet* set = new RooDataSet(
    "",
    "",
    RooArgSet( x(), n(), w() ),
    RooFit::WeightVar( w() )
    );
  AddDataSet( set );
}

/******************************************************************************/

KeysCompMgr::~KeysCompMgr()
{
  for( const auto& name : VarNameList() ){
    RooRealVar* var = Var( name );
    cout << boost::format( "%-50s %8.4lf %8.4f %8.4f" )
      % var->GetName()
      % var->getVal()
      % var->getErrorHi()
      % var->getErrorLo()
         << endl;
  }

  for( const auto& name : FuncNameList() ){
    RooAbsReal* func = Func( name );
    cout << boost::format( "%50s %8.1lf" )
    % func->GetName()
    %func->getVal()
    << endl;
  }
}

/*******************************************************************************
*   Filling by event looping
*******************************************************************************/
void
KeysCompMgr::FillDataSet( const std::string& filename )
{
  fwlite::Event evt( TFile::Open( filename.c_str() ) );
  fwlite::Handle<RecoResult> resulthandle;

  for( evt.toBegin(); !evt.atEnd(); ++evt ){
    resulthandle.getByLabel( evt, Name().c_str(), LabelName( Name() ).c_str(), ProcessName().c_str() );

    const double evtweight = GetEventWeight( evt );
    const double tstarmass = resulthandle.ref().TstarMass();

    if( tstarmass > x().getMax() ){ continue; }
    if( tstarmass < x().getMin() ){ continue; }
    if( evtweight > w().getMax() ){ continue; }
    if( evtweight < w().getMin() ){ continue; }
    x() = tstarmass;
    n() = (double)NumCorrectAssign( resulthandle.ref() );
    DataSet( "" )->add( RooArgSet( x(), n() ), evtweight );
  }

  cout << "Done!" << endl;
}

/******************************************************************************/

void
KeysCompMgr::MakeKeysPdf()
{
  RooDataSet& dataset = *( dynamic_cast<RooDataSet*>( DataSet( "" ) ) );

  double na = dataset.sigma( x() );
  double n2 = dataset.sigma( x(), "n>=3"  );
  double n3 = dataset.sigma( x(), "n>=4"  );
  double n4 = dataset.sigma( x(), "n>=5"  );

  static const vector<double> rholist = {
    1.364,
    1.166,
    1.660
  };
  cout << "Computed Value: "
       << sqrt( na/n2 ) << " "
       << sqrt( na/n3 ) << " "
       << sqrt( na/n4 ) << " "
       << na << endl;

  RooKeysPdf* pdf;

  boost::format nameformat( "rho%.2lf" );
  boost::format titleformat( "#rho = %.2lf" );

  for( const auto rho : rholist ){
    const string name  = str( nameformat % rho );
    const string title = str( titleformat % rho );
    pdf = new RooKeysPdf(
      name.c_str(),
      title.c_str(),
      x(),
      dataset,
      RooKeysPdf::NoMirror,
      rho
      );
    AddPdf( pdf );
  }
}

/******************************************************************************/

void
KeysCompMgr::MakeGauss()
{
  RooRealVar* m  = NewVar( "m", reconamer.GetInput<int>( "mass" ), 400, 3000 );
  RooRealVar* s  = NewVar( "s", 200, 0, 1000 );
  RooAbsPdf* pdf = new RooGaussian( "gauss", "", x(), *m, *s );
  AddPdf( pdf );
  pdf->fitTo(
    *DataSet( "" ),
    RooFit::Minimizer( "Minuit", "Migrad" ),
    RooFit::Minos( kTRUE ),
    RooFit::SumW2Error( kTRUE ),
    RooFit::Verbose( kFALSE ),
    RooFit::PrintLevel( -1 ),
    RooFit::PrintEvalErrors( -1 ),
    RooFit::Warnings( kFALSE )
    );

  RooRealVar* m1         = NewVar( "m1", reconamer.GetInput<int>( "mass" ), 400, 3000 );
  RooRealVar* m2         = NewVar( "m2", reconamer.GetInput<int>( "mass" )+200, 400, 3000 );
  RooRealVar* s1         = NewVar( "s1", 200, 0, 1000 );
  RooRealVar* s2         = NewVar( "s2", 200, 0, 1000 );
  RooRealVar* frac       = NewVar( "frac", 0.5, 0, 1 );
  RooAbsPdf* gauss1      = new RooGaussian( "guass2_1", "", x(), *m1, *s1 );
  RooAbsPdf* gauss2      = new RooGaussian( "guass2_2", "", x(), *m2, *s2 );
  RooAbsPdf* doublegauss = new RooAddPdf( "doublegauss", "", *gauss1, *gauss2, *frac );
  AddPdf( doublegauss );
  doublegauss->fitTo(
    *DataSet( "" ),
    RooFit::Minimizer( "Minuit", "Migrad" ),
    RooFit::Minos( kTRUE ),
    RooFit::SumW2Error( kTRUE ),
    RooFit::Verbose( kFALSE ),
    RooFit::PrintLevel( -1 ),
    RooFit::PrintEvalErrors( -1 ),
    RooFit::Warnings( kFALSE )
    );

}

void
KeysCompMgr::MakeConv()
{
  // Single bifurcation
  RooRealVar* bimean = NewVar( "bimean", 1200, 700, 2000 );
  RooRealVar* bil    = NewVar( "bil", 200, 0, 1000 );
  RooRealVar* bir    = NewVar( "bir", 300, 0, 1000 );
  RooAbsPdf* bipdf   = new RooBifurGauss( "bigauss", "", x(), *bimean, *bil, *bir );

  RooConstVar* gmean = new RooConstVar( "convgmean", "", 0 );
  RooRealVar* gsig   = NewVar( "convgsig", 100, 0, 1000 );
  RooAbsPdf* gpdf    = new RooGaussian( "convgauss", "", x(), *gmean, *gsig );

  RooAbsPdf* convpdf = new RooFFTConvPdf( "conv", "conv", x(), *bipdf, *gpdf );

  convpdf->fitTo(// Forcing fit, and setting to constant
    *DataSet( "" ),
    RooFit::Minimizer( "Minuit", "Migrad" ),
    RooFit::Minos( kTRUE ),
    RooFit::SumW2Error( kTRUE ),
    RooFit::Verbose( kFALSE ),
    RooFit::PrintLevel( -1 ),
    RooFit::PrintEvalErrors( -1 ),
    RooFit::Warnings( kFALSE )
    );
  AddPdf( bipdf );
  AddPdf( gpdf );
  AddPdf( convpdf );
  AddFunc( gmean );

  // Double bifurcation
  RooRealVar* bimean1 = NewVar( "bimean1", 1000, 400, 2000 );
  RooRealVar* bil1    = NewVar( "bil1", 200, 0, 1000 );
  RooRealVar* bir1    = NewVar( "bir1", 300, 0, 1000 );
  RooAbsPdf* bipdf1   = new RooBifurGauss( "bigauss1", "", x(), *bimean1, *bil1, *bir1 );

  RooConstVar* gmean1 = new RooConstVar( "convgmean1", "", 0 );
  RooRealVar* gsig1   = NewVar( "convgsig1", 100, 0, 1000 );
  RooAbsPdf* gpdf1    = new RooGaussian( "convgauss1", "", x(), *gmean1, *gsig1 );

  RooAbsPdf* conv1pdf = new RooFFTConvPdf( "conv1", "", x(), *bipdf1, *gpdf1 );

  RooRealVar* bimean2 = NewVar( "bimean2", 1600, 400, 2000 );
  RooRealVar* bil2    = NewVar( "bil2", 100, 0, 1000 );
  RooRealVar* bir2    = NewVar( "bir2", 200, 0, 1000 );
  RooAbsPdf* bipdf2   = new RooBifurGauss( "bigauss2", "", x(), *bimean2, *bil2, *bir2 );

  RooConstVar* gmean2 = new RooConstVar( "convgmean2", "", 0 );
  RooRealVar* gsig2   = NewVar( "convgsig2", 100, 0, 1000 );
  RooAbsPdf* gpdf2    = new RooGaussian( "convgauss2", "", x(), *gmean2, *gsig2 );

  RooAbsPdf* conv2pdf = new RooFFTConvPdf( "conv2", "", x(), *bipdf2, *gpdf2 );

  RooRealVar* bico = NewVar( "bicoeff2", 0.6, 0, 1 );

  RooAbsPdf* bisum = new RooAddPdf( "convsum2", "", *conv1pdf, *conv2pdf, *bico );

  bisum->fitTo(// Forcing fit, and setting to constant
    *DataSet( "" ),
    RooFit::Minimizer( "Minuit", "Migrad" ),
    RooFit::Minos( kTRUE ),
    RooFit::SumW2Error( kTRUE ),
    RooFit::Verbose( kFALSE ),
    RooFit::PrintLevel( -1 ),
    RooFit::PrintEvalErrors( -1 ),
    RooFit::Warnings( kFALSE )
    );

  // make sample claim ownership of pdfs

  AddPdf( bipdf1 );
  AddPdf( gpdf1 );
  AddPdf( conv1pdf );
  AddPdf( bipdf2 );
  AddPdf( gpdf2 );
  AddPdf( conv2pdf );
  AddPdf( bisum );

  AddFunc( gmean1 );
  AddFunc( gmean2 );
}


#include "bwd_lagrangian.h"

#include <cmath>
#include <Eigen/Dense>

#include "../amr/mesh.h"
#include "../amr/numerics.h"
#include "../amr/refiner.h"

using namespace Eigen;

namespace vlv { namespace tools {
template<typename T>
T modf(T arg, T* iptr) = delete;

template<>
float modf(float arg, float* iptr) { return modff(arg, iptr); }

template<>
double modf(double arg, double* iptr) { return modf(arg, iptr); }

} }

template<typename T, int D, int V>
void vlv::AmrMomentumLagrangianSolver<T,D,V>::solve_mesh(
        toolbox::AdaptiveMesh<T, 3>& mesh0,
        toolbox::AdaptiveMesh<T, 3>& mesh1,
        vec& Einc,
        vec& Binc,
        tools::Params<T>& params)
{

  toolbox::Adapter<T,3> adapter;
  adapter.cells_to_refine.clear();
  adapter.cells_to_unrefine.clear();

  // empty the target mesh
  // TODO: is this efficient; should we recycle instead?
  mesh1.clear();

  std::array<uint64_t, 3> index;
  // std::array<T, 3> grad;

  Vec3E B(Binc.data());
  Vec3E E(Einc.data());

  // level zero fill
  auto val     = T(0);
  T max_val = mesh0.max_value();

  T refine_indicator, unrefine_indicator;
  auto len = mesh0.get_size(0);

  // XXX remove higher-dimension updates (compresses inner loops away)
  // TODO does not work if the fill is not correctly done
  for(int i = 2; i>V-1; i--) len[i] = 1;

  // normalize
  //T norm0 = integrate_moment( mesh0,
  //          [](std::array<T,3>& uvel) -> T { return T(1);}
  //          );

  // cfl bound guards
  auto cids = mesh0.get_cells(true);

  // if mesh is empty, bail out
  if (cids.empty()) { return; }

  auto min_ind = mesh0.get_indices( cids.front() );
  auto max_ind = mesh0.get_indices( cids.back()  );

  int cfl_halo = 10; // how many CFL steps are allowed in backward substitution

  for(uint64_t t=0; t<len[2]; t++) {
    index[2] = t;
    for(uint64_t s=0; s<len[1]; s++) {
      index[1] = s;
      for(uint64_t r=0; r<len[0]; r++) {
        index[0] = r;

        // cfl bound guards
        if( r < (min_ind[0] - cfl_halo) ) continue;
        if( r > (max_ind[0] + cfl_halo) ) continue;

        uint64_t cid = mesh1.get_cell_from_indices(index, 0);
        val = backward_advect(index, 0, mesh0, E, B, params);

        // refinement
        // TODO specialize to value & gradient instead of just value
        refine_indicator   = val/max_val;
        unrefine_indicator = val/max_val;
        adapter.check_cell(mesh1, cid, refine_indicator, unrefine_indicator);

        mesh1.set(cid, val);
      }
    }
  }

  // create new leafs
  // TODO fixme to use max ref. lvl of mesh
  // for(size_t sweep=1; sweep<=mesh1.maximum_refinement_level; sweep++){
  for(int sweep=1; sweep<=mesh1.top_refinement_level; sweep++){

    adapter.refine(mesh1);

    // fmt::print(" sol: cells created {}\n", adapter.cells_to_refine.size());
    // fmt::print(" sol: cells removed {}\n", adapter.cells_removed.size());

    adapter.cells_to_refine.clear();
    adapter.cells_to_unrefine.clear();

    // next we refine
    for(auto&& cid : adapter.cells_created) {
      int rfl = mesh1.get_refinement_level(cid);
      auto index2 = mesh1.get_indices(cid);

      // fmt::print("creating {} at {}\n", cid, rfl);

      val = backward_advect(index2, rfl, mesh0, E, B, params);
      mesh1.set(cid, val);

      refine_indicator   = val/max_val;
      unrefine_indicator = val/max_val;
      adapter.check_cell(mesh1, cid, refine_indicator, unrefine_indicator);
    }

    // now unrefine
    adapter.unrefine(mesh1);

  }

  // normalize back to original weight
  //T norm1 = integrate_moment( mesh1,
  //          [](std::array<T,3>& uvel) -> T { return T(1);}
  //          );
  //mesh1 *= norm0/norm1;

  }

// backward advected Lorentz force
template<typename T, int D, int V>
T vlv::AmrMomentumLagrangianSolver<T,D,V>::backward_advect(
    std::array<uint64_t, 3>& index,
    int rfl,
    const toolbox::AdaptiveMesh<T, 3>& mesh0,
    //toolbox::AdaptiveMesh<T, 3>& mesh1,
    Vec3E& E,
    Vec3E& B,
    tools::Params<T>& params)
{
  T val; // return value
  vec u    = mesh0.get_center(index, rfl);  // velocity
  vec du   = mesh0.get_length(rfl);         // box size, i.e., \Delta u
  auto len = mesh0.get_size(rfl);

  // get shift of the characteristic solution from Lorentz force
  Vec3E uvel( u.data() );
  Vec3E F = lorentz_force(uvel, E, B, params.qm, params.cfl);

  // add other forces; default to zero
  Vec3E Fi = other_forces(uvel, params);
  F += Fi;

  // advection in units of cells
  // NOTE: We normalize with CFL because velocities are in units
  // of c and we need them in units of grid speed.
  // XXX: CHECK
  std::array<T,3> shift, cell_shift;
  for(int i=0; i<V; i++) shift[i] = F(i) / du[i];

  // advected tiles
  std::array<int,3> index_shift;
  for(int i=0; i<V; i++) index_shift[i] = static_cast<int>( trunc(shift[i]) );

  // advection inside cell
  T tmp;
  for(int i=0; i<V; i++) cell_shift[i] = tools::modf<T>(shift[i], &tmp);

  // new grid indices
  std::array<uint64_t, 3> index_new;
  for(int i=0;   i<V;   i++) index_new[i] = index[i] + index_shift[i];
  for(int i = 2; i>V-1; i--) index_new[i] = index[i];

  // set boundary conditions (zero outside the grid limits)
  for(int i=0; i<V; i++){
    if( (index_new[i] <2) || (index_new[i] >= len[i]-2) ) return T(0);
  }

  // interpolation branch
  val = toolbox::interp_cubic<T,V>(mesh0, index_new, cell_shift, rfl);
  //val = toolbox::interp_linear<T,V>(mesh0, index_new, cell_shift, rfl);

  return val;
}

/// Relativistic Lorentz force / Electrostatic version
template<typename T, int D, int V>
inline typename vlv::AmrMomentumLagrangianSolver<T,D,V>::Vec3E
vlv::AmrMomentumLagrangianSolver<T,D,V>::lorentz_force(
  Vec3E& /*uvel*/,
  Vec3E& E,
  Vec3E& /*B*/,
  T qm,
  T cfl)
{
  // electrostatic push
  //
  // Boris scheme for b=0 translates to
  // u = (cfl*u_0 + e + e)/cfl = u_0 + E/cfl
  //
  // with halving taken into account in definition of Ex
  return -qm*E/cfl;
}

/// default zero force for to be overloaded by more complicated solvers
template<typename T, int D, int V>
inline typename vlv::AmrMomentumLagrangianSolver<T,D,V>::Vec3E
vlv::AmrMomentumLagrangianSolver<T,D,V>::other_forces(
    Vec3E& /*uvel*/,
    tools::Params<T>& /*params*/
    )
{
  Vec3E ret = Vec3E::Zero();
  return ret;
}

//--------------------------------------------------
// explicit template instantiation
template class vlv::AmrMomentumLagrangianSolver<Realf, 1, 1>; //1D1V


/*=========================================================================
  The Software is copyright (c) Commonwealth Scientific and Industrial Research Organisation (CSIRO)
  ABN 41 687 119 230.
  All rights reserved.

  Licensed under the CSIRO BSD 3-Clause License
  You may not use this file except in compliance with the License.
  You may obtain a copy of the License in the file LICENSE.md or at

  https://stash.csiro.au/projects/SMILI/repos/smili/browse/license.txt

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
=========================================================================*/
#include "milxQtDiffusionTensorModel.h"

#include <time.h>

#include <vtkMath.h>
#include <vtkLine.h>
#include <vtkSphereSource.h>

#include <boost/math/special_functions/spherical_harmonic.hpp>

milxQtDiffusionTensorModel::milxQtDiffusionTensorModel(QWidget *theParent) : milxQtModel(theParent)
{
    milxQtWindow::prefix = "DTI: ";

    createActions();

    createConnections();
}

milxQtDiffusionTensorModel::~milxQtDiffusionTensorModel()
{
    //dtor
}

void milxQtDiffusionTensorModel::colourByDirection()
{
    vtkSmartPointer<vtkPolyData> currentMesh = model.Result();

    typedef double projectionType;

    ///Determine colours based on axes directions
    vtkSmartPointer<vtkUnsignedCharArray> scalars = vtkSmartPointer<vtkUnsignedCharArray>::New();
        scalars->SetNumberOfComponents(3);
        scalars->SetNumberOfTuples(currentMesh->GetNumberOfPoints());
        scalars->SetName("Fibre Colours");
        scalars->FillComponent(0, 0);
        scalars->FillComponent(1, 0);
        scalars->FillComponent(2, 0);
    vtkSmartPointer<vtkFloatArray> projections = vtkSmartPointer<vtkFloatArray>::New();
        projections->SetNumberOfComponents(3);
        projections->SetNumberOfTuples(currentMesh->GetNumberOfPoints());
        projections->SetName("Fibre Projections");
        projections->FillComponent(0, 0.0);
        projections->FillComponent(1, 0.0);
        projections->FillComponent(2, 0.0);

    emit working(-1);
    if(currentMesh->GetNumberOfLines() == 0)
    {
        printInfo("No lines in model. Computing colours for mesh.");

        ///Use the dot product in each axis
        for(size_t j = 0; j < 3; j ++)
        {
            projectionType axis[3] = {0.0, 0.0, 0.0};
            axis[j] = 1.0;
            cout << "Computing toward axis " << j << endl;

            ///Colour based on each of the gradient values in that direction
            for(vtkIdType k = 0; k < currentMesh->GetNumberOfPoints(); k ++)
            {
                coordinate currentProjection(projections->GetTuple3(k)), position(currentMesh->GetPoint(k));

                projectionType projection = vtkMath::Dot(axis, position.data_block()); //project to axis being done,
                currentProjection[j] = projection; //projection in each direction

                projections->SetTuple3(k, currentProjection[0], currentProjection[1], currentProjection[2]);
            }
        }

        ///Colour based on each of the gradient values in that direction
        for(vtkIdType k = 0; k < currentMesh->GetNumberOfPoints(); k ++)
        {
            coordinate currentProjection(projections->GetTuple3(k));
            coordinate currentProjSquared = element_product(currentProjection, currentProjection);
            projectionType maxProjection = currentProjSquared.max_value();
            currentProjSquared /= maxProjection;

            unsigned char colourOfPoint[3] = {0, 0, 0};
            colourOfPoint[0] = static_cast<unsigned char>( currentProjSquared[0]*255.0 );
            colourOfPoint[1] = static_cast<unsigned char>( currentProjSquared[1]*255.0 );
            colourOfPoint[2] = static_cast<unsigned char>( currentProjSquared[2]*255.0 );

            scalars->SetTupleValue(k, colourOfPoint);
        }

        currentMesh->GetPointData()->SetVectors(projections);
        currentMesh->GetPointData()->SetScalars(scalars);
    }
    else
    {
        printInfo("Re-colouring lines by axes directions");
        //Ensure lines done properly so can colour appropriately
        currentMesh->GetLines()->InitTraversal();
        vtkSmartPointer<vtkPoints> points = vtkSmartPointer<vtkPoints>::New();
        vtkSmartPointer<vtkIdList> idList = vtkSmartPointer<vtkIdList>::New();
        std::vector<vtkIdType> trackLengths;
        for(vtkIdType j = 0; j < currentMesh->GetNumberOfLines(); j ++)
        {
            currentMesh->GetLines()->GetNextCell(idList);
            trackLengths.push_back(idList->GetNumberOfIds());
            for(vtkIdType pointId = 0; pointId < idList->GetNumberOfIds(); pointId ++)
            {
    //            std::cout << idList->GetId(pointId) << " ";

                double position[3];
                currentMesh->GetPoint(idList->GetId(pointId), position);
                points->InsertNextPoint(position);
            }
        }

        //Re-stitch lines together
        vtkIdType step = 0;
        vtkSmartPointer<vtkCellArray> lines = vtkSmartPointer<vtkCellArray>::New();
        for(size_t j = 0; j < trackLengths.size(); j ++)
        {
            for(vtkIdType k = step; k < trackLengths[j]-1; k ++)
            {
                vtkSmartPointer<vtkLine> line = vtkSmartPointer<vtkLine>::New();
                    line->GetPointIds()->SetId(0, k);
                    line->GetPointIds()->SetId(1, k+1);
                lines->InsertNextCell(line);
            }
            qDebug() << trackLengths[j] << endl;
            step += trackLengths[j];
        }

        vtkSmartPointer<vtkPolyData> linesPolyData = vtkSmartPointer<vtkPolyData>::New();
            //Add the points to the dataset
            linesPolyData->SetPoints(points);
            //Add the lines to the dataset
            linesPolyData->SetLines(lines);
            linesPolyData->Modified();
        model.SetInput(linesPolyData);
    }

    emit done(-1);
    generateModel();
}

void milxQtDiffusionTensorModel::harmonics(QString ordersString)
{
    bool ok = false;
    if(ordersString.isEmpty())
    {
        ordersString = QInputDialog::getText(this, tr("Enter order magnitudes of harmonics separated by spaces"),
                                            tr("Orders: "), QLineEdit::Normal,
                                            "1 0 0 0 1 1 1 1 1", &ok);
    }
    if (!ok || ordersString.isEmpty())
        return;

    QStringList orders = ordersString.split(" ");
    std::vector<double> sHarmonicCoefficients;
    for(int i = 0; i < orders.size(); i++)
        sHarmonicCoefficients.push_back(orders[i].toDouble());

    int n_coefs = sHarmonicCoefficients.size();
    int l_max   = milxQtDiffusionTensorModel::LforN( n_coefs );

    vtkSmartPointer<vtkSphereSource> sphereSource = vtkSmartPointer<vtkSphereSource>::New();
        sphereSource->SetThetaResolution( 64 );
        sphereSource->SetPhiResolution( 64 );
        sphereSource->SetRadius( 1.0 );
        sphereSource->Update();
    vtkSmartPointer<vtkPolyData> glyph = sphereSource->GetOutput();

    ///For every point in sphere mesh
    double m_x = 0, m_y = 0, m_z = 0;
    for(int i = 0; i < glyph->GetNumberOfPoints(); i++)
    {
        double point_sphere[3];
        glyph->GetPoint(i, point_sphere);
        double x_s = point_sphere[0];
        double y_s = point_sphere[1];
        double z_s = point_sphere[2];

        double x_g, y_g, z_g;

        ///Compute spherical harmonic amplitude
        double amplitude = computeAmplitude( sHarmonicCoefficients, x_s, y_s, z_s, l_max );

        if(amplitude < 0)
            amplitude = 0;

        ///use this to displace sphere points accordingly
        x_g = x_s * amplitude + m_x;
        y_g = y_s * amplitude + m_y;
        z_g = z_s * amplitude + m_z;

        glyph->GetPoints()->SetPoint(i, x_g, y_g, z_g);
    }

    model.SetInput(glyph);
    model.GenerateNormals(true);

    generateModel();
    milxQtRenderWindow::reset();
}

double milxQtDiffusionTensorModel::computeAmplitude(std::vector<double> SH, double x, double y, double z, int lmax)
{
    double az = atan2(y, x);
    double el = acos(z);
    double val = 0.0;
    for (int l = 0; l <= lmax; l+=2)
    {
        // val += SH[milxQtDiffusionTensorModel::index(l,0)] * boost::math::legendre_p<double>(l, 0, z);
        val += SH[milxQtDiffusionTensorModel::index(l,0)] * boost::math::spherical_harmonic_r<double>(l, 0, el,az);
    }

    for (int m = 1; m <= lmax; m++)
    {
//        float caz = cos(m*az);
//        float saz = sin(m*az);
        for (int l = 2*((m+1)/2); l <= lmax; l+=2)
        {
        // float buf = boost::math::legendre_p<double>(l, 0, z);
        // val += SH[index(l,m)]*buf*caz;
        // val += SH[index(l,-m)]*buf*saz;
            val += SH[milxQtDiffusionTensorModel::index(l,m)]*boost::math::spherical_harmonic_r(l,m,el,az);
            val += SH[milxQtDiffusionTensorModel::index(l,-m)]*boost::math::spherical_harmonic_i(l,m,el,az);
        }
    }

    return val;
}

void milxQtDiffusionTensorModel::createActions()
{
    colourDirectionAct = new QAction(this);
        colourDirectionAct->setText(QApplication::translate("Model", "Colour By Direction", 0, QApplication::UnicodeUTF8));
        colourDirectionAct->setShortcut(tr("Alt+c"));
    harmonicsAct = new QAction(this);
        harmonicsAct->setText(QApplication::translate("Model", "Show Spherical Harmonic ...", 0, QApplication::UnicodeUTF8));
        harmonicsAct->setShortcut(tr("Alt+s"));
}

void milxQtDiffusionTensorModel::createConnections()
{
    //Operations
    connect(colourDirectionAct, SIGNAL(triggered()), this, SLOT(colourByDirection()));
    connect(harmonicsAct, SIGNAL(triggered()), this, SLOT(harmonics()));

    //milxQtModel::createConnections();
}

void milxQtDiffusionTensorModel::contextMenuEvent(QContextMenuEvent *currentEvent)
{
    contextMenu = milxQtModel::basicContextMenu(); //!< Only exists for the duration of the context selection

    contextMenu->addSeparator()->setText(tr("DiffusionTensor"));
    contextMenu->addAction(colourDirectionAct);
    contextMenu->addAction(harmonicsAct);
    contextMenu->addSeparator()->setText(tr("Extensions"));
    foreach(QAction *currAct, extActionsToAdd)
    {
        contextMenu->addAction(currAct);
    }
    contextMenu->addSeparator();
    ///Dont display extensions
    contextMenu->addAction(milxQtModel::scaleAct);
    contextMenu->addAction(milxQtRenderWindow::axesAct);
    contextMenu->addAction(milxQtRenderWindow::refreshAct);

    contextMenu->exec(currentEvent->globalPos());
}


#include <iostream>
#include <fstream>
#include <string>
#include <iomanip>
#include <Eigen/Eigen>
#include <glog/logging.h>
#include <experimental/filesystem>
#include <sys/stat.h>
#include <sys/types.h>
#include <dirent.h>
#include <unistd.h>
#include <opencv2/opencv.hpp>
#include <opencv2/core/eigen.hpp>

#include "json.h"
#include "reprojection_error_cost.h"

using json = nlohmann::json;
namespace fs = std::experimental::filesystem;
using namespace lass;

#define print_var(x) std::cout << #x << " " << x << std::endl;

struct FrameData {
    FrameData() {
        K.setIdentity();
        qcw_gt.setIdentity();
        qcw.setIdentity();
        pcw_gt.setZero();
        pcw.setZero();
    }
    Eigen::Quaterniond qcw_gt, qcw;
    Eigen::Vector3d pcw_gt, pcw;
    Eigen::Matrix3d K;
    std::vector<Eigen::Vector3d> landmarks;
    std::vector<Eigen::Vector2d> keypoints;
    std::string filename;
    std::vector<int> inlier_indices;
};

inline std::vector<std::string> get_filenames(std::string base_dir) {
    std::vector<std::string> filenames;
    DIR *dir;
    struct dirent *entry;
    if ((dir = opendir(base_dir.c_str())) != NULL) {
        /* print all the files and directories within directory */
        while ((entry = readdir(dir)) != NULL) {
            // ignore files and hidden folders
            if (entry->d_type == DT_REG && entry->d_name[0] != '.') {
                filenames.emplace_back(entry->d_name);
                // printf("Find files %s\n", entry->d_name);
            }
        }
        closedir(dir);
    }
    return filenames;
}

inline FrameData load_single_from_json(const std::string &filename, std::vector<Eigen::Vector3d> &id2center) {
    std::ifstream ifs(filename);
    json j_data;
    ifs >> j_data;
    ifs.close();
    // std::cout << j_data;
    CHECK(j_data["label"].size() == j_data["uv"].size()) << "size inconsistency";
    FrameData data;
    Eigen::Matrix3d R;
    for (int row = 0; row < 3; row++) {
        for (int col = 0; col < 4; col++) {
            if (col >= 3) {
                data.pcw_gt(row) = j_data["pose"][row * 4 + col];
            } else {
                R(row, col) = j_data["pose"][row * 4 + col];
            }
        }
    }
    data.qcw_gt = R;
    for (int row = 0; row < 3; ++row) {
        for (int col = 0; col < 3; ++col) {
            data.K(row, col) = j_data["camera_k_matrix"][row * 3 + col];
        }
    }
    for (int idx = 0; idx < j_data["label"].size(); idx++) {
        int id = j_data["label"][idx];
        data.landmarks.push_back(id2center[id]);
        data.keypoints.emplace_back(j_data["uv"][idx][0], j_data["uv"][idx][1]);
    }
    data.filename = filename;
    // print_var(data.filename);
    // print_var(data.K);
    // print_var(data.pcw_gt);
    // print_var(data.qcw_gt.coeffs());
    // print_var(data.keypoints.size());
    // print_var(data.landmarks.size());
    return data;
}

inline std::vector<FrameData> load_data_from_json(const std::string &base_dir) {
    std::ifstream cfs(base_dir + "/id2centers.json");
    json j_centers;
    cfs >> j_centers;
    cfs.close();
    std::vector<Eigen::Vector3d> centers(j_centers.size());
    for (int idx = 0; idx < j_centers.size(); idx++) {
        json &j_center = j_centers[idx];
        centers[idx] = {j_center[0], j_center[1], j_center[2]};
    }

    auto filenames = get_filenames(base_dir + "/prediction_json");
    std::sort(filenames.begin(), filenames.end());
    std::vector<FrameData> frames;
    for (const auto &fn : filenames) {
        frames.push_back(load_single_from_json(base_dir + "/prediction_json/" + fn, centers));
    }
    return frames;
}

double compute_reprojection_error(const FrameData &frame) {
    double rpe_sum = 0;
    double count = 0;
    for (const auto &idx : frame.inlier_indices) {
        Eigen::Vector3d pt_c = frame.qcw * frame.landmarks[idx] + frame.pcw;
        Eigen::Vector2d pt_2d = pt_c.hnormalized();
        CHECK(std::abs(frame.K(0, 0) - frame.K(1, 1)) < 1e-5);
        pt_2d = {frame.K(0, 0) * pt_2d.x() + frame.K(0, 2), frame.K(0, 0) * pt_2d.y() + frame.K(1, 2)};
        double rpe = (pt_2d - frame.keypoints[idx]).norm();
        rpe_sum += rpe;
        count++;
    }
    return rpe_sum / count;
}

void solve_pose(FrameData &frame) {
    // get initial pose and inlier points via opencv solvePnPRansac
    std::vector<cv::Point2d> image_pts;
    std::vector<cv::Point3d> world_pts;
    for (int i = 0; i < frame.landmarks.size(); ++i) {
        image_pts.emplace_back(frame.keypoints[i].x(), frame.keypoints[i].y());
        world_pts.emplace_back(frame.landmarks[i].x(), frame.landmarks[i].y(), frame.landmarks[i].z());
    }
    cv::Mat K = cv::Mat::eye(cv::Size(3, 3), CV_64F);
    K.at<double>(0, 0) = frame.K(0, 0);
    K.at<double>(0, 2) = frame.K(0, 2);
    K.at<double>(1, 1) = frame.K(1, 1);
    K.at<double>(1, 2) = frame.K(1, 2);
    cv::Mat rotation, translation;
    std::vector<int> inlier_indices;
    cv::solvePnPRansac(world_pts, image_pts,
                       K, cv::Mat(),
                       rotation, translation, false,
                       200, 8.0, 0.99, inlier_indices,
                       cv::SOLVEPNP_P3P);
    cv::Rodrigues(rotation, rotation);
    Eigen::Matrix3d R;
    cv::cv2eigen(rotation, R);
    frame.qcw = R;
    cv::cv2eigen(translation, frame.pcw);
    // print_var(rotation);
    // print_var(translation);

    // check frustum
    for (auto it = inlier_indices.begin(); it != inlier_indices.end();) {
        Eigen::Vector3d pt_c = frame.qcw * frame.landmarks[*it] + frame.pcw;
        bool is_in_frustum = pt_c.z() > 0;
        Eigen::Vector2d pt_2d = pt_c.hnormalized();
        pt_2d = {frame.K(0, 0) * pt_2d.x() + frame.K(0, 2), frame.K(0, 0) * pt_2d.y() + frame.K(1, 2)};
        if (pt_2d.x() < 0 || pt_2d.x() > 480 || pt_2d.y() < 0 || pt_2d.y() > 270) {
            is_in_frustum = false;
        }
        if (!is_in_frustum) {
            std::cerr << "find pt not in frustum!\n";
            print_var(pt_c);
            print_var(pt_2d);
            it = inlier_indices.erase(it);
        } else {
            ++it;
        }
    }

    // prepare optimize data
    if (inlier_indices.size() < 5) {
        std::cerr << "inliers not enough! current inlier count " << inlier_indices.size() << std::endl;
    }

    frame.inlier_indices = inlier_indices;
    std::vector<Eigen::Vector3d> inlier_landmarks;
    std::vector<Eigen::Vector2d> inlier_keypoints;
    // print_var(inlier_indices.size());
    for (const auto &idx : inlier_indices) {
        inlier_keypoints.push_back(frame.keypoints[idx]);
        inlier_landmarks.push_back(frame.landmarks[idx]);
    }

    // print_var(frame.qcw_gt.coeffs());
    // print_var(frame.qcw.coeffs());
    // print_var(frame.pcw_gt);
    // print_var(frame.pcw);
    // print_var(compute_reprojection_error(frame));
    double rpe_before = compute_reprojection_error(frame);

    // optimize via ceres
    ceres::Problem problem;
    std::array<double, 7> camera = {frame.qcw.w(), frame.qcw.x(), frame.qcw.y(), frame.qcw.z(), frame.pcw.x(), frame.pcw.y(), frame.pcw.z()};
    auto loss_function = new ceres::HuberLoss(1.0);
    ceres::LocalParameterization *camera_parameterization =
        new ceres::ProductParameterization(
            new ceres::QuaternionParameterization(),
            new ceres::IdentityParameterization(3));
    problem.AddParameterBlock(camera.data(), 7, camera_parameterization);
    for (int i = 0; i < inlier_landmarks.size(); ++i) {
        problem.AddParameterBlock(inlier_landmarks[i].data(), 3);
        problem.SetParameterBlockConstant(inlier_landmarks[i].data());
        ceres::CostFunction *cost_function = ReprojectionErrorWithQuaternions::Create(
            inlier_keypoints[i], frame.K);
        problem.AddResidualBlock(cost_function, loss_function, camera.data(), inlier_landmarks[i].data());
    }
    ceres::Solver::Options solver_options;
    solver_options.linear_solver_type = ceres::DENSE_SCHUR;
    solver_options.trust_region_strategy_type = ceres::LEVENBERG_MARQUARDT;
    solver_options.use_explicit_schur_complement = true;
    // solver_options.minimizer_progress_to_stdout = true;
    solver_options.minimizer_progress_to_stdout = false;
    ceres::Solver::Summary solver_summary;
    ceres::Solve(solver_options, &problem, &solver_summary);
    // print_var(camera[0]);
    // print_var(camera[1]);
    // print_var(camera[2]);
    // print_var(camera[3]);
    // print_var(camera[4]);
    // print_var(camera[5]);
    // print_var(camera[6]);
    // std::cout << solver_summary.BriefReport();
    // assign valie
    frame.pcw = {camera[4], camera[5], camera[6]};
    frame.qcw = Eigen::Quaterniond(camera[0], camera[1], camera[2], camera[3]);
    double rpe_after = compute_reprojection_error(frame);
    printf("reprojection error : %.4f -> %.4f\n", rpe_before, rpe_after);
}

int main(int argc, char **argv) {
    std::string base_dir = argv[1];
    auto frames = load_data_from_json(base_dir);
    std::vector<double> APEs, AREs;

    for (auto &f : frames) {
        // solve and optimize
        solve_pose(f);
        // evaluate
        if (f.inlier_indices.empty()) {
            APEs.push_back(std::numeric_limits<double>::max());
            AREs.push_back(std::numeric_limits<double>::max());
            continue;
        }
        APEs.push_back((f.pcw - f.pcw_gt).norm() * 1e3);
        Eigen::Quaterniond q_diff = f.qcw * f.qcw_gt.conjugate();
        Eigen::AngleAxisd aa(q_diff);
        AREs.push_back(aa.angle() * 180 / M_PI);
    }
    std::sort(APEs.begin(), APEs.end());
    std::sort(AREs.begin(), AREs.end());
    printf("APE median: %.3f [mm]\nARE median %.3f[DEG]\n", APEs[APEs.size() / 2], AREs[AREs.size() / 2]);
    return 0;
}


/* Siconos is a program dedicated to modeling, simulation and control
 * of non smooth dynamical systems.
 *
 * Copyright 2021 INRIA.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
*/

#include "Lagrangian2d1DR.hpp"
#include <boost/math/quaternion.hpp>
#include "NewtonEulerDS.hpp"
#include "Interaction.hpp"
#include "BlockVector.hpp"
#include <boost/math/quaternion.hpp>

// #define DEBUG_NOCOLOR
// #define DEBUG_STDOUT
// #define DEBUG_MESSAGES
#include "siconos_debug.h"

void Lagrangian2d1DR::initialize(Interaction& inter)
{
  LagrangianR::initialize(inter);
  //proj_with_q  _jachqProj.reset(new SimpleMatrix(_jachq->size(0),_jachq->size(1)));

  if((inter.getSizeOfDS() !=3) and (inter.getSizeOfDS() !=6))
  {
    THROW_EXCEPTION("Lagrangian2d1DR::initialize(Interaction& inter). The size of ds must of size 3");
  }
  unsigned int qSize = 3 * (inter.getSizeOfDS() / 3);
  _jachq.reset(new SimpleMatrix(1, qSize));
}

void Lagrangian2d1DR::computeJachq(const BlockVector& q, BlockVector& z)
{
  DEBUG_BEGIN("Lagrangian2d1DR::computeJachq(Interaction& inter, SP::BlockVector q0 \n");

  double Nx = _Nc->getValue(0);
  double Ny = _Nc->getValue(1);
  double Px = _Pc1->getValue(0);
  double Py = _Pc1->getValue(1);
  double G1x = q.getValue(0);
  double G1y = q.getValue(1);

  _jachq->setValue(0,0,Nx);
  _jachq->setValue(0,1,Ny);
  _jachq->setValue(0,2,(G1y-Py)*Nx - (G1x-Px)*Ny);

  if(q.size() ==6)
  {
    DEBUG_PRINT("take into account second ds\n");
    double G2x = q.getValue(3);
    double G2y = q.getValue(4);

    _jachq->setValue(0,3,-Nx);
    _jachq->setValue(0,4,-Ny);
    _jachq->setValue(0,5,- ((G2y-Py)*Nx - (G2x-Px)*Ny));
  }
  DEBUG_EXPR(_jachq->display(););
  DEBUG_END("Lagrangian2d1DR::computeJachq(Interaction& inter, SP::BlockVector q0) \n");

}

double Lagrangian2d1DR::distance() const
{
  DEBUG_BEGIN("Lagrangian2d1DR::distance(...)\n")
  SiconosVector dpc(*_Pc2 - *_Pc1);
  DEBUG_END("Lagrangian2d1DR::distance(...)\n")
  return dpc.norm2() * (inner_prod(*_Nc, dpc) >= 0 ? -1 : 1);

}

void Lagrangian2d1DR::computeh(const BlockVector& q, BlockVector& z, SiconosVector& y)
{
  DEBUG_BEGIN("Lagrangian2d1DR::computeh(...)\n");
  DEBUG_EXPR(q.display());
  // Contact points and normal are stored as relative to q1 and q2, if
  // no q2 then pc2 and normal are absolute.

  // Update pc1 based on q and relPc1

  double angle= q(2);
  (*_Pc1)(0) = q(0) + cos(angle) * (*_relPc1)(0)+ sin(angle) * (*_relPc1)(1);
  (*_Pc1)(1) = q(1) + sin(angle) * (*_relPc1)(0)- cos(angle) * (*_relPc1)(1);
  if(q.size() == 6)
  {
    // To be checked
    DEBUG_PRINT("take into account second ds\n");
    angle = q(5);
    (*_Pc2)(0) = q(3) + cos(angle) * (*_relPc2)(0)+ sin(angle) * (*_relPc2)(1);
    (*_Pc2)(1) = q(4) + sin(angle) * (*_relPc2)(0)- cos(angle) * (*_relPc2)(1);
    (*_Nc)(0) =  cos(angle) * (*_relNc)(0)+ sin(angle) * (*_relNc)(1);
    (*_Nc)(1) =  sin(angle) * (*_relNc)(0)+ cos(angle) * (*_relNc)(1);
  }
  else
  {
    *_Pc2 = *_relPc2;
    *_Nc = *_relNc;
  }

  // SP::SiconosVector q1 = (q0.getAllVect())[0];
  // ::boost::math::quaternion<double> qq1((*q1)(3), (*q1)(4), (*q1)(5), (*q1)(6));
  // ::boost::math::quaternion<double> qpc1(0,(*_relPc1)(0),(*_relPc1)(1),(*_relPc1)(2));

  // // apply q1 rotation and add
  // qpc1 = qq1 * qpc1 / qq1;
  // (*_Pc1)(0) = qpc1.R_component_2() + (*q1)(0);
  // (*_Pc1)(1) = qpc1.R_component_3() + (*q1)(1);
  // (*_Pc1)(2) = qpc1.R_component_4() + (*q1)(2);

  // if (q0.numberOfBlocks() > 1)
  // {
  //   // Update pc2 based on q0 and relPc2
  //   SP::SiconosVector q2 = (q0.getAllVect())[1];
  //   ::boost::math::quaternion<double> qq2((*q2)(3), (*q2)(4), (*q2)(5), (*q2)(6));
  //   ::boost::math::quaternion<double> qpc2(0,(*_relPc2)(0),(*_relPc2)(1),(*_relPc2)(2));

  //   // apply q2 rotation and add
  //   qpc2 = qq2 * qpc2 / qq2;
  //   (*_Pc2)(0) = qpc2.R_component_2() + (*q2)(0);
  //   (*_Pc2)(1) = qpc2.R_component_3() + (*q2)(1);
  //   (*_Pc2)(2) = qpc2.R_component_4() + (*q2)(2);

  //   // same for normal
  //   ::boost::math::quaternion<double> qnc(0, (*_relNc)(0), (*_relNc)(1), (*_relNc)(2));
  //   qnc = qq2 * qnc / qq2;
  //   (*_Nc)(0) = qnc.R_component_2();
  //   (*_Nc)(1) = qnc.R_component_3();
  //   (*_Nc)(2) = qnc.R_component_4();
  // }
  // else
  // {
  //   *_Pc2 = *_relPc2;
  //   *_Nc = *_relNc;
  // }

  LagrangianScleronomousR::computeh(q, z, y);
  y.setValue(0, distance());
  DEBUG_EXPR(y.display(););
  DEBUG_EXPR(display(););
  DEBUG_END("Lagrangian2d1DR::computeh(...)\n")
}
void Lagrangian2d1DR::display() const
{
  LagrangianR::display();

  std::cout << " _Pc1 :" << std::endl;
  if(_Pc1)
    _Pc1->display();
  else
    std::cout << " nullptr :" << std::endl;

  std::cout << " _Pc2 :" << std::endl;
  if(_Pc2)
    _Pc2->display();
  else
    std::cout << " nullptr :" << std::endl;

  std::cout << " _relPc1 :" << std::endl;
  if(_relPc1)
    _relPc1->display();
  else
    std::cout << " nullptr :" << std::endl;

  std::cout << " _relPc2 :" << std::endl;
  if(_relPc2)
    _relPc2->display();
  else
    std::cout << " nullptr :" << std::endl;

  std::cout << " _Nc :" << std::endl;
  if(_Nc)
    _Nc->display();
  else
    std::cout << " nullptr :" << std::endl;
  std::cout << " _relNc :" << std::endl;
  if(_relNc)
    _relNc->display();
  else
    std::cout << " nullptr :" << std::endl;

}