//
// Created by martin on 18. 8. 2021.
//

#include "meanShift.h"



const double radius = 5.0;

double distance(const std::pair<int, int> &left, const std::pair<int, int> &right) {
    return sqrt((left.first - right.first) * (left.first - right.first) +
                (left.second - right.second) * (left.second - right.second));
}



std::pair<int, int> centre(const std::set<std::pair<int, int>> &points) {
    int x = 0, y = 0;
    for (auto point: points) {
        x += point.first;
        y += point.second;
    }
    x /= points.size();
    y /= points.size();

    return std::make_pair(x,y);
}

double centre(const std::set<double> &points) {
    double angle = 0;

    for (auto point: points) angle += point;

    return angle/points.size();
}

// group ellipses depending on which cluster they belong to and set their center coordinates to the cluster centre
std::vector<std::vector<ellipseParameters>> collect (const std::set<std::pair<int,int>> & clusters, const std::vector<ellipseParameters> & ellipses) {
    std::vector<std::vector<ellipseParameters>> clustered (clusters.size(),std::vector<ellipseParameters>());
    for ( ellipseParameters ellipse : ellipses) {
        int i = -1;
        for (auto center : clusters) {
            i++;
            if (distance(center, std::make_pair(ellipse.center.x,ellipse.center.y)) <= radius) {
                ellipse.center = cv::Point2i(center.first,center.second);

                clustered[i].emplace_back(ellipse);
                break;
            }
        }
    }

    return clustered;
}

// group ellipses by their angle and set the ellipse angle to the cluster angle they belong to
std::vector<std::vector<ellipseParameters>> collect (const std::set<double> & clusters, const std::vector<ellipseParameters> & ellipses) {
    std::vector<std::vector<ellipseParameters>> clustered (clusters.size(),std::vector<ellipseParameters>());
    for ( ellipseParameters ellipse : ellipses) {
        int i = -1;
        for (auto angle : clusters) {
            i++;
            if (abs(angle - ellipse.theta) <= radius/180*M_PI) {
                ellipse.theta = angle;
                clustered[i].emplace_back(ellipse);
                break;
            }
        }
    }

    return clustered;
}

//cluster ellipses depending on their center points
std::vector<std::vector<ellipseParameters>> clusterCenter(std::vector<ellipseParameters> &ellipses) {
    std::set<std::pair<int, int>> clusters, previous;
    for (ellipseParameters ellipse: ellipses) {
        clusters.emplace(std::make_pair(ellipse.center.x, ellipse.center.y));
    }

    while (previous != clusters) {
        previous = clusters;
        std::set<std::pair<int, int>> nextClusters;
        for (auto currentCentroid : clusters) {
            std::set<std::pair<int, int>> cluster;
            for (auto centroid: clusters) {
                if (distance(currentCentroid, centroid) <= radius) {
                    cluster.emplace(centroid);
                }
            }
            nextClusters.emplace(centre(cluster));

        }
        clusters = nextClusters;
    }

    return collect(clusters,ellipses);
}


//cluster ellipses depending on their angle
std::vector<std::vector<ellipseParameters>> clusterAngle(std::vector<ellipseParameters> &ellipses) {

    std::set<double> clusters, previous;
    for (ellipseParameters ellipse: ellipses) {
        clusters.emplace(ellipse.theta);
    }

    while (previous != clusters) {
        previous = clusters;
        std::set<double> nextClusters;
        for (auto currentCentroid : clusters) {
            std::set<double> cluster;
            for (auto centroid: clusters) {
                if (abs(currentCentroid- centroid) <= radius*M_PI/180) {
                    cluster.emplace(centroid);
                }
            }
            nextClusters.emplace(centre(cluster));

        }
        clusters = nextClusters;
    }

    return collect(clusters,ellipses);

}

// cluster ellipses depending on their axis
std::vector<ellipseParameters> clusterAxis(std::vector<ellipseParameters> &ellipses) {
    std::set<std::pair<int, int>> clusters, previous;
    for (ellipseParameters ellipse: ellipses) {
        clusters.emplace(std::make_pair(ellipse.major, ellipse.minor));
    }

    while (previous != clusters) {
        previous = clusters;
        std::set<std::pair<int, int>> nextClusters;
        for (auto currentCentroid : clusters) {
            std::set<std::pair<int, int>> cluster;
            for (auto centroid: clusters) {
                if (distance(currentCentroid, centroid) <= radius) {
                    cluster.emplace(centroid);
                }
            }
            nextClusters.emplace(centre(cluster));

        }
        clusters = nextClusters;
    }

    cv::Point2i centerPoint = ellipses[0].center;
    double angle = ellipses[0].theta;

    std::vector<ellipseParameters> clusteredEllipses;

    for (auto [major,minor] : clusters) {
        ellipseParameters params;
        params.center = centerPoint;
        params.major = major;
        params.minor = minor;
        params.theta = angle;
        clusteredEllipses.emplace_back(params);
    };

    return clusteredEllipses;

}



std::vector<ellipseParameters> clusterAll(std::vector<ellipseParameters> &candidates) {
    std::vector<ellipseParameters> ellipses;
    if (!candidates.empty()) {
        auto centers = clusterCenter(candidates); //cluster by center
        for (std::vector<ellipseParameters> clusteredAngle: centers) { // each center may have ellipses of different orientation
            auto anglesCentroid = clusterAngle(clusteredAngle);
            for (std::vector<ellipseParameters> angled: anglesCentroid) { // each ellipse orientation may have ellipses of different size
                auto clusteredEllipses = clusterAxis(angled);
                ellipses.insert(ellipses.end(), clusteredEllipses.begin(),
                                clusteredEllipses.end());
            }
        }
    }
    return ellipses;
}
