#include <limits> // std::numeric_limits
#include <opencv2/imgproc/imgproc.hpp> // cv::line, cv::circle
#include <openpose/utilities/errorAndLog.hpp>
#include <openpose/utilities/fastMath.hpp>
#include <openpose/utilities/keypoint.hpp>

namespace op
{
    const std::string errorMessage = "The Array<float> is not a RGB image. This function is only for array of dimension: [sizeA x sizeB x 3].";

    float getDistance(const float* keypointPtr, const int elementA, const int elementB)
    {
        try
        {
            const auto pixelX = keypointPtr[elementA*3] - keypointPtr[elementB*3];
            const auto pixelY = keypointPtr[elementA*3+1] - keypointPtr[elementB*3+1];
            return std::sqrt(pixelX*pixelX+pixelY*pixelY);
        }
        catch (const std::exception& e)
        {
            error(e.what(), __LINE__, __FUNCTION__, __FILE__);
            return -1.f;
        }
    }

    void scaleKeypoints(Array<float>& keypoints, const float scale)
    {
        try
        {
            scaleKeypoints(keypoints, scale, scale);
        }
        catch (const std::exception& e)
        {
            error(e.what(), __LINE__, __FUNCTION__, __FILE__);
        }
    }

    void scaleKeypoints(Array<float>& keypoints, const float scaleX, const float scaleY)
    {
        try
        {
            if (scaleX != 1. && scaleY != 1.)
            {
                // Error check
                if (!keypoints.empty() && keypoints.getSize(2) != 3)
                    error(errorMessage, __LINE__, __FUNCTION__, __FILE__);
                // Get #people and #parts
                const auto numberPeople = keypoints.getSize(0);
                const auto numberParts = keypoints.getSize(1);
                // For each person
                for (auto person = 0 ; person < numberPeople ; person++)
                {
                    // For each body part
                    for (auto part = 0 ; part < numberParts ; part++)
                    {
                        const auto finalIndex = 3*(person*numberParts + part);
                        keypoints[finalIndex] *= scaleX;
                        keypoints[finalIndex+1] *= scaleY;
                    }
                }
            }
        }
        catch (const std::exception& e)
        {
            error(e.what(), __LINE__, __FUNCTION__, __FILE__);
        }
    }

    void scaleKeypoints(Array<float>& keypoints, const float scaleX, const float scaleY, const float offsetX, const float offsetY)
    {
        try
        {
            if (scaleX != 1. && scaleY != 1.)
            {
                // Error check
                if (!keypoints.empty() && keypoints.getSize(2) != 3)
                    error(errorMessage, __LINE__, __FUNCTION__, __FILE__);
                // Get #people and #parts
                const auto numberPeople = keypoints.getSize(0);
                const auto numberParts = keypoints.getSize(1);
                // For each person
                for (auto person = 0 ; person < numberPeople ; person++)
                {
                    // For each body part
                    for (auto part = 0 ; part < numberParts ; part++)
                    {
                        const auto finalIndex = keypoints.getSize(2)*(person*numberParts + part);
                        keypoints[finalIndex] = keypoints[finalIndex] * scaleX + offsetX;
                        keypoints[finalIndex+1] = keypoints[finalIndex+1] * scaleY + offsetY;
                    }
                }
            }
        }
        catch (const std::exception& e)
        {
            error(e.what(), __LINE__, __FUNCTION__, __FILE__);
        }
    }

    void renderKeypointsCpu(Array<float>& frameArray, const Array<float>& keypoints, const std::vector<unsigned int>& pairs,
                            const std::vector<float> colors, const float thicknessCircleRatio, const float thicknessLineRatioWRTCircle,
                            const float threshold)
    {
        try
        {
            if (!frameArray.empty())
            {
                // Array<float> --> cv::Mat
                auto frame = frameArray.getCvMat();

                // Security check
                if (frame.dims != 3 || frame.size[0] != 3)
                    error(errorMessage, __LINE__, __FUNCTION__, __FILE__);

                // Get frame channels
                const auto width = frame.size[2];
                const auto height = frame.size[1];
                const auto area = width * height;
                cv::Mat frameB{height, width, CV_32FC1, &frame.data[0]};
                cv::Mat frameG{height, width, CV_32FC1, &frame.data[area * sizeof(float) / sizeof(uchar)]};
                cv::Mat frameR{height, width, CV_32FC1, &frame.data[2 * area * sizeof(float) / sizeof(uchar)]};

                // Parameters
                const auto lineType = 8;
                const auto shift = 0;
                const auto numberColors = colors.size();
                const auto thresholdRectangle = 0.1f;
                const auto numberKeypoints = keypoints.getSize(1);
                const auto areaKeypoints = numberKeypoints * keypoints.getSize(2);

                // Keypoints
                for (auto person = 0 ; person < keypoints.getSize(0) ; person++)
                {
                    const auto personRectangle = getKeypointsRectangle(&keypoints[person*areaKeypoints], numberKeypoints, thresholdRectangle);
                    if (personRectangle.area() > 0)
                    {
                        const auto ratioAreas = fastMin(1.f, fastMax(personRectangle.width/(float)width, personRectangle.height/(float)height));
                        // Size-dependent variables
                        const auto thicknessRatio = fastMax(intRound(std::sqrt(area)*thicknessCircleRatio * ratioAreas), 2);
                        // Negative thickness in cv::circle means that a filled circle is to be drawn.
                        const auto thicknessCircle = (ratioAreas > 0.05 ? thicknessRatio : -1);
                        const auto thicknessLine = intRound(thicknessRatio * thicknessLineRatioWRTCircle);
                        const auto radius = thicknessRatio / 2;

                        // Draw lines
                        for (auto pair = 0 ; pair < pairs.size() ; pair+=2)
                        {
                            const auto index1 = (person * numberKeypoints + pairs[pair]) * keypoints.getSize(2);
                            const auto index2 = (person * numberKeypoints + pairs[pair+1]) * keypoints.getSize(2);
                            if (keypoints[index1+2] > threshold && keypoints[index2+2] > threshold)
                            {
                                const auto colorIndex = pairs[pair+1]*3; // Before: colorIndex = pair/2*3;
                                const cv::Scalar color{colors[colorIndex % numberColors],
                                                       colors[(colorIndex+1) % numberColors],
                                                       colors[(colorIndex+2) % numberColors]};
                                const cv::Point keypoint1{intRound(keypoints[index1]), intRound(keypoints[index1+1])};
                                const cv::Point keypoint2{intRound(keypoints[index2]), intRound(keypoints[index2+1])};
                                cv::line(frameR, keypoint1, keypoint2, color[0], thicknessLine, lineType, shift);
                                cv::line(frameG, keypoint1, keypoint2, color[1], thicknessLine, lineType, shift);
                                cv::line(frameB, keypoint1, keypoint2, color[2], thicknessLine, lineType, shift);
                            }
                        }

                        // Draw circles
                        for (auto part = 0 ; part < numberKeypoints ; part++)
                        {
                            const auto faceIndex = (person * numberKeypoints + part) * keypoints.getSize(2);
                            if (keypoints[faceIndex+2] > threshold)
                            {
                                const auto colorIndex = part*3;
                                const cv::Scalar color{colors[colorIndex % numberColors],
                                                       colors[(colorIndex+1) % numberColors],
                                                       colors[(colorIndex+2) % numberColors]};
                                const cv::Point center{intRound(keypoints[faceIndex]), intRound(keypoints[faceIndex+1])};
                                cv::circle(frameR, center, radius, color[0], thicknessCircle, lineType, shift);
                                cv::circle(frameG, center, radius, color[1], thicknessCircle, lineType, shift);
                                cv::circle(frameB, center, radius, color[2], thicknessCircle, lineType, shift);
                            }
                        }
                    }
                }
            }
        }
        catch (const std::exception& e)
        {
            error(e.what(), __LINE__, __FUNCTION__, __FILE__);
        }
    }

    Rectangle<float> getKeypointsRectangle(const float* keypointPtr, const int numberKeypoints, const float threshold)
    {
        try
        {
            if (numberKeypoints < 1)
                error("Number body parts must be > 0", __LINE__, __FUNCTION__, __FILE__);

            float minX = std::numeric_limits<float>::max();
            float maxX = 0.f;
            float minY = minX;
            float maxY = maxX;
            for (auto part = 0 ; part < numberKeypoints ; part++)
            {
                const auto score = keypointPtr[3*part + 2];
                if (score > threshold)
                {
                    const auto x = keypointPtr[3*part];
                    const auto y = keypointPtr[3*part + 1];
                    // Set X
                    if (maxX < x)
                        maxX = x;
                    if (minX > x)
                        minX = x;
                    // Set Y
                    if (maxY < y)
                        maxY = y;
                    if (minY > y)
                        minY = y;
                }
            }
            if (maxX >= minX && maxY >= minY)
                return Rectangle<float>{minX, minY, maxX-minX, maxY-minY};
            else
                return Rectangle<float>{};
        }
        catch (const std::exception& e)
        {
            error(e.what(), __LINE__, __FUNCTION__, __FILE__);
            return Rectangle<float>{};
        }
    }

    float getKeypointsArea(const float* keypointPtr, const int numberKeypoints, const float threshold)
    {
        try
        {
            return getKeypointsRectangle(keypointPtr, numberKeypoints, threshold).area();
        }
        catch (const std::exception& e)
        {
            error(e.what(), __LINE__, __FUNCTION__, __FILE__);
            return 0.f;
        }
    }

    int getBiggestPerson(const Array<float>& keypoints, const float threshold)
    {
        try
        {
            if (!keypoints.empty())
            {
                const auto numberPeople = keypoints.getSize(0);
                const auto numberKeypoints = keypoints.getSize(1);
                const auto area = numberKeypoints * keypoints.getSize(2);
                auto biggestPoseIndex = -1;
                auto biggestArea = -1.f;
                for (auto person = 0 ; person < numberPeople ; person++)
                {
                    const auto newPersonArea = getKeypointsArea(&keypoints[person*area], numberKeypoints, threshold);
                    if (newPersonArea > biggestArea)
                    {
                        biggestArea = newPersonArea;
                        biggestPoseIndex = person;
                    }
                }
                return biggestPoseIndex;
            }
            else
                return -1;
        }
        catch (const std::exception& e)
        {
            error(e.what(), __LINE__, __FUNCTION__, __FILE__);
            return -1;
        }
    }
}
