//
//  Lynkeos
//  $Id: MyListProcessing.m,v 1.18 2005/02/01 22:59:48 j-etienne Exp $
//
//  Created by Jean-Etienne LAMIAUD on Fri Nov 07 2003.
//  Copyright (c) 2003-2005. Jean-Etienne LAMIAUD
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// 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.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
//

#include <stdlib.h>
#include <assert.h>

#include "MyListProcessing.h"

#include "fourier.h"
#include "corelation.h"
#include "stack.h"

//==============================================================================
// Generic processing functions
//==============================================================================
static void getImageSample( MyImageListItem* item, RGB *dark, RGB *flat,
                            MyIntegerRect rect, 
                            REAL *spectrum, RGB *pixels, short expand )
{
   NSImage* srcImage;
   NSRect wRect;

   // Keep the original rectangle unchanged
   wRect.origin.x = rect.origin.x;
   wRect.origin.y = rect.origin.y;
   wRect.size.width = rect.size.width;
   wRect.size.height = rect.size.height;

   srcImage = [item getImage];
   if ( srcImage != nil )
   {
      // Create an image to draw the NSImage in
      NSImageRep* srcImageRep = [srcImage bestRepresentationForDevice:nil];
      int width = [srcImageRep pixelsWide],
         height = [srcImageRep pixelsHigh];
      NSImage* offScreenImage = [[[NSImage alloc] initWithSize:NSMakeSize(width,
                                                                        height)]
                                 autorelease];
      NSBitmapImageRep* bitMap;
      u_char *plane[5];
      u_short x, y, w, h, wp, ox, oy, rowSize, pixelSize;
      bool planar;
      RGB color;

      // Draw the desired image in it
      [offScreenImage lockFocus];
      [srcImageRep drawInRect:NSMakeRect(0,0,width,height)];
      // Retrieve the part of the rectangle that is in the image
      if ( wRect.origin.x < 0 )
      {
         wRect.size.width += wRect.origin.x;
         wRect.origin.x = 0;
      }
      if ( wRect.origin.y < 0 )
      {
         wRect.size.height += wRect.origin.y;
         wRect.origin.y = 0;
      }
      if ( wRect.origin.x + wRect.size.width > width )
         wRect.size.width = width - wRect.origin.x;
      if ( wRect.origin.y + wRect.size.height > height )
         wRect.size.height = height - wRect.origin.y;
      bitMap = [[[NSBitmapImageRep alloc] initWithFocusedViewRect:wRect] 
                  autorelease];
      [offScreenImage unlockFocus];

      // Access the data
      [bitMap getBitmapDataPlanes:plane];
      rowSize = [bitMap bytesPerRow];
      pixelSize = [bitMap bitsPerPixel]/8;

      planar = [bitMap isPlanar];
      w = rect.size.width;
      h = rect.size.height;
      ox = wRect.origin.x - rect.origin.x;
      // Beware of the Y flip between bitmap and screen
      oy = rect.origin.y + rect.size.height 
           - wRect.origin.y - wRect.size.height;
      // Padded width for in place transform
      wp = (w/2 + 1)*sizeof(COMPLEX)/sizeof(REAL);

      for ( y = 0; y < h; y++ )
      {
         for ( x = 0; x < w; x++ )
         {
            short i, j;

            // Fill with black outside the image (and take care of the flip, 
            // again and again...)
            if ( rect.origin.x + x < 0 || 
                 rect.origin.x + x >= wRect.origin.x + wRect.size.width ||
                 rect.origin.y + h - 1 - y < 0 || 
                 rect.origin.y + h - 1 - y >= 
                                            wRect.origin.y + wRect.size.height )
            {
               color.red = 0;
               color.green = 0;
               color.blue = 0;
            }
            else
            {
               // Read the data in the bitmap
               if ( planar )
               {
                  color.red = plane[0][(y-oy)*rowSize+x-ox];
                  color.green = plane[1][(y-oy)*rowSize+x-ox];
                  color.blue = plane[2][(y-oy)*rowSize+x-ox];
               }
               else
               {
                  color.red = plane[0][(y-oy)*rowSize+(x-ox)*pixelSize];
                  color.green = plane[0][(y-oy)*rowSize+(x-ox)*pixelSize+1];
                  color.blue = plane[0][(y-oy)*rowSize+(x-ox)*pixelSize+2];
               }

               // Apply dark frame and flat field if any
               if ( dark != NULL )
               {
                  RGB darkColor = dark[(height-rect.origin.y-h+1 + y)*width 
                                       + rect.origin.x + x];
                  color.red -= darkColor.red;
                  color.green -= darkColor.green;
                  color.blue -= darkColor.blue;
               }

               if ( flat != NULL )
               {
                  RGB flatColor = flat[(height-rect.origin.y-h+1 + y)*width 
                                       + rect.origin.x + x];
                  // No check for too small value
                  // The user will soon see if the flat field is bad
                  color.red /= flatColor.red;
                  color.green /= flatColor.green;
                  color.blue /= flatColor.blue;
               }
            }

            // Fill in the spectrum
            if ( spectrum != NULL )
               spectrum[y*wp+x] = (color.red + color.green + color.blue) / 3.0;

            // Expand the cropped part
            if ( pixels != NULL )
            {
               for ( j = 0; j < expand; j++ )
                  for ( i = 0; i < expand; i++ )
                     pixels[(y*expand + j)*w*expand + x*expand + i] = color;
            }
         }
      }

   }
}

// Cut the highest frequencies from the spectrum to suppress noise
static void cutoffSpectrum( SPECTRUM spectrum, u_short width, u_short height, 
                            u_short cutoff )
{
   u_short x, y;
   u_short h_2 = height/2;
   u_long cut2 = cutoff*cutoff;

   // Save time if there is no cutoff at all
   if ( cutoff >= sqrt(width*width+height*height) )
      return;

   for ( y = 0; y < height; y++ )
   {
      for ( x = 0; x < width/2 + 1; x++ )
      {
         short dx = x, dy = y;
         u_long f2; 
         if ( dy >= h_2 )
            dy -= height;
         f2 = dx*dx + dy*dy;

         if ( f2 > cut2 )
            spectrum[y*(width/2+1)+x] = 0.0;
      }
   }
}

static double quality( SPECTRUM spectrum, u_short width, u_short height,
                       u_short down, u_short up )
{
    u_short x, y;
    double q = 0.0;
    u_long d2 = down*down, u2 = up*up;
    u_long n = 0;
    double lum = (__real__ spectrum[0])/(double)width/(double)height;

    for( y = 0; y < height; y++ )
    {
        for ( x = 0; x < width/2 +1; x++ )
        {
            short dx = x, dy = y;
            long f2;
            if ( dy >= height/2 )
                dy -= height;
            f2 = dx*dx + dy*dy;
            if ( f2 > d2 && f2 < u2 )
            {
                COMPLEX s = spectrum[y*(width/2+1)+x];
                q += sqrt( __real__ s * __real__ s + __imag__ s * __imag__ s );
                n++;
            }
        }
    }

    return( q/lum/(double)n );
}

static double entropy( REAL *image, u_short width, u_short height )
{
   // Padded width for in place transform
   u_short pw = (width/2+1)*sizeof(COMPLEX)/sizeof(REAL);
   double e = 0, bmax = 0;
   u_long x, y;

   // Compute the quadratic pixel sum
   for( y = 0; y < height; y++ )
      for( x = 0; x < width; x++ )
         bmax += image[y*pw+x] * image[y*pw+x];

   bmax = sqrt(bmax);

   // Compute the entropy
   for( y = 0; y < height; y++ )
   {
      for( x = 0; x < width; x++ )
      {
         double b = image[y*pw+x]/bmax;
         if ( b > 0.0 )
            e -= b * log(b);
      }
   }

   return( e );
}

//==============================================================================
// Private common part of list processing
//==============================================================================
@interface MyListProcessing(Private)
- (id) initWithDelegate :(id)delegate ;
- (void) setEnumerator:(NSData*)list 
             darkFrame:(NSData*)dark flatField:(NSData*)flat ;
- (void) processList ;
- (void) processNextItem ;
- (void) processItem :(MyImageListItem*)item ; // To be overriden in subclasses
@end

@implementation MyListProcessing(Private)

- (id) initWithDelegate :(id)delegate
{
   [self init];

   _result = nil;
   _delegate = delegate;
   _processEnded = NO;
   _list = nil;
   _cropRectangle = MyMakeIntegerRect(0,0,0,0);

   [_delegate processDidCreate:self];

   return( self );
}

- (void) setEnumerator:(NSData*)list 
             darkFrame:(NSData*)dark flatField:(NSData*)flat
{
   // No need to retain, it is not autoreleased in this thread
   [list getBytes:&_list];

   [dark getBytes:&_darkFrame];
   [flat getBytes:&_flatField];
}

- (void) processList
{
   // Create a run loop for this thread
   NSRunLoop* runLoop = [NSRunLoop currentRunLoop];

   while ( ! _processEnded )
   {
      // Process the run loop to handle inter-threads messaging
      if ( _list == nil )
         // Infinite timeout when there is no list to process yet
         [runLoop runMode: NSDefaultRunLoopMode 
                  beforeDate:[NSDate distantFuture]];
      else
      {
         // Null timeout and process next item immediately after
         [runLoop runMode :NSDefaultRunLoopMode beforeDate:[NSDate date]];
         if ( ! _processEnded  )
            [self processNextItem];
      }
   }

   [_delegate processDidFinish:self data:_result];
}

- (void) processNextItem
{
   NSAutoreleasePool* pool = [[NSAutoreleasePool alloc] init];
   MyImageListItem *item = [_list nextObject];

   if ( item == nil )
      // Process is finished
      [self stopProcessing];
   else
      [self processItem:item];

   [pool release];
}

- (void) processItem :(MyImageListItem*)item
{
   NSAssert( NO, @"MyListProcessing doesn't respond to processItem" );
}

@end

//==============================================================================
// Root class for list processing
//==============================================================================
@implementation MyListProcessing

- (void) dealloc
{
   [_result release];
   [super dealloc];
}

+ (void) threadWithAttributes :(NSDictionary*)attr 
{
   NSAutoreleasePool *pool;
   NSPort *rxPort, *txPort;
   NSConnection *cnx;
   MyListProcessing *process;
   id delegate;

   pool = [[NSAutoreleasePool alloc] init];

   rxPort = [attr objectForKey:K_RX_PORT_KEY];
   txPort = [attr objectForKey:K_TX_PORT_KEY];
   NSAssert( rxPort != nil && txPort != nil, 
             @"Missing connection for MyListProcessing thread" );

   cnx = [NSConnection connectionWithReceivePort:rxPort sendPort:txPort];

   delegate = [cnx rootProxy];
   process = [[self alloc] initWithDelegate:delegate];
   // The main thread has sent retain to "process" proxy, but it seems to retain
   // our "process" also, while when the main threads sends "release" to the 
   // proxy, it will not release our "process", therefore we correct the retain 
   // count :
   [process release];

   [process processList];

   // As the "release" in the main thread is sent only to the proxy "process"
   // This is the real release for this thread
   [process release];

   [pool release];
}

- (void) stopProcessing
{
   _processEnded = YES;
}

@end

//==============================================================================
// Image aligner derived class
//==============================================================================
@implementation MyImageAligner

- (id) init
{
   if ( (self = [super init]) != nil )
   {
      _referenceItem = nil;
      FFT_DATA_INIT( &_bufferSpectrum );
      _referenceSpectrum = nil;
      _referenceOrigin = MyMakeIntegerPoint(0,0);
      _side = 0;
      _cutoff = 0;
      _precisionThreshold = 0;
      _refSpectrumLock = nil;
      _refSpectrumAvailable = NO;
   }

   return( self );
}

- (void) dealloc
{
   free_spectrum( &_bufferSpectrum );

   [super dealloc];
}

- (oneway void) alignWithList :(NSData*)list reference:(NSData*)refItem 
                     refBuffer:(NSData*)refSpectrum
                          lock:(NSData*)lock
                    darkFrame:(NSData*)dark flatField:(NSData*)flat
                     rectangle:(MyIntegerRect)rect 
                        cutOff:(double)cutoff 
            precisionThreshold:(double)threshold
{
   NSPoint offset;
   MyIntegerRect r;

   [self setEnumerator:list darkFrame:dark flatField:flat];
   _cropRectangle = rect;

   [lock getBytes:&_refSpectrumLock];
   
   [refItem getBytes:&_referenceItem];
   NSAssert( _referenceItem != nil, @"No reference item to align from" );
   [_referenceItem retain];

   _side = _cropRectangle.size.width ;

   // Get the other attributes
   _cutoff = _side*cutoff;
   _precisionThreshold = _side*threshold;
   [refSpectrum getBytes :&_referenceSpectrum];

   // Extract the data
   r = _cropRectangle;
   if ( [_referenceItem hasSearchSquare] )
      r.origin = [_referenceItem searchSquareOrigin];
   _referenceOrigin = r.origin;

   // Prepare the reference spectrum in only one thread
   if ( [_refSpectrumLock tryLock] )
   {
      getImageSample( _referenceItem, _darkFrame, _flatField, r, 
                      (REAL*)_referenceSpectrum->spectrum, nil, 1 );

      // And transform it into a spectrum
      fourier( *_referenceSpectrum );

      // Cut the highest frequencies
      cutoffSpectrum( _referenceSpectrum->spectrum, _side, _side, _cutoff );

      // Set the reference item to 0,0 offset
      offset.x = 0.0;
      offset.y = 0.0;
      [_delegate processDidProgress:[NSData dataWithBytes:&_referenceItem 
                                            length:sizeof(MyImageListItem*)]
                                            data:[NSData dataWithBytes:&offset 
                                            length:sizeof(NSPoint)]];
      _refSpectrumAvailable = YES;
      [_refSpectrumLock unlock];
   }

   // Allocate the buffer for each other images
   allocate_spectrum( &_bufferSpectrum, _side, _side, 1, 
                      FOR_DIRECT|FOR_INVERSE );
}

- (void) processItem :(MyImageListItem*)item
{
   if ( item != _referenceItem && [item getSelectionState] == NSOnState )
   {
      MyIntegerRect r = _cropRectangle;
      NSPoint offset;
      CORRELATION_PEAK peak;

      // Get the spectrum of that other image
      if ( [item hasSearchSquare] )
         r.origin = [item searchSquareOrigin];
      getImageSample( item, _darkFrame, _flatField, r, 
                      (REAL*)_bufferSpectrum.spectrum, nil, 1 );
      fourier( _bufferSpectrum );
      cutoffSpectrum( _bufferSpectrum.spectrum, _side, _side, _cutoff );

      // Check the reference spectrum availability before corelating against it
      if ( !_refSpectrumAvailable )
      {
         // Rendez vous with the "reference" thread
         [_refSpectrumLock lock];
         _refSpectrumAvailable = YES;
         [_refSpectrumLock unlock];
      }

      // correlate it against the reference
      correlate_spectrums( *_referenceSpectrum, _bufferSpectrum, 
                           _bufferSpectrum );
      corelation_peak( _bufferSpectrum, &peak );

      if ( peak.sigma_x < _precisionThreshold 
           && peak.sigma_y < _precisionThreshold )
      {
         // Beware, there is a y-flip between the bitmap and the screen
         offset.x = peak.x - r.origin.x + _referenceOrigin.x;
         offset.y = -peak.y - r.origin.y + _referenceOrigin.y;

         // Inform the delegate
         [_delegate processDidProgress:[NSData dataWithBytes:&item 
                                               length:sizeof(MyImageListItem*)]
                                  data:[NSData dataWithBytes:&offset 
                                               length:sizeof(NSPoint)]];
      }
      else{
         // Inform the delegate with no alignment
         [_delegate processDidProgress:[NSData dataWithBytes:&item 
                                               length:sizeof(MyImageListItem*)]
                                  data:nil];
      }
   }
}

@end

//==============================================================================
// Image analyzer derived class
//==============================================================================
@implementation MyImageAnalyzer

- (id) init
{
   [super init];

   _side = 0;
   _lowerCutoff = 0;
   _upperCutoff = 0;
   FFT_DATA_INIT( &_bufferSpectrum );

   return( self );
}

- (void) dealloc
{
   free_spectrum( &_bufferSpectrum );
   [super dealloc];
}

- (oneway void) analyzeWithList :(NSData*)list 
                       darkFrame:(NSData*)dark flatField:(NSData*)flat
                       rectangle:(MyIntegerRect)rect
                          method:(MyAnalysisMethod)method
                       lowCutoff:(double)lCutoff highCutoff:(double)hCutoff
{
   [self setEnumerator:list darkFrame:dark flatField:flat];
   _cropRectangle = rect;

   _side = _cropRectangle.size.width ;

   _method = method;
   _lowerCutoff = _side*lCutoff;
   _upperCutoff= _side*hCutoff;

   // Allocate the buffer for each image
   allocate_spectrum( &_bufferSpectrum, _side, _side, 1,
                      _method == SpectrumAnalysis ? FOR_DIRECT : 0 );
}

- (void) processItem :(MyImageListItem*)item
{
   if ( [item getSelectionState] == NSOnState )
   {
      MyIntegerRect r = _cropRectangle;
      double q;

      // Get the spectrum of that image
      if ( [item hasSearchSquare] )
         r.origin = [item searchSquareOrigin];
      getImageSample( item, _darkFrame, _flatField, r, 
                      (REAL*)_bufferSpectrum.spectrum, NULL, 1 );
      if ( _method == SpectrumAnalysis )
          fourier( _bufferSpectrum );

      // Analyze its quality
      switch ( _method )
      {
         case SpectrumAnalysis:
          q = quality( _bufferSpectrum.spectrum, _side, _side, 
                       _lowerCutoff, _upperCutoff );
            break;
         case EntropyAnalysis:
            // Maximum entropy of N pixels is sqrt(N)*log(sqrt(N))
            q = (_side*log(_side) / 
                 entropy((REAL*)_bufferSpectrum.spectrum, _side, _side) 
                  -  1.0) * 10.0;
            break;
         default:
            NSAssert(NO, @"Invalid analysis method");
      }

      if ( isnan( q ) )
         printf( "NaN quality !\n" );
      // Inform the delegate
      [_delegate processDidProgress:[NSData dataWithBytes:&item 
                                            length:sizeof(MyImageListItem*)]
                               data:[NSData dataWithBytes:&q 
                                            length:sizeof(double)]];
   }
}

@end

//==============================================================================
// Image stacker derived class
//==============================================================================
@implementation MyImageStack

- (id) init
{
   [super init];

   _rgbSum = NULL;
   _rgbBuffer = NULL;
   _factor = 0;

   return( self );
}

- (void) dealloc
{
   if ( _rgbBuffer != NULL )
      free( _rgbBuffer );

   [super dealloc];
}

- (oneway void) stackWithList:(NSData*)list 
                    darkFrame:(NSData*)dark flatField:(NSData*)flat
                    rectangle:(MyIntegerRect)rect
                   sizeFactor:(u_short)factor
{
   const RGB Black = {0.0,0.0,0.0};
   u_short x, y;

   [self setEnumerator:list darkFrame:dark flatField:flat];
   _cropRectangle = rect;
   _factor = factor;

   // Prepare an empty image to stack in
   _rgbSum = (RGB*)malloc( _cropRectangle.size.width*_factor
                           * _cropRectangle.size.height*_factor
                           * sizeof(RGB) );
   for ( y = 0; y < _cropRectangle.size.height*_factor; y++ )
      for ( x = 0; x < _cropRectangle.size.width*_factor; x++ )
         _rgbSum[y*_cropRectangle.size.width*_factor+x] = Black;

   _result = [[NSData dataWithBytes:&_rgbSum length:sizeof(RGB*)] retain];

   _rgbBuffer = (RGB*)malloc( _cropRectangle.size.width*_factor
                           * _cropRectangle.size.height*_factor
                           * sizeof(RGB) );
}

- (void) processItem :(MyImageListItem*)item
{
   if ( [item getSelectionState] == NSOnState && [item isAligned] )
   {
      MyIntegerRect r = _cropRectangle;
      MyIntegerPoint shift;
      NSPoint p = [item alignOffset];

      // Get the image part to add
      p.x *= -1;	// Shift the crop rectangle in the opposite side
      p.y *= -1;
      shift.x = (p.x < 0 ? (int)(p.x-1) : (int)p.x);   // Crop at integer pixels
      shift.y = (p.y < 0 ? (int)(p.y-1) : (int)p.y);
      r.origin.x += shift.x;
      r.origin.y += shift.y;
      getImageSample( item, _darkFrame, _flatField, r, nil, _rgbBuffer, 
                      _factor );

      // Accumulate (Warning, there is a Y flip between screen and bitmap)
      stack_layer( _rgbSum, _rgbBuffer, (p.x - shift.x)*(double)_factor, 
                   (1.0 - p.y + shift.y)*(double)_factor,
                   r.size.width*_factor, r.size.height*_factor );

      // Inform the delegate
      [_delegate processDidProgress:[NSData dataWithBytes:&item 
                                            length:sizeof(MyImageListItem*)]
                               data:nil];
   }
}

@end

void normalize_rgb( RGB *rgb, u_long length, double scale, BOOL mono )
{
   double s;
   u_long i;

   // Calculate max value if needed
   if ( scale == 0.0 )
   {
      double max = 0;

      // The scale shall be 1/max
      for( i = 0; i < length; i++ )
      {
         RGB v = rgb[i];

         assert( v.red >= 0 && v.green >= 0 && v.blue >= 0 );

         // Make the image monochromatic if needed
         if ( mono )
         {
            double m = (v.red + v.green + v.blue)/3.0;

            v.red = m;
            v.green = m;
            v.blue = m;
            rgb[i] = v;
         }

         if ( v.red > max )
            max = v.red;
         if ( v.green > max )
            max = v.green;
         if ( v.blue > max )
            max = v.blue;
      }
      s = 1.0/max;
   }
   else
      s = scale;

   // Apply the scale
   for( i = 0; i < length; i++ )
   {
      rgb[i].red *= s;
      rgb[i].green *= s;
      rgb[i].blue *= s;
   }
}