/*AUTHOR:  Ravi Prakash
           University of Southern California, Los Angeles
*/

#include<cv.h>
#include<highgui.h>
#include<cxcore.h>
#include<stdio.h>

#define DOHSV 1 //0->Cluster RGB, 1->Cluster HSV
#define NOOFCLUSTERS 4
#define CHANNEL 2 //If DOHSV=0 then 0->B, 1->G, 2->R Defines which channel to use for clustering
                  //If DOHSV=1 then 0->H, 1->S, 2->V

float angDiff(float a, float b)
{
     int temp=a-b;
     if(temp<0) temp*=-1;
     if(temp>90) return 180-temp;
     else return temp;     
}

//This code is from Open CV, slightly modified by me.
void cvKMeans2forH( const CvArr* samples_arr, int cluster_count,
           CvArr* labels_arr, CvTermCriteria termcrit )
{
    CvMat* centers = 0;
    CvMat* old_centers = 0;
    CvMat* counters = 0;

    CV_FUNCNAME( "cvKMeans2" );

    __BEGIN__;

    CvMat samples_stub, labels_stub;
    CvMat* samples = (CvMat*)samples_arr;
    CvMat* labels = (CvMat*)labels_arr;
    CvMat* temp = 0;
    CvRNG rng = CvRNG(-1);
    int i, j, k, sample_count, dims;
    int ids_delta, iter;
    double max_dist;

    if( !CV_IS_MAT( samples ))
        CV_CALL( samples = cvGetMat( samples, &samples_stub ));

    if( !CV_IS_MAT( labels ))
        CV_CALL( labels = cvGetMat( labels, &labels_stub ));

    if( cluster_count < 1 )
        CV_ERROR( CV_StsOutOfRange, "Number of clusters should be positive" );

    if( CV_MAT_DEPTH(samples->type) != CV_32F || CV_MAT_TYPE(labels->type) != CV_32SC1 )
        CV_ERROR( CV_StsUnsupportedFormat,
        "samples should be floating-point matrix, cluster_idx - integer vector" );

    if( labels->rows != 1 && (labels->cols != 1 || !CV_IS_MAT_CONT(labels->type)) ||
        labels->rows + labels->cols - 1 != samples->rows )
        CV_ERROR( CV_StsUnmatchedSizes,
        "cluster_idx should be 1D vector of the same number of elements as samples' number of rows" );

    CV_CALL( termcrit = cvCheckTermCriteria( termcrit, 1e-6, 100 ));

    termcrit.epsilon *= termcrit.epsilon;
    sample_count = samples->rows;

    if( cluster_count > sample_count )
        cluster_count = sample_count;

    dims = samples->cols*CV_MAT_CN(samples->type);
    ids_delta = labels->step ? labels->step/(int)sizeof(int) : 1;

    CV_CALL( centers = cvCreateMat( cluster_count, dims, CV_64FC1 ));
    CV_CALL( old_centers = cvCreateMat( cluster_count, dims, CV_64FC1 ));
    CV_CALL( counters = cvCreateMat( 1, cluster_count, CV_32SC1 ));

    // init centers
    for( i = 0; i < sample_count; i++ )
        labels->data.i[i] = cvRandInt(&rng) % cluster_count;

    counters->cols = cluster_count; // cut down counters
    max_dist = termcrit.epsilon*2;

    for( iter = 0; iter < termcrit.max_iter; iter++ )
    {
        // computer centers
        cvZero( centers );
        cvZero( counters );

        for( i = 0; i < sample_count; i++ )
        {
            float* s = (float*)(samples->data.ptr + i*samples->step);
            k = labels->data.i[i*ids_delta];
            double* c = (double*)(centers->data.ptr + k*centers->step);
            for( j = 0; j <= dims - 4; j += 4 )
            {                
                double t0 = angDiff(c[j],s[j]);
                double t1 = angDiff(c[j+1],s[j+1]);

                c[j] = t0;
                c[j+1] = t1;

                t0 = angDiff(c[j+2],s[j+2]);
                t1 = angDiff(c[j+3],s[j+3]);

                c[j+2] = t0;
                c[j+3] = t1;
            }
            for( ; j < dims; j++ )
                c[j] = (c[j]+s[j]);
            counters->data.i[k]++;
        }

        if( iter > 0 )
            max_dist = 0;

        for( k = 0; k < cluster_count; k++ )
        {
            double* c = (double*)(centers->data.ptr + k*centers->step);
            if( counters->data.i[k] != 0 )
            {
                double scale = 1./counters->data.i[k];
                for( j = 0; j < dims; j++ )
                    c[j] *= scale;
            }
            else
            {
                i = cvRandInt( &rng ) % sample_count;
                float* s = (float*)(samples->data.ptr + i*samples->step);
                for( j = 0; j < dims; j++ )
                    c[j] = s[j];
            }
            
            if( iter > 0 )
            {
                double dist = 0;
                double* c_o = (double*)(old_centers->data.ptr + k*old_centers->step);
                for( j = 0; j < dims; j++ )
                {
                    double t = c[j] - c_o[j];
                    dist += t*t;
                }
                if( max_dist < dist )
                    max_dist = dist;
            }
        }

        // assign labels
        for( i = 0; i < sample_count; i++ )
        {
            float* s = (float*)(samples->data.ptr + i*samples->step);
            int k_best = 0;
            double min_dist = DBL_MAX;

            for( k = 0; k < cluster_count; k++ )
            {
                double* c = (double*)(centers->data.ptr + k*centers->step);
                double dist = 0;
                
                j = 0;
                for( ; j <= dims - 4; j += 4 )
                {
                    double t0 = c[j] - s[j];
                    double t1 = c[j+1] - s[j+1];
                    dist = dist+t0*t0 + t1*t1;
                    t0 = c[j+2] - s[j+2];
                    t1 = c[j+3] - s[j+3];
                    dist = ((int)(dist+t0*t0 + t1*t1))%90;
                }

                for( ; j < dims; j++ )
                {
                    double t = c[j] - s[j];
                    dist = ((int)(dist+t*t))%90;
                }

                if( min_dist > dist )
                {
                    min_dist = dist;
                    k_best = k;
                }
            }

            labels->data.i[i*ids_delta] = k_best;
        }

        if( max_dist < termcrit.epsilon )
            break;

        CV_SWAP( centers, old_centers, temp );
    }

    cvZero( counters );
    for( i = 0; i < sample_count; i++ )
        counters->data.i[labels->data.i[i]]++;

    // ensure that we do not have empty clusters
    for( k = 0; k < cluster_count; k++ )
        if( counters->data.i[k] == 0 )
            for(;;)
            {
                i = cvRandInt(&rng) % sample_count;
                j = labels->data.i[i];
                if( counters->data.i[j] > 1 )
                {
                    labels->data.i[i] = k;
                    counters->data.i[j]--;
                    counters->data.i[k]++;
                    break;
                }
            }

    __END__;

    cvReleaseMat( &centers );
    cvReleaseMat( &old_centers );
    cvReleaseMat( &counters );
}


int main()
{
     CvRNG rng = cvRNG(0xffffffff);
     
     IplImage *im=cvLoadImage("football.jpg", 1);
     IplImage *toShow=cvCloneImage(im);
     cvNamedWindow("Image Window", 1);     
     cvShowImage("Image Window", im);
     if(DOHSV==1) cvCvtColor(im, im, CV_BGR2HSV);
     printf("\n\nInput image depth = %d", im->depth);


   
     CvMat *arr    =cvCreateMat(im->height*im->width, 1, CV_32FC1 );
     CvMat *retVect=cvCreateMat(im->height*im->width, 1, CV_32SC1 );
     
     printf("\n\nCalculating random color pallete");
     int colorArr[3][NOOFCLUSTERS];
     for(int k=0; k<NOOFCLUSTERS; k++)
            {
                 colorArr[0][k]=cvRandInt(&rng)%255;
                 colorArr[1][k]=cvRandInt(&rng)%255;
                 colorArr[2][k]=cvRandInt(&rng)%255;
            }
     printf("\nCalculated random color pallete");
    
    printf("\n\nConverting image data to cvMat");
    for(int i=0; i<im->height; i++)
            for(int j=0; j<im->width; j++)
                    {
                         ((float*)(arr->data.ptr + arr->step* (i*im->width+j)))[0]= ((uchar*)(im->imageData + im->widthStep*i))[j*3+CHANNEL];
                    }
     printf("\nConverted image data to cvMat");

     printf("\n\nDoing K-means clustering");
     if(DOHSV==1 && CHANNEL==0) cvKMeans2forH(arr, NOOFCLUSTERS, retVect, cvTermCriteria( CV_TERMCRIT_ITER, 20, 1.0 ) );
     else cvKMeans2(arr, NOOFCLUSTERS, retVect, cvTermCriteria( CV_TERMCRIT_ITER, 20, 1.0 ) );
     printf("\nDone K-means clustering");
  
     
     printf("\n\nForming output image from returned index array");
     for(int i=0; i<im->height; i++)
            for(int j=0; j<im->width; j++)
                    {
                         int index=(retVect->data.ptr + retVect->step* (i*im->width+j))[0];
                         ((uchar*)(toShow->imageData + toShow->widthStep*i))[j*3  ]= colorArr[0] [index];
                         ((uchar*)(toShow->imageData + toShow->widthStep*i))[j*3+1]= colorArr[1] [index];
                         ((uchar*)(toShow->imageData + toShow->widthStep*i))[j*3+2]= colorArr[2] [index];
                    }
     printf("\nFormed output image from returned index array");
     
     
     cvNamedWindow("Clustered", 1);     
     cvShowImage("Clustered", toShow);
     cvSaveImage("Output.jpg", toShow);     

     cvWaitKey(0);

     cvReleaseMat(&arr);
     cvReleaseMat(&retVect);
          
     cvDestroyWindow("Image Window");
     cvDestroyWindow("Clustered");     
     
     cvReleaseImage(&im);
     cvReleaseImage(&toShow);
     
     return 1;
}