//
//  NSImageToRimage.m
//  AUCoderDemo
//
//  Created by Andrew Salamon on 5/31/06.
//  Copyright (c) 2006 Machine Perception Laboratory 
//  University of California San Diego.//

#import "NSImageToRimage.h"

#ifdef USE_IMAGEMAGICK
#import <Magick++.h>
#endif

/// divide image values by this to normalize (usually 255.0 or 1.0)
static const float norm = 255.0;

@implementation NSImage(CERTMethods)
+ (NSImage *)imageFromRImage:(RImage<float> *)rimage
{
	return [self imageFromRImage:rimage withNorm:norm];
}

+ (NSImage *)imageFromRImage:(RImage<float> *)rimage withNorm:(float)_norm
{
	NSImage *newImage = NULL;
	NSBitmapImageRep *rep = [[NSBitmapImageRep alloc] initWithBitmapDataPlanes:NULL
																	pixelsWide:rimage->width
																	pixelsHigh:rimage->height
																 bitsPerSample:8
															   samplesPerPixel:3
																	  hasAlpha:NO
																	  isPlanar:NO
																colorSpaceName:NSCalibratedRGBColorSpace
																  bitmapFormat:(NSBitmapFormat)0 
																   bytesPerRow:0
																  bitsPerPixel:0]; 
	if( rep )
	{
		if( NO )
		{
			for( int row=0; row < [rep pixelsHigh]; ++row )
			{
				//				unsigned int i = row * [imageRep bytesPerRow];
				for( int col=0; col < [rep pixelsWide]; ++col )
				{
					NSUInteger pvalues[3];
					unsigned char rgb = (unsigned char)(rimage->getPixel( col, row ) * _norm );
					pvalues[0] = rgb;
					pvalues[1] = rgb;
					pvalues[2] = rgb;
					
					[rep setPixel:pvalues atX:col y:row];
				}
			}
		}
		else
		{
			unsigned char *data = [rep bitmapData];
			//int pixelCnt = rimage->width * rimage->height;
			for( int row=0; row < [rep pixelsHigh]; ++row )
			{
				unsigned char *rowData = data + (row *[rep bytesPerRow]);
				for( int col=0; col < [rep pixelsWide]; ++col )
				{
					unsigned char rgb = (unsigned char)(rimage->getPixel( col, row ) * _norm );
					*rowData = rgb; ++rowData;
					*rowData = rgb; ++rowData;
					*rowData = rgb; ++rowData;
				}
			}
//			for( int i=0, repi=0; i < pixelCnt; ++i )
//			{
//				unsigned char rgb = (unsigned char)(rimage->array[i] * _norm);
//				data[repi++] = rgb;
//				data[repi++] = rgb;
//				data[repi++] = rgb;
//			}
		}
		newImage = [[NSImage alloc] initWithSize:NSMakeSize(rimage->width,rimage->height)];
		[newImage addRepresentation:rep];
		[rep release];
	}
	
	return [newImage autorelease];
}

#define ALPHASHIFT    0x18
#define ALPHAMASK     0xff000000
#define NOTALPHAMASK  0x00ffffff

#define REDSHIFT      0x10
#define REDMASK       0x00ff0000
#define NOTREDMASK    0xff00ffff

#define GREENSHIFT    0x08
#define GREENMASK     0x0000ff00
#define NOTGREENMASK  0xffff00ff

#define BLUESHIFT     0x00
#define BLUEMASK      0x000000ff
#define NOTBLUEMASK   0xffffff00

const RGBPIXEL_ aMask = ALPHAMASK;
const RGBPIXEL_ rMask = REDMASK;
const RGBPIXEL_ gMask = GREENMASK;
const RGBPIXEL_ bMask = BLUEMASK;

+ (NSImage *)imageFromColorRImage:(RImage<RGBPIXEL_> *)rimage;
{
	NSImage *newImage = NULL;
	NSBitmapImageRep *rep = [[NSBitmapImageRep alloc] initWithBitmapDataPlanes:NULL
																	pixelsWide:rimage->width
																	pixelsHigh:rimage->height
																 bitsPerSample:8
															   samplesPerPixel:3
																	  hasAlpha:NO
																	  isPlanar:NO
																colorSpaceName:NSCalibratedRGBColorSpace
																  bitmapFormat:(NSBitmapFormat)0 
																   bytesPerRow:0
																  bitsPerPixel:0]; 
	if( rep )
	{
		if( NO )
		{
			for( int row=0; row < [rep pixelsHigh]; ++row )
			{
				//				unsigned int i = row * [imageRep bytesPerRow];
				for( int col=0; col < [rep pixelsWide]; ++col )
				{
					NSUInteger pvalues[3];
					RGBPIXEL_ rgb = (RGBPIXEL_)(rimage->getPixel( col, row ) );
					
					//				unsigned int alpha = rgb & aMask;
					pvalues[0] = (rgb & rMask) >> REDSHIFT;
					pvalues[1] = (rgb & gMask) >> GREENSHIFT;
					pvalues[2] = rgb & bMask;
					
					[rep setPixel:pvalues atX:col y:row];
				}
			}
		}
		else
		{
			unsigned char *data = [rep bitmapData];
			//int pixelCnt = rimage->width * rimage->height;
			for( int row=0; row < [rep pixelsHigh]; ++row )
			{
				unsigned char *rowData = data + (row *[rep bytesPerRow]);
				for( int col=0; col < [rep pixelsWide]; ++col )
				{
					RGBPIXEL_ rgb = (RGBPIXEL_)(rimage->getPixel( col, row ) );
					*rowData = (rgb & rMask) >> REDSHIFT;
					++rowData;
					*rowData = (rgb & gMask) >> GREENSHIFT;
					++rowData;
					*rowData = rgb & bMask;
					++rowData;
				}
			}			
		}
		
		newImage = [[NSImage alloc] initWithSize:NSMakeSize(rimage->width,rimage->height)];
		[newImage addRepresentation:rep];
		[rep release];
	}
	
	return [newImage autorelease];
}

- (NSBitmapImageRep *)bitmapImageRep
{
	NSArray *reps = [self representations];
	NSEnumerator *repEnum = [reps objectEnumerator];
	NSBitmapImageRep *imageRep = NULL;
	while( (imageRep = [repEnum  nextObject]) && ![imageRep isKindOfClass:[NSBitmapImageRep class]] )
	{
		
	}
	if( !imageRep )
	{
		// If it contains an NSPICTImageRep then we have to use this hack to get the bitmap rep (slow)
		NSData	*tiffData = [self TIFFRepresentation];
		imageRep = [NSBitmapImageRep imageRepWithData:tiffData];
	}		
	return imageRep;
}

- (RImage<float> *)rimageRep
{
	return rimageFromImage( self );
}
//{
//	NSBitmapImageRep *imageRep = [self bitmapImageRep];
//	if( !imageRep )
//	{
//		// If it contains an NSPICTImageRep then we have to use this hack to get the bitmap rep (slow)
//		NSData	*tiffData = [self TIFFRepresentation];
//		imageRep = [NSBitmapImageRep imageRepWithData:tiffData];
//		
//		if( !imageRep )
//			return NULL;
//	}
//	
//	NSSize size = [imageRep size];
//	// From reading the documentation it looks like maybe we should be using pixelsWide/High all the time, instead of the size method.
//	if( (size.width != [imageRep pixelsWide])  || (size.height != [imageRep pixelsHigh]) )
//	{
//#ifdef NSIMAGE_DEBUG
//		NSLog( @"Mismatch between [imageRep size] and [imageRep pixelsWide/pixelsHigh]: size %@; width/height %d/%d",
//			   NSStringFromSize(size), [imageRep pixelsWide], [imageRep pixelsHigh] );
//#endif
//		size.width = [imageRep pixelsWide];
//		size.height = [imageRep pixelsHigh];
//	}
//	
//	unsigned char *data = [imageRep bitmapData];
//	unsigned int dataInd = 0;
//	BOOL canUseRaw = NO;
//	BOOL skipRawAlpha = [imageRep hasAlpha];
//	if( [imageRep numberOfPlanes] != 1 )
//	{
//		NSLog( @"Planar image type, can't handle it, yet" );
//		return NULL;
//	}
//	
//	RImage<float> *rimage = new RImage<float>( static_cast<int>(ceil(size.width)), static_cast<int>(ceil(size.height)) );
//	
//	NSAssert( fabs(size.width - rimage->width) < 1.0, @"Bad width in RImage" );
//	NSAssert( fabs(size.height - rimage->height) < 1.0, @"Bad height in RImage" );
//	
//	unsigned int byteCount = [imageRep pixelsWide] * [imageRep pixelsHigh] * [imageRep samplesPerPixel]; // 
//	unsigned int pixel = 0;
//	BOOL alphaFirst = [imageRep hasAlpha] && ([imageRep bitmapFormat] & NSAlphaFirstBitmapFormat);
//	unsigned int inc = [imageRep samplesPerPixel]; // [imageRep hasAlpha]?4:3;
//	
//	if( [imageRep samplesPerPixel] == 3 )
//	{
//		if( ([imageRep pixelsWide] * [imageRep samplesPerPixel]) == [imageRep bytesPerRow] )
//		{
//			canUseRaw = YES;
//			skipRawAlpha = NO;
//		}
//		else if( ([imageRep pixelsWide] * ([imageRep samplesPerPixel]+1)) == [imageRep bytesPerRow] )
//		{
//			canUseRaw = YES;
//			skipRawAlpha = YES;
//		}
//	}
//		
//	//NSLog( @"Row byte padding: %d", [imageRep bytesPerRow] - ([imageRep pixelsWide] * [imageRep samplesPerPixel]) );
//#ifdef DEBUG
//	if( !canUseRaw )
//	{
//		NSLog( @"Unable to use raw image data: %d x %d x %d (%d). %d, %@, %@",
//			   [imageRep pixelsWide], [imageRep pixelsHigh], [imageRep samplesPerPixel], [imageRep bitsPerPixel], [imageRep bytesPerRow],
//			   ([imageRep hasAlpha]?@"Alpha":@"No Alpha"),
//			   (([imageRep bitmapFormat] & NSAlphaFirstBitmapFormat)?@"Alpha First":@"Alpha Last") );
//	}
//#endif
//	
//	if( YES || !canUseRaw )
//	{
//		// I think this is general enough that we could use it for all images, but it hasn't been tested enough to be sure.
//		for( unsigned int row=0; row < size.height; ++row )
//		{
//			dataInd = row * [imageRep bytesPerRow];
//			for( unsigned int col=0; col < size.width; ++col )
//			{
//				int theGrayValue = 0;
//
//				if( inc == 1 )
//				{
//					theGrayValue = data[dataInd++];
//				}
//				else
//				{
//					int r = 0;
//					int g = 0;
//					int b = 0;
//					
//					if( skipRawAlpha && alphaFirst )
//						++dataInd; // skip the alpha
//					r = data[dataInd++];
//					g = data[dataInd++];
//					b = data[dataInd++];
//					if( skipRawAlpha && !alphaFirst )
//						++dataInd; // skip the alpha
//
//					theGrayValue = ((306 * (r & 0xff)) +
//									(601 * (g & 0xff)) +
//									(117 * (b & 0xff))) >> 10; 
//				}
//				//					float pixVal2 = theGrayValue / 255.0;
//				float pixVal2 = theGrayValue / norm;
//				//			float pixVal = (data[i] + data[i+1] + data[i+2]) / 3.0;
//				//			pixVal /= 255.0;
//				//rimage->setPixel( pixel++, pixVal2 );
//				NSAssert( (pixVal2 <= 1.0), @"Pixel value is greater than 1.0 but shouldn't be." );
//				rimage->setPixel( col, row, pixVal2 );
//			}
//		}
//	}
//	else if( YES || ([imageRep bytesPerRow] != ([imageRep pixelsWide] * [imageRep samplesPerPixel])) )
//	{
//		for( unsigned int row=0; row < size.height; ++row )
//		{
//			//				unsigned int i = row * [imageRep bytesPerRow];
//			for( unsigned int col=0; col < size.width; ++col )
//			{
//				int theGrayValue = 0;
//				unsigned int pixValues[inc];
//				if( !canUseRaw )
//					[imageRep getPixel:pixValues atX:col y:row];
//				if( inc == 1 )
//				{
//					theGrayValue = pixValues[0];
//				}
//				else
//				{
//					int r = 0;
//					int g = 0;
//					int b = 0;
//
//					if( canUseRaw )
//					{
//						if( skipRawAlpha )
//							++dataInd; // skip the alpha
//						r = data[dataInd++];
//						g = data[dataInd++];
//						b = data[dataInd++];
//					}
//					else
//					{
//						unsigned int tmpi = 0;
//						if( alphaFirst )
//							++tmpi;
//						r = pixValues[tmpi];
//						g = pixValues[tmpi+1];
//						b = pixValues[tmpi+2];
//						
//					}
//					
//					theGrayValue = ((306 * (r & 0xff)) +
//									(601 * (g & 0xff)) +
//									(117 * (b & 0xff))) >> 10; 
//				}
//				//					float pixVal2 = theGrayValue / 255.0;
//				float pixVal2 = theGrayValue / norm;
//				//			float pixVal = (data[i] + data[i+1] + data[i+2]) / 3.0;
//				//			pixVal /= 255.0;
//				//rimage->setPixel( pixel++, pixVal2 );
//				NSAssert( (pixVal2 <= 1.0), @"Pixel value is greater than 1.0 but shouldn't be." );
//				rimage->setPixel( col, row, pixVal2 );
//			}
//		}
//	}
//	else
//	{
//		for( unsigned int i=0; i < byteCount; i+=inc )
//		{
//			unsigned int tmpi = i;
//			if( alphaFirst )
//				++tmpi;
//			int r = data[tmpi];
//			int g = data[tmpi+1];
//			int b = data[tmpi+2];
//			int theGrayValue = ((306 * (r & 0xff)) +
//								(601 * (g & 0xff)) +
//								(117 * (b & 0xff))) >> 10; 
//			float pixVal2 = theGrayValue / norm;
//			//			float pixVal = (data[i] + data[i+1] + data[i+2]) / 3.0;
//			//			pixVal /= 255.0;
//			rimage->setPixel( pixel++, pixVal2 );
//		}
//	}
//	return rimage;
//}

- (NSImage *)imageFromRect:(NSRect)rect
{
    NSAffineTransform * xform = [NSAffineTransform transform];
	
    // translate reference frame to map rectangle 'rect' into first quadrant
    [xform translateXBy: -rect.origin.x
                    yBy: -rect.origin.y];
    
    NSSize canvas_size = [xform transformSize: rect.size];
	
    NSImage * canvas = [[NSImage alloc] initWithSize: canvas_size];
    [canvas lockFocus];
	
    [xform concat];
    
    // Get NSImageRep of image
    NSImageRep * rep = [self bestRepresentationForDevice: nil];
	
    [rep drawAtPoint: NSZeroPoint];
    
    [canvas unlockFocus];
    return [canvas autorelease];
}

#ifdef USE_IMAGEMAGICK
- (RImage<float> *)rimageFromImageMagick
{
	using namespace Magick;

	NSBitmapImageRep *imageRep = [self bitmapImageRep];
	if( !imageRep )
	{
		return NULL;
	}
	
	NSDictionary *properties =
        [NSDictionary dictionaryWithObjectsAndKeys:
			[NSNumber numberWithFloat:1.0],
			NSImageCompressionFactor, NULL];    
	
	NSData *bitmapData = [imageRep representationUsingType:NSJPEGFileType properties:properties];
	
	Blob blob( [bitmapData bytes], [bitmapData length] );
	Image image( blob );
	
	if (image.columns() == 0 && image.rows() == 0)
	{
//		NSLog( @"Image contains 0 rows and 0 columns!" );
		return NULL;
	}
	
	RImage<float> *pixels = new RImage<float>( image.columns(), image.rows() );
	
//	image.quantizeColors(256); // convert to 256 colors
	image.quantize(); // convert to grayscale
	
	image.write(0,0, image.columns(), image.rows(), "I", FloatPixel, pixels->array);
	return pixels;
}
#endif
@end

RImage<float> *rimageFromImage( NSImage *image )
{
	NSBitmapImageRep *imageRep = [image bitmapImageRep];
	if( !imageRep )
	{
		// If it contains an NSPICTImageRep then we have to use this hack to get the bitmap rep (slow)
		NSData	*tiffData = [image TIFFRepresentation];
		imageRep = [NSBitmapImageRep imageRepWithData:tiffData];
		
		if( !imageRep )
			return NULL;
	}
	
	NSSize size = [imageRep size];
	// From reading the documentation it looks like maybe we should be using pixelsWide/High all the time, instead of the size method.
	if( (size.width != [imageRep pixelsWide])  || (size.height != [imageRep pixelsHigh]) )
	{
#ifdef NSIMAGE_DEBUG
		NSLog( @"Mismatch between [imageRep size] and [imageRep pixelsWide/pixelsHigh]: size %@; width/height %d/%d",
			   NSStringFromSize(size), [imageRep pixelsWide], [imageRep pixelsHigh] );
#endif
		size.width = [imageRep pixelsWide];
		size.height = [imageRep pixelsHigh];
        [imageRep setSize:size]; // This makes sure that how the image gets displayed matches how we draw the face boxes and features.
	}
	
	unsigned char *data = [imageRep bitmapData];
	unsigned int dataInd = 0;
	BOOL canUseRaw = NO;
	BOOL skipRawAlpha = [imageRep hasAlpha];
	if( [imageRep numberOfPlanes] != 1 )
	{
		NSLog( @"Planar image type, can't handle it, yet" );
		return NULL;
	}
	
	RImage<float> *rimage = new RImage<float>( static_cast<int>(ceil(size.width)), static_cast<int>(ceil(size.height)) );
	
	BOOL alphaFirst = [imageRep hasAlpha] && ([imageRep bitmapFormat] & NSAlphaFirstBitmapFormat);
	unsigned int inc = [imageRep samplesPerPixel]; // [imageRep hasAlpha]?4:3;
	
	if( [imageRep samplesPerPixel] == 3 )
	{
		if( ([imageRep pixelsWide] * [imageRep samplesPerPixel]) == [imageRep bytesPerRow] )
		{
			canUseRaw = YES;
			skipRawAlpha = NO;
		}
		else if( ([imageRep pixelsWide] * ([imageRep samplesPerPixel]+1)) == [imageRep bytesPerRow] )
		{
			canUseRaw = YES;
			skipRawAlpha = YES;
		}
	}
	
	//NSLog( @"Row byte padding: %d", [imageRep bytesPerRow] - ([imageRep pixelsWide] * [imageRep samplesPerPixel]) );
//#ifdef DEBUG
//	if( !canUseRaw )
//	{
//		NSLog( @"Unable to use raw image data: %d x %d x %d (%d). %d, %@, %@",
//			   [imageRep pixelsWide], [imageRep pixelsHigh], [imageRep samplesPerPixel], [imageRep bitsPerPixel], [imageRep bytesPerRow],
//			   ([imageRep hasAlpha]?@"Alpha":@"No Alpha"),
//			   (([imageRep bitmapFormat] & NSAlphaFirstBitmapFormat)?@"Alpha First":@"Alpha Last") );
//	}
//#endif
	
	if( YES )
	{
		// I think this is general enough that we could use it for all images, but it hasn't been tested enough to be sure.
		for( unsigned int row=0; row < size.height; ++row )
		{
			dataInd = row * [imageRep bytesPerRow];
			for( unsigned int col=0; col < size.width; ++col )
			{
				int theGrayValue = 0;
				
				if( inc == 1 )
				{
					theGrayValue = data[dataInd++];
				}
				else
				{
					int r = 0;
					int g = 0;
					int b = 0;
					
					if( skipRawAlpha && alphaFirst )
						++dataInd; // skip the alpha
					r = data[dataInd++];
					g = data[dataInd++];
					b = data[dataInd++];
					if( skipRawAlpha && !alphaFirst )
						++dataInd; // skip the alpha
					
					theGrayValue = ((306 * (r & 0xff)) +
									(601 * (g & 0xff)) +
									(117 * (b & 0xff))) >> 10; 
					// In OpenCV the formula is: Y = (0.299)R + (0.587)G + (0.114)B
				}
				//					float pixVal2 = theGrayValue / 255.0;
				float pixVal2 = theGrayValue / norm;
				//			float pixVal = (data[i] + data[i+1] + data[i+2]) / 3.0;
				//			pixVal /= 255.0;
				//rimage->setPixel( pixel++, pixVal2 );
				rimage->setPixel( col, row, pixVal2 );
			}
		}
	}
	return rimage;
}

RImage<RGBPIXEL_> *colorRimageFromImage( NSImage *image )
{
	NSBitmapImageRep *imageRep = [image bitmapImageRep];
	if( !imageRep )
	{
		// If it contains an NSPICTImageRep then we have to use this hack to get the bitmap rep (slow)
		NSData	*tiffData = [image TIFFRepresentation];
		imageRep = [NSBitmapImageRep imageRepWithData:tiffData];
		
		if( !imageRep )
			return NULL;
	}
	
	NSSize size = [imageRep size];
	// From reading the documentation it looks like maybe we should be using pixelsWide/High all the time, instead of the size method.
	if( (size.width != [imageRep pixelsWide])  || (size.height != [imageRep pixelsHigh]) )
	{
#ifdef NSIMAGE_DEBUG
		NSLog( @"Mismatch between [imageRep size] and [imageRep pixelsWide/pixelsHigh]: size %@; width/height %d/%d",
			  NSStringFromSize(size), [imageRep pixelsWide], [imageRep pixelsHigh] );
#endif
		size.width = [imageRep pixelsWide];
		size.height = [imageRep pixelsHigh];
	}
	
	unsigned char *data = [imageRep bitmapData];
	unsigned int dataInd = 0;
	BOOL canUseRaw = NO;
	BOOL skipRawAlpha = [imageRep hasAlpha];
	if( [imageRep numberOfPlanes] != 1 )
	{
		NSLog( @"Planar image type, can't handle it, yet" );
		return NULL;
	}
	
	RImage<RGBPIXEL_> *rimage = new RImage<RGBPIXEL_>( static_cast<int>(ceil(size.width)), static_cast<int>(ceil(size.height)) );
	
	BOOL alphaFirst = [imageRep hasAlpha] && ([imageRep bitmapFormat] & NSAlphaFirstBitmapFormat);
	unsigned int inc = [imageRep samplesPerPixel]; // [imageRep hasAlpha]?4:3;
	
	if( [imageRep samplesPerPixel] == 3 )
	{
		if( ([imageRep pixelsWide] * [imageRep samplesPerPixel]) == [imageRep bytesPerRow] )
		{
			canUseRaw = YES;
			skipRawAlpha = NO;
		}
		else if( ([imageRep pixelsWide] * ([imageRep samplesPerPixel]+1)) == [imageRep bytesPerRow] )
		{
			canUseRaw = YES;
			skipRawAlpha = YES;
		}
	}
	
	if( YES )
	{
		// I think this is general enough that we could use it for all images, but it hasn't been tested enough to be sure.
		for( unsigned int row=0; row < size.height; ++row )
		{
			dataInd = row * [imageRep bytesPerRow];
			for( unsigned int col=0; col < size.width; ++col )
			{
				RGBPIXEL_ theGrayValue = 0;
				
				if( inc == 1 )
				{
					RGBPIXEL_ tmp = data[dataInd++];
					
					theGrayValue = (tmp << REDSHIFT) & rMask;
					theGrayValue = theGrayValue | ((tmp << GREENSHIFT) & gMask);
					theGrayValue = theGrayValue | (tmp & bMask);
				}
				else
				{
					RGBPIXEL_ r = 0;
					RGBPIXEL_ g = 0;
					RGBPIXEL_ b = 0;
					
					if( skipRawAlpha && alphaFirst )
						++dataInd; // skip the alpha
					r = data[dataInd++];
					g = data[dataInd++];
					b = data[dataInd++];
					if( skipRawAlpha && !alphaFirst )
						++dataInd; // skip the alpha
					
					theGrayValue = (r << REDSHIFT) & rMask;
					theGrayValue = theGrayValue | ((g << GREENSHIFT) & gMask);
					theGrayValue = theGrayValue | (b & bMask);
				}
				rimage->setPixel( col, row, theGrayValue );
			}
		}
	}
	return rimage;
}