//
//  Lynkeos
//  $Id: MyPostProcessing.m,v 1.8 2005/01/27 23:12:42 j-etienne Exp $
//
//  Created by Jean-Etienne LAMIAUD on Wed Dec 10 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.
//

#import "MyPostProcessing.h"

#include "fourier.h"

#ifdef GNUSTEP
#import <AppKit/AppKit.h>
#else
#endif

#define K_CUTOFF 2	/* The Gaussian is 1/K_CUTOFF at argument radius */

static void prepareImageSpectrum( FFT_DATA s, RGB* src, REAL *minv, REAL *maxv)
{
   u_long x, y;
   REAL min = HUGE, max = -HUGE;

   // Copy the pixels
   for( y = 0; y < s.h; y++ )
   {
      for ( x = 0; x < s.w; x++ )
      {
         RGB c = src[y*s.w+x];

         *redValue(s,x,y) = c.red;
         *greenValue(s,x,y) = c.green;
         *blueValue(s,x,y) = c.blue;

         // Update the range
         if ( c.red < min )
            min = c.red;
         if ( c.green < min )
            min = c.green;
         if ( c.blue < min )
            min = c.blue;
         if ( c.red > max )
            max = c.red;
         if ( c.green > max )
            max = c.green;
         if ( c.blue > max )
            max = c.blue;
      }
   }

   if ( minv != NULL )
      *minv = min;
   if ( maxv != NULL )
      *maxv = max;
}

inline static u_short scale( REAL v, REAL vmin, REAL factor )
{
   REAL vs = (v - vmin)*factor;
   u_short vres;
   
   if ( vs < 0.0 )
      vres = 0;
   else if ( vs > 65535.0 )
      vres = 65535;
   else
      vres = vs;
   
   return( vres );
}

static void makeGaussianSpectrum( FFT_DATA buffer, double r )
{
   const double k = log(K_CUTOFF)/r/r;
   double weight;
   u_short x, y;
   u_long i;

   // Fill the buffer with the gaussian
   for( y = 0; y <= (buffer.h+1)/2; y++ )
   {
      for( x = 0; x <= (buffer.w+1)/2; x++ )
      {
	 double d2, v;

	 d2 = x*x + y*y;
	 v = exp( -k*d2 );

	 *colorValue(buffer,x,y,0) = v;

	 if ( x != 0 )
	    *colorValue(buffer,buffer.w-x,y,0) = v;

	 if ( y != 0 )
	    *colorValue(buffer,x,buffer.h-y,0) = v;

	 if ( x != 0 && y != 0 )
	    *colorValue(buffer,buffer.w-x,buffer.h-y,0) = v;
      }
   }

   // Transform in a spectrum
   fourier(buffer);

   // Normalize the result
   // The (0,0) sample of the Gaussian spectrum is the gaussian integral 
   // on the domain
   weight = __real__ buffer.spectrum[0];
   for( i = 0; i < (buffer.w/2+1)*buffer.h; i++ )
      buffer.spectrum[i] /= weight;
}

static void processSpectrums( SPECTRUM result, SPECTRUM src,
                              u_short width, u_short height,
                              SPECTRUM dGauss, double threshold,
                              SPECTRUM uGauss, double gain )
{
   COMPLEX  dTerm, uTerm;
   REAL module, s2 = threshold*threshold;
   
   u_long n;
   
   if ( threshold == 1.0 && gain == 0.0 )
      return; /* Nothing to do */
   
   for( n = 0; n < (width/2+1)*height; n++ )
   {
      COMPLEX i = src[n];
      
      /* Deconvolution term = source/gauss when gauss > threshold, 
         source/threshold otherwise */
      /* For the memory : 1/(a+ib) = (a-ib)/(a2+b2) */
      if ( s2 < 1.0 )
      {
         COMPLEX dg = dGauss[n];

         module = (__real__ dg * __real__ dg) + (__imag__ dg * __imag__ dg);
	    
         if ( module > s2 )
            dTerm = i/dg;

         else
            dTerm = i/threshold;
      }
      else
         dTerm = i;
	 
      /* Unsharp term = 1 + gain*(1 - gauss) */
      COMPLEX ug = uGauss[n];

      uTerm = 1 + gain - gain*ug;

      /* Complex product of the two terms */	 
      result[n] = uTerm*dTerm;
   }
}

@implementation MyPostProcessing
- (id) init
{
   _source = NULL;
   FFT_DATA_INIT(&_originalSpectrum);
   _spectrum_sequence = 0;
   FFT_DATA_INIT(&_result);
   FFT_DATA_INIT(&_unsharpGauss);
   _unsharpRadius = -1.0;
   FFT_DATA_INIT(&_deconvGauss);
   _deconvRadius = -1.0;
   
   _width = 0;
   _height = 0;
   _minValue = 0.0;
   _maxValue = 0.0;
   
   return( self );
}

- (void) dealloc
{
   free_spectrum( &_originalSpectrum );
   free_spectrum( &_unsharpGauss );
   free_spectrum( &_deconvGauss );
   free_spectrum( &_result );
}

- (REAL) minValue { return( _minValue ); }
- (REAL) maxValue { return( _maxValue ); }

- (void) process :(RGB*)src seqnb:(long)sequence
            width:(u_short)width height:(u_short)height
     deconvRadius:(double)dRadius deconvThreshold:(double)threshold
    unsharpRadius:(double)uRadius unsharpGain:(double)gain
{
   // Invalidate the result if there is no source
   if ( src == NULL )
   {
      free_spectrum( &_originalSpectrum );
      free_spectrum( &_deconvGauss );
      free_spectrum( &_unsharpGauss );
      free_spectrum( &_result );
      return;
   }
   
   // Shortcut for "no processing"
   else if ( threshold == 1.0 && gain == 0.0 )
   {
      // Force complete processing at next call
      free_spectrum( &_originalSpectrum );
      free_spectrum( &_deconvGauss );
      free_spectrum( &_unsharpGauss );
      
      // Allocate the "no processing" buffer
      _width = width;
      _height = height;
      allocate_spectrum( &_result, width, height, 3, 0 );
      
      // And just copy the source in it
      prepareImageSpectrum( _result, src, &_minValue, &_maxValue );
   }
   
   // Standard processing
   else
   {
      // Reconstruct the source spectrum, only when the source changes
      if ( sequence != _spectrum_sequence 
           || _originalSpectrum.spectrum == NULL )
      {
	 // Allocate FFT buffers 
	 allocate_spectrum( &_originalSpectrum, width, height, 3, FOR_DIRECT  );
	 allocate_spectrum( &_result, width, height, 3, FOR_INVERSE  );

	 // Compute the image spectrum
	 prepareImageSpectrum( _originalSpectrum, src, NULL, NULL );
         fourier(_originalSpectrum);

	 _source = src;
      }
      
      // Construct the deconvolution Gaussian (if it needs to be changed)
      if ( width != _width || height != _height 
           || _deconvGauss.spectrum == NULL )
      {
         allocate_spectrum( &_deconvGauss, width, height, 1, FOR_DIRECT );
	 _deconvRadius = -1.0;
      }
      if ( _deconvRadius != dRadius )
      {
	 _deconvRadius = dRadius;
	 makeGaussianSpectrum( _deconvGauss, dRadius );
      }
      
      // Construct the unsharp Gaussian
      if ( width != _width || height != _height 
           || _unsharpGauss.spectrum == NULL )
      {
         allocate_spectrum( &_unsharpGauss, width, height, 1, FOR_DIRECT );
	 _unsharpRadius = -1.0;
      }
      if ( _unsharpRadius != uRadius )
      {
	 _unsharpRadius = uRadius;
	 makeGaussianSpectrum( _unsharpGauss, uRadius );
      }
      
      // Perform the deconvolution + unsharp processing
      processSpectrums( rgbPlane(_result,RED_PLANE), 
                        rgbPlane(_originalSpectrum,RED_PLANE), 
                        width, height, _deconvGauss.spectrum, threshold, 
                        _unsharpGauss.spectrum, gain );
      processSpectrums( rgbPlane(_result,GREEN_PLANE), 
                        rgbPlane(_originalSpectrum,GREEN_PLANE), 
                        width, height, _deconvGauss.spectrum, threshold, 
                        _unsharpGauss.spectrum, gain );
      processSpectrums( rgbPlane(_result,BLUE_PLANE), 
                        rgbPlane(_originalSpectrum,BLUE_PLANE), 
                        width, height, _deconvGauss.spectrum, threshold, 
                        _unsharpGauss.spectrum, gain );

      // Return to spatial domain
      fourier_inverse( _result, &_minValue, &_maxValue );
   }
   
   _width = width;
   _height = height;
   _spectrum_sequence = sequence;
}

- (NSBitmapImageRep*) makeImageWithMin :(REAL)minV Max:(REAL)maxV
{
   NSBitmapImageRep* bitmap;
   u_short *pixels;
   u_short x, y, c;
   REAL f;
   
   // Return no image if no process done
   if ( _result.spectrum == NULL )
      return( nil );
   
   // Create a bitmap with it
   bitmap = [[[NSBitmapImageRep alloc] initWithBitmapDataPlanes:nil
						     pixelsWide:_width
						     pixelsHigh:_height
						  bitsPerSample:16
						samplesPerPixel:3
						       hasAlpha:NO
						       isPlanar:NO
                                       colorSpaceName:NSCalibratedRGBColorSpace
						    bytesPerRow:0
						   bitsPerPixel:0]
                                       autorelease];
   
   // Copy the pixels
   pixels = (u_short*)[bitmap bitmapData];
   f = 65536.9/(maxV - minV);
   
   for( y = 0; y < _height; y++ )
      for( x = 0; x < _width; x++ )
         for( c = 0; c < 3; c++ )
            pixels[y*_width*3+x*3+c] = scale( *colorValue(_result,x,y,c), 
                                              minV, f );
   
   return( bitmap );
}

@end
