#include "RiaImageTools.h"

#include "cvfAssert.h"

#include <QColor>
#include <QImage>

#include <algorithm>
#include <limits>

//--------------------------------------------------------------------------------------------------
/// Meijster, Roerdink, Hesselink
/// A  GENERAL  ALGORITHM  FOR  COMPUTING  DISTANCE TRANSFORMS IN LINEAR TIME
/// http://fab.cba.mit.edu/classes/S62.12/docs/Meijster_distance.pdf
/// Currently Euclidean only, but can be easily extended by replacing the lambda functions.
//--------------------------------------------------------------------------------------------------
void RiaImageTools::distanceTransform2d( std::vector<std::vector<unsigned int>>& image )
{
    if ( image.empty() )
    {
        return;
    }
    if ( image.front().empty() )
    {
        return;
    }
    const int64_t M = (int64_t)image.size();
    const int64_t N = (int64_t)image.front().size();

    int64_t infVal = M + N;
    CVF_ASSERT( infVal <= std::numeric_limits<unsigned int>::max() );

    // First phase
    std::vector<std::vector<int64_t>> g( M );

#pragma omp parallel for
    for ( int64_t x = 0; x < M; ++x )
    {
        g[x].resize( N, infVal );
        if ( image[x][0] )
        {
            g[x][0] = 0;
        }
        for ( int64_t y = 1; y < N - 1; ++y )
        {
            if ( image[x][y] )
            {
                g[x][y] = 0;
            }
            else
            {
                g[x][y] = 1 + g[x][y - 1];
            }
        }
        for ( int64_t y = N - 2; y >= 0; --y )
        {
            if ( g[x][y + 1] < g[x][y] )
            {
                g[x][y] = 1 + g[x][y + 1];
            }
        }
    }

    auto f = []( int64_t x, int64_t i, const std::vector<std::vector<int64_t>>& g, int64_t y ) {
        return ( x - i ) * ( x - i ) + g[i][y] * g[i][y];
    };

    auto sep = []( int64_t i, int64_t u, const std::vector<std::vector<int64_t>>& g, int64_t y ) {
        if ( i == u ) return (int64_t)0;

        int64_t numerator = u * u - i * i + g[u][y] * g[u][y] - g[i][y] * g[i][y];
        int64_t divisor   = 2 * ( u - i );
        return numerator / divisor;
    };

    // Second phase
#pragma omp parallel for
    for ( int64_t y = 0; y < N; ++y )
    {
        int64_t              q = 0;
        std::vector<int64_t> s( std::max( N, M ), (int64_t)0 );
        std::vector<int64_t> t( std::max( N, M ), (int64_t)0 );

        for ( int64_t u = 1; u < M - 1; ++u )
        {
            while ( q >= 0 && f( t[q], s[q], g, y ) > f( t[q], u, g, y ) )
            {
                q--;
            }
            if ( q < 0 )
            {
                q    = 0;
                s[0] = u;
            }
            else
            {
                int64_t w = 1 + sep( (int64_t)s[q], u, g, y );
                if ( w < M )
                {
                    q++;
                    s[q] = u;
                    t[q] = w;
                }
            }
        }
        for ( int64_t u = M - 1; u >= 0; --u )
        {
            int64_t fVal = f( u, s[q], g, y );
            CVF_ASSERT( fVal <= std::numeric_limits<double>::max() );
            image[u][y] = static_cast<unsigned int>( fVal );
            if ( u == t[q] )
            {
                q = q - 1;
            }
        }
    }
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RiaImageTools::makeGrayScale( QImage& image )
{
    for ( int i = 0; i < image.height(); i++ )
    {
        uchar* scanLine = image.scanLine( i );
        for ( int j = 0; j < image.width(); j++ )
        {
            QRgb* pixel = reinterpret_cast<QRgb*>( scanLine + j * 4 );
            int   gray  = qGray( *pixel );
            int   alpha = qAlpha( *pixel );
            *pixel      = QColor( gray, gray, gray, alpha ).rgba();
        }
    }
}