//
//  PREqualize.m
//  PRICE
//  Image level equalization
//
//  Created by Riccardo Mottola on Fri Dec 05 2003.
//  Copyright (c) 2003-2006 Carduus. All rights reserved.
//
// This program is free software; you can redistribute it and/or modify it under the terms of the version 2 of the GNU General Public License as published by the Free Software Foundation.
// This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.

#import <math.h>
#import <limits.h>

#import "PREqualize.h"


@implementation PREqualize

- (PRImage *)equalizeImage:(PRImage *)srcImage :(int)space
{
    NSBitmapImageRep *srcImageRep;
    PRImage          *destImage;
    NSBitmapImageRep *destImageRep;
    int              w, h;
    int              x, y;
    int              i;
    unsigned char    *srcData;
    unsigned char    *destData;
    int              bytesPerPixel;
    int              pixNum;
    BOOL             isColor;
    
    /* some trace */
    printf("levels: %d\n", UCHAR_MAX);
    printf("space: %d\n", space);
    
    /* get source image representation and associated information */
    srcImageRep = [srcImage tiffRep];
    
    w = [srcImageRep pixelsWide];
    h = [srcImageRep pixelsHigh];
    pixNum = h * w;
    printf("pixels: %d\n", pixNum);
    bytesPerPixel = [srcImageRep bitsPerPixel] /8;
    
    /* check bith depth and color/greyscale image */
    if ([srcImageRep hasAlpha])
    {
        if ([srcImageRep samplesPerPixel] == 2)
        {
            printf("Grayscale image\n");
            isColor = NO;
        }
        else
        {
            printf("Color image\n");
            isColor = YES;
        }
    }
    else
    {
        if ([srcImageRep samplesPerPixel] == 1)
        {
            printf("Grayscale image\n");
            isColor = NO;
        }
        else
        {
            printf("Color image\n");
            isColor = YES;
        }
    }
    
    /* allocate destination image and its representation */
    destImage = [[PRImage alloc] initWithSize:NSMakeSize(w, h)];
    if (isColor)
    {
        destImageRep = [[NSBitmapImageRep alloc]
                    initWithBitmapDataPlanes:NULL
                    pixelsWide:w
                    pixelsHigh:h
                    bitsPerSample:8
                    samplesPerPixel:3
                    hasAlpha:NO
                    isPlanar:NO
                    colorSpaceName:NSCalibratedRGBColorSpace
                    bytesPerRow:0
                    bitsPerPixel:0];
    } else
    {
        destImageRep = [[NSBitmapImageRep alloc]
                    initWithBitmapDataPlanes:NULL
                    pixelsWide:w
                    pixelsHigh:h
                    bitsPerSample:8
                    samplesPerPixel:1
                    hasAlpha:NO
                    isPlanar:NO
                    colorSpaceName:NSCalibratedWhiteColorSpace
                    bytesPerRow:0
                    bitsPerPixel:0];
    }
    srcData = [srcImageRep bitmapData];
    destData = [destImageRep bitmapData];
    
    if (isColor)
    {
        if (space == COLOR_SPACE_RGB)
        {
            unsigned long int histogramDenormR[UCHAR_MAX+1]; /* not normalized pixel count for each level */
            unsigned long int histogramDenormG[UCHAR_MAX+1]; /* not normalized pixel count for each level */
            unsigned long int histogramDenormB[UCHAR_MAX+1]; /* not normalized pixel count for each level */
            float histogramR[UCHAR_MAX+1];                   /* normalized histogram */
            float histogramG[UCHAR_MAX+1];                   /* normalized histogram */
            float histogramB[UCHAR_MAX+1];                   /* normalized histogram */
            float cumulativeHistogramR[UCHAR_MAX+1];         /* cumulative histogram */
            float cumulativeHistogramG[UCHAR_MAX+1];         /* cumulative histogram */
            float cumulativeHistogramB[UCHAR_MAX+1];         /* cumulative histogram */
            
            /* calculate the histogram */
            for (i = 0; i <= UCHAR_MAX; i++)
                histogramDenormR[i] = histogramDenormG[i] = histogramDenormB[i] =  0;
            for (y = 0; y < h; y++)
                for (x = 0; x < w*3; x += 3)
                {
                    histogramDenormR[srcData[y*(w*3) + x]]++;
                    histogramDenormG[srcData[y*(w*3) + x + 1]]++;
                    histogramDenormB[srcData[y*(w*3) + x + 2]]++;
                }
        
            /* normalize histogram */
            for (i = 0; i <= UCHAR_MAX; i++)
            {
                histogramR[i] = (float)histogramDenormR[i] / (float)pixNum;
                histogramG[i] = (float)histogramDenormG[i] / (float)pixNum;
                histogramB[i] = (float)histogramDenormB[i] / (float)pixNum;
            }
            
            /* cumulative histogram */
            cumulativeHistogramR[0] = histogramR[0];
            cumulativeHistogramG[0] = histogramG[0];
            cumulativeHistogramB[0] = histogramB[0];
            for (i = 1; i <= UCHAR_MAX; i++)
            {
                cumulativeHistogramR[i] = cumulativeHistogramR[i-1] + histogramR[i];
                cumulativeHistogramG[i] = cumulativeHistogramG[i-1] + histogramG[i];
                cumulativeHistogramB[i] = cumulativeHistogramB[i-1] + histogramB[i];
            }
            
            /* equalize */
            for (y = 0; y < h; y++)
                for (x = 0; x < w*3; x += 3)
                {
                    destData[y*(w*3) + x] = floor((UCHAR_MAX+0.9)*cumulativeHistogramR[srcData[y*(w*3) + x]]);
                    destData[y*(w*3) + x + 1] = floor((UCHAR_MAX+0.9)*cumulativeHistogramG[srcData[y*(w*3) + x + 1]]);
                    destData[y*(w*3) + x + 2] = floor((UCHAR_MAX+0.9)*cumulativeHistogramB[srcData[y*(w*3) + x + 2]]);
                }
        } else if (space == COLOR_SPACE_YUV)
        {
            unsigned long int histogramDenormY[UCHAR_MAX+1]; /* not normalized pixel count for each level */
            float histogramY[UCHAR_MAX+1];                   /* normalized histogram */
            float cumulativeHistogramY[UCHAR_MAX+1];         /* cumulative histogram */
            register int r, g, b;
            unsigned char yy, cb, cr;
            
/*
 *            JPEG-YCbCr (601) from "digital 8-bit R'G'B'  "
 *            ========================================================================
 *            Y' =       + 0.299    * R'd + 0.587    * G'd + 0.114    * B'd
 *            Cb = 128   - 0.168736 * R'd - 0.331264 * G'd + 0.5      * B'd
 *            Cr = 128   + 0.5      * R'd - 0.418688 * G'd - 0.081312 * B'd
 *            ........................................................................
 *            R'd, G'd, B'd   in {0, 1, 2, ..., 255}
 *            Y', Cb, Cr      in {0, 1, 2, ..., 255}
 * R = Y                    + 1.402   (Cr-128)
 * G = Y - 0.34414 (Cb-128) - 0.71414 (Cr-128)
 * B = Y + 1.772   (Cb-128)
 */
            /* first we convert the whole image to YCC-d */
            for (y = 0; y < h; y++)
                for (x = 0; x < w*3; x += 3)
                {
                    r = srcData[y*(w*3) + x];
                    g = srcData[y*(w*3) + x + 1];
                    b = srcData[y*(w*3) + x + 2];
                    yy = rintf(       + 0.2990f*r + 0.5870f*g + 0.1140f*b);
                    cb = rintf(128.0f - 0.1687f*r - 0.3313f*g + 0.5000f*b);
                    cr = rintf(128.0f + 0.5000f*r - 0.4187f*g - 0.0813f*b);

                    destData[y*(w*3) + x]     = yy;
                    destData[y*(w*3) + x + 1] = cb;
                    destData[y*(w*3) + x + 2] = cr;
                }

            /* calculate the histogram */
            for (i = 0; i <= UCHAR_MAX; i++)
                histogramDenormY[i] =  0;
            for (y = 0; y < h; y++)
                for (x = 0; x < w*3; x += 3)
                    histogramDenormY[destData[y*(w*3) + x]]++;

            /* normalize histogram */
            for (i = 0; i <= UCHAR_MAX; i++)
                histogramY[i] = (float)histogramDenormY[i] / (float)pixNum;

            /* cumulative histogram */
            cumulativeHistogramY[0] = histogramY[0];
            for (i = 1; i <= UCHAR_MAX; i++)
                cumulativeHistogramY[i] = cumulativeHistogramY[i-1] + histogramY[i];


            /* equalize */
            for (y = 0; y < h; y++)
                for (x = 0; x < w*3; x += 3)
                {
                    destData[y*(w*3) + x] = floor((UCHAR_MAX+0.9)*cumulativeHistogramY[destData[y*(w*3) + x]]);
                }
                    
            /* now we convert back */
            /* first we convert the whole image to YCC-d */
            for (y = 0; y < h; y++)
                for (x = 0; x < w*3; x += 3)
                {
                    yy = destData[y*(w*3) + x];
                    cb = destData[y*(w*3) + x + 1];
                    cr = destData[y*(w*3) + x + 2];
                    r = yy                     + (int)rintf(1.40200f*(cr-128));
                    g = yy - (int)rintf(0.34414f*(cb-128) + 0.71414f*(cr-128));
                    b = yy + (int)rintf(1.77200f*(cb-128));
                    r = r > UCHAR_MAX ? UCHAR_MAX : r;
                    g = g > UCHAR_MAX ? UCHAR_MAX : g;
                    b = b > UCHAR_MAX ? UCHAR_MAX : b;
                    r = r < 0 ? 0 : r;
                    g = g < 0 ? 0 : g;
                    b = b < 0 ? 0 : b;
                    destData[y*(w*3) + x]     = r;
                    destData[y*(w*3) + x + 1] = g;
                    destData[y*(w*3) + x + 2] = b;
                } 
        }
    } else
    {
        unsigned long int histogramDenorm[UCHAR_MAX+1]; /* not normalized pixel count for each level */
        float histogram[UCHAR_MAX+1];                   /* normalized histogram */
        float cumulativeHistogram[UCHAR_MAX+1];         /* cumulative histogram */
        /* calculate the histogram */
        for (i = 0; i <= UCHAR_MAX; i++)
            histogramDenorm[i] = 0;
        for (y = 0; y < h; y++)
            for (x = 0; x < w; x++)
                histogramDenorm[srcData[y*w + x]]++;
    
        /* normalize histogram */
        for (i = 0; i <= UCHAR_MAX; i++)
            histogram[i] = (float)histogramDenorm[i] / (float)pixNum;
        
        /* cumulative histogram */
        cumulativeHistogram[0] = histogram[0];
        for (i = 1; i <= UCHAR_MAX; i++)
            cumulativeHistogram[i] = cumulativeHistogram[i-1] + histogram[i];
        
        /* equalize */
        for (y = 0; y < h; y++)
            for (x = 0; x < w; x++)
                destData[y*w + x] = floor((UCHAR_MAX+0.9)*cumulativeHistogram[srcData[y*w + x]]);
    }
    
    [destImage addRepresentation:destImageRep];
    [destImageRep release];
    [destImage autorelease];
    return destImage;
}


@end
