// Autor: Lenka Sokova
// xlogin: xsokov01

#include "hough_transform.h"

// Modify Hough transform
// by https://github.com/opencv/opencv/blob/4.x/modules/imgproc/src/hough.cpp
int computeNumangle( double min_theta, double max_theta, double theta_step )
{
    int numangle = cvFloor((max_theta - min_theta) / theta_step) + 1;
    // If the distance between the first angle and the last angle is
    // approximately equal to pi, then the last angle will be removed
    // in order to prevent a line to be detected twice.
    if ( numangle > 1 && fabs(CV_PI - (numangle-1)*theta_step) < theta_step/2 )
        --numangle;
    return numangle;
}

void countAngles(vector<float> &trigtab, int numAngle, float theta, float irho){

    for( int n = 0; n < numAngle; n++ )
    {
        trigtab[n*2] = (float)(cos((double)n*theta) * irho);
        trigtab[n*2+1] = (float)(sin((double)n*theta) * irho);
    }
}

void processFindLine(int maxN, int xFlag,  int k, int i, int j, int width, int height, uchar *ptrMask, Point (&lineEnd)[2], int lineGap, int dx0, int dy0, int y0, int x0){
    vector<uchar> deleted;
    int shift = 16;

    int gap = 0, x = x0, y = y0, dx = dx0, dy = dy0;

    if (k > 0)
        dx = -dx, dy = -dy;

    // walk along the line using fixed-point arithmetic,
    // stop at the image border or in case of too big gap
    for (;; x += dx, y += dy) {

        uchar *mdata;
        int i1, j1;

        if (xFlag) {
            j1 = x;
            i1 = y >> shift;
        }
        else {
            j1 = x >> shift;
            i1 = y;
        }

        if (j1 < 0 || j1 >= width || i1 < 0 || i1 >= height)
            break;

        mdata = ptrMask + i1 * width + j1;

        if (*mdata)
        {

            gap = 0;
            lineEnd[k].y = i1;
            lineEnd[k].x = j1;
            deleted.push_back(*mdata);
        }
        else if (++gap > lineGap){
            break;
        }
    }

    // delete processed points
    for(auto &sur : deleted)
    {
        sur = 0;
    }
}




void houghProbabilisticTransform(Mat &image, vector<Vec4i> &lines, float rho, float theta,
                                         int threshold, int lineLength, int lineGap, int linesMax){
    Point pt;
    float irho = 1 / rho;

    CV_Assert( image.type() == CV_8UC1 );

    int width = image.cols;
    int height = image.rows;

    int numAngle = computeNumangle(0.0, CV_PI, theta);
    int numRho = cvRound(((width + height) * 2 + 1) / rho);

    Mat accum = Mat::zeros( numAngle, numRho, CV_32SC1 );

    Mat mask( height, width, CV_8UC1 );
    std::vector<float> trigtab(numAngle*2);


    countAngles(trigtab, numAngle, theta, irho);

    uchar* ptrMask = mask.ptr();
    std::vector<Point> nonZeroPoints;

    // collect non-zero image points
    for( pt.y = 0; pt.y < height; pt.y++ )
    {
        const uchar* data = image.ptr(pt.y);
        uchar* mdata = mask.ptr(pt.y);
        for( pt.x = 0; pt.x < width; pt.x++ )
        {
            if( data[pt.x] )
            {
                mdata[pt.x] = (uchar)1;
                nonZeroPoints.push_back(pt);
            }
            else
                mdata[pt.x] = 0;
        }
    }


    const float* ttab = &trigtab[0];
    int count = nonZeroPoints.size();

    RNG rng((uint64)-1);


    // process non zero points
    while (count > 0){

        // choose random point out of the remaining ones
        int idx = rng.uniform(0, count);
        int max_val = threshold-1, max_n = 0;
        Point point = nonZeroPoints[idx];

        nonZeroPoints[idx] = nonZeroPoints[count-1];

        if (--count < 0)
            return;

        int maxValue = threshold-1, maxN = 0;
        Point lineEnd[2];
        const int shift = 16;

        int i = point.y, j = point.x, k, x0, y0, dx0, dy0, xFlag;
        float a, b;
        int good_line;

        if (!ptrMask[i* width + j])
            continue;

// Update intersections in param space
//-------------------------------------------------------
        int* adata = accum.ptr<int>();

        for (int n = 0; n < numAngle; n++, adata += numRho)
        {
            int r = cvRound(j * ttab[n*2] + i * ttab[n*2+1]);
            r += (numRho - 1)/2;
            int value = ++adata[r];


            if (maxValue < value){
                maxValue = value;
                maxN = n;
            }
        }
//-------------------------------------------------------

        if(maxValue < threshold){
            continue;
        }

        a = -ttab[maxN*2 + 1];
        b = ttab[maxN*2];
        x0 = j;
        y0 = i;


        if( fabs(a) > fabs(b) )
        {
            xFlag = 1;
            dx0 = a > 0 ? 1 : -1;
            dy0 = cvRound( b*(1 << shift)/fabs(a) );
            y0 = (y0 << shift) + (1 << (shift-1));
        }
        else
        {
            xFlag = 0;
            dy0 = b > 0 ? 1 : -1;
            dx0 = cvRound( a*(1 << shift)/fabs(b) );
            x0 = (x0 << shift) + (1 << (shift-1));
        }

        // process Line in both directions
        processFindLine(maxN, xFlag, 0, i, j, width, height, ptrMask, lineEnd, lineGap, dx0, dy0, y0, x0);
        processFindLine(maxN, xFlag, 1, i, j, width, height, ptrMask, lineEnd, lineGap, dx0, dy0, y0, x0);

        // check the size of the line
        good_line = std::abs(lineEnd[1].x - lineEnd[0].x) >= lineLength ||
                    std::abs(lineEnd[1].y - lineEnd[0].y) >= lineLength;


        if( good_line ) {
            Vec4i lr(lineEnd[0].x, lineEnd[0].y, lineEnd[1].x, lineEnd[1].y);
            lines.push_back(lr);

            if ((int) lines.size() >= linesMax){
                return;
            }
        }
    }
}
