#include "BinocularPolyFit.h"

using namespace std;
using namespace cv;

BinocularPolyFit::BinocularPolyFit(const PlType& plType)
    : PolyFit(plType)
    , binocularUnknowns(0)
    , binocularCalibrated(false)
{
    binocularUnknowns = 2 * (PolyFit::unknowns - 1) + 1;
}

string BinocularPolyFit::description() const
{
    return "BI" + PolyFit::description();
}

bool BinocularPolyFit::calibrate(vector<CollectionTuple>& calibrationTuples, QString& errorMsg)
{
    PointVector leftPupil;
    PointVector rightPupil;
    PointVector gaze;

    for (const auto& tuple : calibrationTuples) {
        leftPupil.insert(PolyFit::normalize(tuple.lEye.pupil.center, tuple.lEye.input));
        rightPupil.insert(PolyFit::normalize(tuple.rEye.pupil.center, tuple.rEye.input));
        gaze.insert(PolyFit::normalize(tuple.field.collectionMarker.center, tuple.field.input));
    }

    if (gaze.size() < binocularUnknowns) {
        errorMsg = QString("Not enough calibration points (%1/%2).").arg(gaze.size()).arg(binocularUnknowns);
        return false;
    }

    binocularCalibrated = calibrate(plType, leftPupil.x, leftPupil.y, rightPupil.x, rightPupil.y, gaze.x, bcx, errorMsg)
        && calibrate(plType, leftPupil.x, leftPupil.y, rightPupil.x, rightPupil.y, gaze.y, bcy, errorMsg);

    if (!binocularCalibrated)
        return binocularCalibrated;

    const auto& ref = calibrationTuples[0];
    float threshold = 0.01f * hypot(ref.field.width, ref.field.height);
    vector<CollectionTuple> inliers;
    for (const auto& tuple : calibrationTuples) {
        GazeEstimate l, r, b;
        l.pupilConfidence = tuple.lEye.pupil.confidence;
        r.pupilConfidence = tuple.rEye.pupil.confidence;
        estimateBinocular2d(tuple, l, r, b);
        Point2f gt = { tuple.field.collectionMarker.center.x, tuple.field.collectionMarker.center.y };
        if (cv::norm(gt - b.gp) < threshold)
            inliers.push_back(tuple);
    }

    binocularCalibrated = calibrateWithInliers(inliers, errorMsg);
    return binocularCalibrated;
}

bool BinocularPolyFit::calibrateWithInliers(std::vector<CollectionTuple>& calibrationTuples, QString& errorMsg)
{
    if (!PolyFit::calibrate(calibrationTuples, errorMsg))
        return false;

    PointVector leftPupil;
    PointVector rightPupil;
    PointVector gaze;

    for (const auto& tuple : calibrationTuples) {
        leftPupil.insert(PolyFit::normalize(tuple.lEye.pupil.center, tuple.lEye.input));
        rightPupil.insert(PolyFit::normalize(tuple.rEye.pupil.center, tuple.rEye.input));
        gaze.insert(PolyFit::normalize(tuple.field.collectionMarker.center, tuple.field.input));
    }

    if (gaze.size() < binocularUnknowns) {
        errorMsg = QString("Not enough calibration points (%1/%2).").arg(gaze.size()).arg(binocularUnknowns);
        return false;
    }

    return calibrate(plType, leftPupil.x, leftPupil.y, rightPupil.x, rightPupil.y, gaze.x, bcx, errorMsg)
        && calibrate(plType, leftPupil.x, leftPupil.y, rightPupil.x, rightPupil.y, gaze.y, bcy, errorMsg);
}

bool BinocularPolyFit::calibrate(const PlType& type, const Mat& lx, const Mat& ly, const Mat& rx, const Mat& ry, const Mat& z, Mat1d& c, QString& errorMsg)
{

    Mat ones = Mat::ones(lx.rows, 1, DataType<double>::type);
    Mat lpl = buildMat(type, lx, ly);
    Mat rpl = buildMat(type, rx, ry);
    Mat tmp;
    Mat A;
    hconcat(ones, lpl, tmp);
    hconcat(tmp, rpl, A);

    z.convertTo(z, A.type());

    bool result = true;
    try {
        result &= solve(A, z, c, DECOMP_SVD);
    } catch (...) {
        errorMsg = "cv::solve exception";
        result = false;
    }

    if (result) {
        // TODO: is this still necessary?
        for (int i = 0; i < c.rows; i++)
            if (std::isnan(c(i))) {
                errorMsg = "found NaN coefficients.";
                result = false;
            }
    }

    return result;
}

void BinocularPolyFit::estimateBinocular2d(const DataTuple& tuple, GazeEstimate& left, GazeEstimate& right, GazeEstimate& binocular)
{
    (void)tuple;
    bool leftValid = left.pupilConfidence >= minPupilConfidence;
    bool rightValid = right.pupilConfidence >= minPupilConfidence;

    auto evaluate = [&](DataTuple tuple, const Mat1d& c) {
        Point2d nlp = PolyFit::normalize(tuple.lEye.pupil.center, tuple.lEye.input);
        Point2d nrp = PolyFit::normalize(tuple.rEye.pupil.center, tuple.rEye.input);
        double res = c.at<double>(0); // constant contribution
        res += PolyFit::evaluateVariable(plType, nlp.x, nlp.y, c.rowRange(1, PolyFit::unknowns)); // left eye contribution
        res += PolyFit::evaluateVariable(plType, nrp.x, nrp.y, c.rowRange(PolyFit::unknowns, binocularUnknowns)); // right eye contribution
        return res;
    };

    if (leftValid && rightValid && binocularCalibrated) {
        binocular.gp = { static_cast<float>(tuple.field.width * evaluate(tuple, bcx)), static_cast<float>(tuple.field.height * evaluate(tuple, bcy)) };
        binocular.valid = true;
        binocular.pupilConfidence = 0.5 * (left.pupilConfidence + right.pupilConfidence);
    } else {
        if (leftValid)
            binocular = left;
        if (rightValid)
            binocular = right;
    }
}