//
// Created by Martin Gargalovič on 11. 8. 2021.
//

#include <vector>
#include <opencv2/core.hpp>
#include <opencv2/core/types.hpp>
#include "ellipseSet.h"



#define min(a, b) a < b ? a : b
#define max(a, b) a > b ? a : b

const double M_1_2_PI = M_PI / 2;
const int MINPOINTS = 10;

/**
 * from lsd.cpp
 */
double angle_difff(double a, double b) {
    a -= b;
    while (a <= -M_PI) a += 2 * M_PI;
    while (a > M_PI) a -= 2 * M_PI;
    if (a < 0.0) a = -a;
    return a;
}


bool
lineAppr(int supportInliersCount, const std::vector<rectangle> &rectangles,
         const std::vector<int> &group, double inliersRatio) {
    double len = 0;
    for (auto line: group) {
        len += std::sqrt(
                std::pow(rectangles[line].x1 - rectangles[line].x2, 2) +
                std::pow(rectangles[line].y1 - rectangles[line].y2, 2));
    }
    return (supportInliersCount > len * inliersRatio);
}


std::vector<std::vector<cv::Point2i>>
getArc(const std::vector<std::vector<cv::Point2i>> &lines,
       std::vector<int> &group) {
    std::vector<std::vector<cv::Point2i>> arc;
    for (auto line: group) arc.emplace_back(lines[line]);
    return arc;
}


Eigen::Matrix<double, 6, 6>
scatterMatrixGroup(const std::vector<int> &group,
                   const std::vector<Eigen::Matrix<double, 6, 6>> &matrices) {
    Eigen::Matrix<double, 6, 6> out;
    out.setZero();
    for (int segment: group) {
        out += matrices[segment];
    }
    return out;
}

Eigen::Matrix<double, 6, 6>
scatterMatrix(const std::vector<cv::Point2i> &supportGroup) {

    Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> D, DT;
    D.resize(supportGroup.size(), 6);
    DT.resize(6, supportGroup.size());

    int current = 0;
    for (auto point: supportGroup) {
        D(current, 0) = point.x * point.x;
        D(current, 1) = point.x * point.y;
        D(current, 2) = point.y * point.y;
        D(current, 3) = point.x;
        D(current, 4) = point.y;
        D(current, 5) = 1;
        current++;
    }

    Eigen::Matrix<double, 6, 6> S;
    S.setZero();

    //S = D * D^T
    for (int i = 0; i < 6; i++)
        for (int j = 0; j < 6; j++) {
            for (size_t k = 0; k < supportGroup.size(); k++)
                S(i, j) +=
                        D(k, i) *
                        D(k, j);
        }
    return S;
}


Eigen::Matrix<double, 6, 1> fitEllipse(const Eigen::Matrix<double, 6, 6> &S) {
    Eigen::Matrix<double, 6, 6> C;
    C.setZero();
    C(2, 0) = -1;
    C(0, 2) = -1;
    C(1, 1) = 2;

    auto decomp = S.inverse() * C;

    Eigen::EigenSolver<Eigen::Matrix<double, 6, 6>> solver;
    solver.compute(decomp, true);

    auto vectors = solver.eigenvectors().real();


    Eigen::Matrix<double, 6, 1> eq;
    eq = vectors.col(0);

    return eq;
}

double semiAxis(Eigen::Matrix<double, 6, 1> &ellipseEQ, bool major) {
    //general ellipse equation : Ax^2 + Bxy + Cy^2 + Dx + Ey + F
    double A, B, C, D, E, F;
    A = ellipseEQ(0);
    B = ellipseEQ(1);
    C = ellipseEQ(2);
    D = ellipseEQ(3);
    E = ellipseEQ(4);
    F = ellipseEQ(5);

    //(1*major) + (-1 * ! major) to represent the +- in the equation
    double left = 2 * ((A * E * E) + (C * D * D) - (B * D * E) +
                       (B * B - 4 * A * C) * F);
    double right = (A + C) + (((1 * major) - (1 * !major)) *
                              sqrt((A - C) * (A - C) + B * B));

    double top = -std::sqrt(left * right);

    double bot = B * B - 4 * A * C;

    return top / bot;
}

cv::Point2i centre(Eigen::Matrix<double, 6, 1> &ellipseEQ) {
    double A = ellipseEQ(0);
    double B = ellipseEQ(1);
    double C = ellipseEQ(2);
    double D = ellipseEQ(3);
    double E = ellipseEQ(4);

    int x = (2 * C * D - B * E) / (B * B - 4 * A * C);
    int y = (2 * A * E - B * D) / (B * B - 4 * A * C);
    return cv::Point2i(x, y);
}

double angle(Eigen::Matrix<double, 6, 1> &ellipseEQ) {
    double A = ellipseEQ(0);
    double B = ellipseEQ(1);
    double C = ellipseEQ(2);

    if (B == 0) {
        if (A > C) {
            return M_PI / 2; // if A > C then return 90 degrees
        } else {
            return 0;
        }
    }
    return atan((1 / B) * (C - A - sqrt((A - C) * (A - C) + B * B)));
}


bool getParam(Eigen::Matrix<double, 6, 1> &eq, ellipseParameters &params) {
    double major = semiAxis(eq, true);
    double minor = semiAxis(eq, false);
    double ellipse_angle = angle(eq);
    cv::Point2i center = centre(eq);


    if (std::isnan(major) || std::isnan(minor) ||
        std::isnan(ellipse_angle))
        return false;
    params.center = center;
    params.major = major;
    params.minor = minor;
    params.theta = ellipse_angle;
    if (minor > major) {
        params.major = minor;
        params.minor = major;
        if (ellipse_angle < 0)
            params.theta += M_1_2_PI;
        else
            params.theta -= M_1_2_PI;
        if (ellipse_angle < -M_1_2_PI)
            params.theta += M_PI;
        if (ellipse_angle > M_1_2_PI)
            params.theta -= M_PI;

    }
    return true;
}


bool
rosinDistance(ellipseParameters &params, const cv::Point2i &point,
              double threshold) {
    double a_e = params.major, b_e = params.minor;
    double F = a_e * a_e - b_e * b_e;

    auto relativePoint = point - params.center;

    //rotating the point around the centre of the ellipse as rosinDistance is counted in canonical ellipse
    double x = cos(-params.theta) * relativePoint.x -
               sin(-params.theta) * relativePoint.y;
    double y = sin(-params.theta) * relativePoint.x +
               cos(-params.theta) * relativePoint.y;

    double X = x * x;
    double Y = y * y;

    double discriminant = (F + X + Y) * (F + X + Y) - 4 * X * F;
    double A = ((F + X + Y) - sqrt(discriminant)) / 2.0;
    double a_h = sqrt(A);
    double b_h = sqrt(F - A);

    double x_i = a_h * sqrt((a_e * a_e * (b_e * b_e + b_h * b_h)) /
                            (a_h * a_h * b_e * b_e + a_e * a_e * b_h * b_h));
    double y_i = (b_e * b_h * sqrt(a_e * a_e - a_h * a_h)) /
                 sqrt(a_h * a_h * b_e * b_e + a_e * a_e * b_h * b_h);

    std::vector<double> distances{distance(x, y, x_i, y_i),
                                  distance(x, y, x_i, -y_i),
                                  distance(x, y, -x_i, y_i),
                                  distance(x, y, -x_i, -y_i)};
    return *min_element(distances.begin(), distances.end()) < threshold;
}

bool calcInliers(Eigen::Matrix<double, 6, 1> &coefficients,
                 std::vector<std::vector<cv::Point2i>> &arc,
                 cv::Mat &angles, int &supportInliersCount,
                 double distanceTolerance,
                 int polarity, double inliersRatio,
                 const std::vector<rectangle> &rectangles,
                 const std::vector<int> &group) {

    double a = coefficients(0);
    double b = coefficients(1);
    double c = coefficients(2);
    double d = coefficients(3);
    double e = coefficients(4);

    int inliersCount = 0;
    ellipseParameters params;
    getParam(coefficients, params);
    for (auto line: arc) {

        for (cv::Point2i point: line) {

            double normX = (a * point.x + b * point.y / 2 + d / 2);
            double normY = (c * point.y + b * point.x / 2 + e / 2);


            double pointNorm = atan2(normY, normX);
            ++inliersCount;
            double angleAt = angles.at<float>(point);

            if (!rosinDistance(params, point, distanceTolerance)) continue;

            if (angle_difff((polarity > 0) * -M_PI + pointNorm, angleAt) <=
                M_PI / 8) {
                ++supportInliersCount;
            }
        }
    }
    return lineAppr(supportInliersCount, rectangles, group,
                    inliersRatio);
}


bool
validateEllipse(std::vector<std::vector<cv::Point2i>> &arc,
                Eigen::Matrix<double, 6, 6> &S, cv::Mat &angles,
                int &support_inliers_count,
                double distanceTolerance,
                ellipseParameters &parameters, int polarity,
                double inliersRatio, const std::vector<int> &group,
                const std::vector<rectangle> &rectangles, bool setParams) {
    assert (!arc.empty());

    Eigen::Matrix<double, 6, 1> ellipseEQ = fitEllipse(S);
    ellipseParameters params;

    if (!getParam(ellipseEQ, params)) return false;

// check whether the ellipse is smaller then the size of the image and if the midpoint of the ellipse in the image
    if (min(params.major, params.minor)
        <3 * distanceTolerance || (max(params.major, params.minor))>(min(
                                                                             angles.size().width,
                                                                             angles.size().width)) ||
        params.center.x < 0 || params.center.x > angles.size().width ||
        params.center.y < 0 || params.center.y > angles.size().height) {
        return false;
    }

    // calc the support inliers
    if (!calcInliers(ellipseEQ, arc, angles, support_inliers_count,
                     distanceTolerance, polarity,
                     inliersRatio, rectangles, group))
        return false;

    if (setParams) {
        parameters = params;
    }

    return true;
}


bool
validateEllipse(std::vector<std::vector<cv::Point2i>> &lines,
                std::vector<std::vector<int>> &groups,
                size_t group1, size_t group2, cv::Mat &angles,
                int &support_inliers_group1_count,
                int &support_inliers_group2_count,
                double distanceTolerance,
                ellipseParameters &parameters, int polarity,
                double inliersRatio,
                std::vector<Eigen::Matrix<double, 6, 6>> &matrices,
                const std::vector<rectangle> &rectangles) {

    auto arc1 = getArc(lines, groups[group1]);
    auto arc2 = getArc(lines, groups[group2]);

    Eigen::Matrix<double, 6, 6> S = matrices[group1] + matrices[group2];

    std::vector<cv::Point2i> upInliers, downInliers;
    if (!validateEllipse(arc1, S, angles, support_inliers_group1_count,
                         distanceTolerance,
                         parameters, polarity, inliersRatio, groups[group1],
                         rectangles, false))
        return false;

    if (!validateEllipse(arc2, S, angles, support_inliers_group2_count,
                         distanceTolerance,
                         parameters, polarity, inliersRatio, groups[group2],
                         rectangles, true))
        return false;

    return true;
}


cv::Point2i normalize(cv::Point2i &vector) {
    double inversqrt = 1 / sqrt(vector.x * vector.x + vector.y * vector.y);
    return cv::Point2i(vector.x * inversqrt, vector.y * inversqrt);
}

void rotate(cv::Point2i &point, double angle) {
    int x = point.x;
    int y = point.y;

    point.x = cos(angle) * x - sin(angle) * y;
    point.y = sin(angle) * x + cos(angle) * y;
}


bool regionConstraint(std::vector<std::vector<cv::Point2i>> &lines,
                      std::vector<std::vector<int>> &groups, int group1,
                      int group2, const std::vector<rectangle> &rectangles,
                      double threshold) {
    auto arc1 = getArc(lines, groups[group1]);
    auto arc2 = getArc(lines, groups[group2]);


    size_t gie = groups[group1].size() - 1;
    cv::Point2i Psi(rectangles[groups[group1][0]].x2,
                    rectangles[groups[group1][0]].y2);
    cv::Point2i Pei(rectangles[groups[group1][gie]].x1,
                    rectangles[groups[group1][gie]].y1);

    std::pair<int, int> vecLs = std::make_pair(rectangles[groups[group1][0]].dx,
                                               rectangles[groups[group1][0]].dy);
    std::pair<int, int> vecLe = std::make_pair(
            rectangles[groups[group1][gie]].dx,
            rectangles[groups[group1][gie]].dy);

    size_t group2_size = groups[group2].size() - 1;
    cv::Point2i Psj(rectangles[groups[group2][0]].x2,
                    rectangles[groups[group2][0]].y2);
    cv::Point2i Pej(rectangles[groups[group2][group2_size]].x1,
                    rectangles[groups[group2][group2_size]].y1);

    auto vec1 = Pej - Psi;
    auto vec2 = Psj - Pei;

    auto PsPe = Pei - Psi;
    //rotate PsPe by 90 deg clockwise
    rotate(PsPe, -M_1_2_PI);
    PsPe = normalize(PsPe);
    cv::Point2i M1 = (Psi + Pei) / 2;
    cv::Point2i M2 = (Psj + Pej) / 2;
    auto vecM = M2 - M1;
    double scalPM = PsPe.x * vecM.x + PsPe.y * vecM.y;


    if (vecLs.first * vec1.x + vecLs.second * vec1.y < threshold) return false;
    if (vecLe.first * vec2.x + vecLe.second * vec2.y < threshold) return false;
    return (((-1) * rectangles[group1].polarity) * scalPM) > threshold;
}


int getPolarity(std::vector<std::vector<int>> &groups,
                std::vector<rectangle> &rectangles, size_t group) {
    int polarity = 0;
    for (int i: groups[group]) {
        polarity += rectangles[i].polarity;
    }

    return polarity > 0 ? 1 : polarity != 0 ? -1 : 0;
}


double
goodness(ellipseParameters param, int support_inliers_count, double coverage) {

    double B = M_PI * (1.5 * (param.major + param.minor) -
                       std::sqrt(param.major * param.minor));

    return std::sqrt(support_inliers_count / B * coverage / (2 * M_PI));
}


std::vector<ellipseParameters>
ellipseSet(std::vector<std::vector<cv::Point2i>> &lines,
           std::vector<rectangle> &rectangles, cv::Mat &angles,
           std::vector<std::vector<int>> &groups, double distanceTolerance,
           double scale, int specificPolarity, std::vector<double> &coverages,
           double inliersRatio, double Tr, double Tcover) {

    std::vector<ellipseParameters> candidates;
    std::vector<Eigen::Matrix<double, 6, 6>> matrixGroups;


    {
        std::vector<Eigen::Matrix<double, 6, 6>> matrices;
        for (auto line: lines) {
            matrices.emplace_back(scatterMatrix(line));
        }

        for (auto group: groups)
            matrixGroups.push_back(scatterMatrixGroup(group, matrices));
    }

    // calculates equation of ellipse if the coverage of the line is larger than treshold: 2/3*M_PI
    for (size_t i = 0; i < groups.size(); i++) {
        if (coverages[i] < Tcover) continue;
        int polarity = getPolarity(groups, rectangles, i);
        if (specificPolarity != 0 && specificPolarity != polarity) continue;
        ellipseParameters ellipse;
        auto arc = getArc(lines, groups[i]);
        int arcSize = 0;
        for (auto line: arc) arcSize += line.size();
        if (arcSize < MINPOINTS) continue;
        auto S = matrixGroups[i];
        std::vector<cv::Point2i> inliers;
        int support_inliers_cnt = 0;
        if (validateEllipse(arc, S, angles, support_inliers_cnt,
                            distanceTolerance,
                            ellipse, polarity, inliersRatio, groups[i],
                            rectangles, true)) {
            if (scale != 1.0) {
                ellipse.major /= scale;
                ellipse.minor /= scale;
                ellipse.center.x /= scale;
                ellipse.center.y /= scale;
            }
            candidates.emplace_back(ellipse);
        }
    }

    // calculates ellipse equation with by using 2 groups of arc-segments
    for (size_t i = 0; i < groups.size(); i++) {
        for (size_t j = i + 1; j < groups.size(); j++) {

            if (getPolarity(groups, rectangles, i) !=
                getPolarity(groups, rectangles, j))
                continue;
            int polarity = getPolarity(groups, rectangles, i);

            if (specificPolarity != 0 && polarity != specificPolarity)
                continue;
            // verification of the groups of line segments can be put into 1 ellipse still needed (polarity and region restriction)


            if (!regionConstraint(lines, groups, i, j, rectangles,
                                  -3 * distanceTolerance))
                continue;

            if (double coverage = coverages[i] + coverages[j]; coverage >
                                                               Tcover) {
                ellipseParameters ellipse;
                int support_inliers_group1_count = 0;
                int support_inliers_group2_count = 0;
                if (validateEllipse(lines, groups, i, j, angles,
                                    support_inliers_group1_count,
                                    support_inliers_group2_count,
                                    distanceTolerance,
                                    ellipse, polarity, inliersRatio,
                                    matrixGroups,
                                    rectangles)) {
                    if (scale != 1.0) {
                        ellipse.major /= scale;
                        ellipse.minor /= scale;
                        ellipse.center.x /= scale;
                        ellipse.center.y /= scale;
                    }

                    if (goodness(ellipse, support_inliers_group1_count +
                                          support_inliers_group2_count,
                                 coverage) >= Tr)
                        candidates.emplace_back(ellipse);
                }
            }
        }
    }
    return candidates;
}