//
//  Lynkeos
//  $Id$
//
//  Created by Jean-Etienne LAMIAUD on Tue Feb 28 2023.
//  Copyright (c) 2023. 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 <limits.h>

#include "FITSRawReader.h"

static int colorSelector(void *opaque, u_short x, u_short y)
{
   const u_short *bayer = (u_short*)opaque;
   return bayer[(x % 2) + (y % 2)*2];
}

@interface FITSRawReader(Private)
@end

@implementation FITSRawReader(Private)
@end

@implementation FITSRawReader

+ (void) load
{
   // Nothing to do, this is just to force the runtime to load this class
}

+ (void) lynkeosFileTypes:(NSArray**)fileTypes
{
   // Open with priority on the standard FITS reader
   NSNumber *pri = [NSNumber numberWithInt:2];
   *fileTypes = [NSArray arrayWithObjects: pri, @"fits",
                                           pri, @"fts",
                                           pri, @"fit",
                                           nil];
}

+ (BOOL) hasCustomImageBuffer { return( YES ); }

- (id) init
{
   self = [super init];
   if ( self != nil )
   {
      _fits = NULL;
      _width = 0.0;
      _height = 0.0;
      _imageScale = 1.0;
      _imageZero = 0.0;
      _minValue = HUGE_VAL;
      _maxValue = -HUGE_VAL;
      _imageScale = 1.0;
      _bottomUpRows = YES;
      for (int i = 0; i < 2; i++)
         for (int j = 0; j < 2; j++)
            _bayerArray[i][j] = 0;
   }
   return( self );
}

- (id) initWithURL:(NSURL*)url
{
   int err = 0;
   int dimension, nbits;
   long size[2];
   
   self = [self init];
   
   if ( self != nil )
   {
      const char *file;
      char buffer[80];
      BOOL isBayer = NO;

      // Unfortunately, CFITSIO does not handle correctly the
      // file://localhost/... URL given by Cocoa
      if ( [url isFileURL] )
         file = [[url path] fileSystemRepresentation];
      else
         file = [[url absoluteString] UTF8String];
      
      fits_open_file( &_fits, file, READONLY, &err );
      if (err == 0)
      {
         fits_get_img_param( _fits, 2, &nbits, &dimension, size, &err );
         if (err == 0)
         {
            _width = size[0];
            _height = size[1];
            
            // Get the row ordering, if stated
            if ( fits_read_key( _fits, TSTRING, "ROWORDER", buffer, NULL, &err) == 0 )
            {
               if (strcmp(buffer,"TOP-DOWN")==0)
                  _bottomUpRows = NO;
            }
            err = 0;
            
            if (fits_read_key ( _fits, TSTRING, "BAYERPAT", buffer, NULL, &err) == 0)
            {
               isBayer = YES;
               for (int i = 0; i < 2; i++)
               {
                  for (int j = 0; j < 2; j++)
                  {
                     int row = j;
                     if (_bottomUpRows)
                        row = (_height - j - 1) % 2;
                     switch(buffer[row*2+i])
                     {
                        case 'R':
                           _bayerArray[i][j] = 0;
                           break;
                        case 'G':
                           _bayerArray[i][j] = 1;
                           break;
                        case 'B':
                           _bayerArray[i][j] = 2;
                           break;
                        default:
                           NSLog(@"Inconsistent color name %c in RAW FITS", buffer[j*2+i]);
                           isBayer = NO;
                     }
                  }
               }
            }
            err = 0;

            if ( dimension == 2 && isBayer )
            {
               // Save the image scale and zero for sample read
               if ( fits_read_key(_fits, TDOUBLE, "BSCALE", &_imageScale, NULL, &err) != 0 )
                  _imageScale = 1.0;
               err = 0.0;
               if (fits_read_key(_fits, TDOUBLE, "BZERO", &_imageZero, NULL, &err) != 0 )
                  _imageZero = 0.0;
               err = 0.0;
               
               if ( fits_read_key( _fits, TDOUBLE, "DATAMIN", &_minValue, NULL, &err)
                   != 0 )
                  _minValue = HUGE_VAL;
               err = 0.0;
               if ( fits_read_key( _fits, TDOUBLE, "DATAMAX", &_maxValue, NULL, &err)
                   != 0 )
                  _maxValue = -HUGE_VAL;
               err = 0.0;
               
               // Determine the scale and zero to use when converting to a NSImage
               if ( _minValue >= _maxValue )
               {
                  // No information, we need to read all the data
                  double *buf = (double*)malloc( sizeof(double)*_width );
                  u_short x, y;
                  
                  fits_set_bscale( _fits, _imageScale, _imageZero, &err );
                  
                  for( y = 1; y <= _height && err == 0; y++ )
                  {
                     long first[2] = {1,y};
                     fits_read_pix( _fits, TDOUBLE, first, _width,
                                   NULL, buf, NULL, &err );
                     for( x = 0; x < _width; x++ )
                     {
                        if ( buf[x] < _minValue )
                           _minValue = buf[x];
                        if ( buf[x] > _maxValue )
                           _maxValue = buf[x];
                     }
                  }
                  
                  // Discard an impossible to understand error
                  if ( err != 0 )
                     fits_report_error( stderr, err );
                  err = 0;
                  free( buf );
               }
            }
         }
      }

      if ( err != 0 || dimension != 2 || !isBayer )
      {
         NSLog(@"Unable to open RAW FITS image : %@", [url absoluteString] );
         fits_report_error( stderr, err );
         [self release];
         self = nil;
      }
   }

   return( self );
}

- (void) dealloc
{
   int err = 0;
   if ( _fits != NULL )
      fits_close_file( _fits, &err );
   
   NSAssert( err == 0, @"FITS closing error" );
   [super dealloc];
}

- (void) imageWidth:(u_short*)w height:(u_short*)h
{
   *w = _width;
   *h = _height;
}

- (u_short) numberOfPlanes
{
   return( 3 );
}

- (void) getMinLevel:(double*)vmin maxLevel:(double*)vmax
{
   if ( _minValue < _maxValue && _imageScale != 0.0 )
   {
      *vmin = _minValue;
      *vmax = _maxValue;
   }
   else
   {
      *vmin = 0.0;
      *vmax = 255.0;
   }
}

- (NSImage*) getNSImage
{
   NSPoint nullOffsets[] = {NSZeroPoint, NSZeroPoint, NSZeroPoint};
   double vmin[] = {_minValue, _minValue, _minValue, _minValue};
   double vmax[] = {_maxValue, _maxValue, _maxValue, _maxValue};
   double g[] = {1.0, 1.0, 1.0, 2.5};
   LynkeosImageBuffer *limg =  [self getCustomImageSampleAtX:0 Y:0 W:_width H:_height
                                               withTransform:[[NSAffineTransform transform] transformStruct]
                                                 withOffsets:nullOffsets];
      
#if !GNUSTEP
   return [[[NSImage alloc] initWithCGImage:[limg getImageInRect:LynkeosMakeIntegerRect(0, 0, _width, _height)
                                                       withBlack:vmin white:vmax gamma:g]
                                       size:NSZeroSize]
               autorelease];
#else
   return [[[NSImage alloc] initWithData:[[limg getNSImageWithBlack:vmin
                                                              white:vmax
                                                              gamma:g]
                                           TIFFRepresentation]] autorelease];
#endif
}

- (void) setMode:(ListMode_t)mode { _mode = mode; }

- (void) setDarkFrame:(LynkeosImageBuffer*)dark
{
   if (_darkFrame != nil)
   {
      [_darkFrame release];
      _darkFrame = nil;
   }
   NSAssert(dark == nil || _mode == ImageMode || _mode == UnsetListMode,
            @"Bad reader mode to set dark frame %d", _mode);
   NSAssert(dark == nil || [dark isKindOfClass:[BayerImageBuffer class]],
            @"RAW TUFF dark frame is not in Bayer format");
   _darkFrame = (BayerImageBuffer*)[dark copy];
   
   if (_darkFrame != nil)
   {
      // Weight will be substracted during calibration, therefore, set all weight to zero, except for dead pixels
      // Start by getting mean and standard deviation of pixels in the bayer matrix
      u_short c, x, y;
      double s = 0.0, s2 = 0.0, n = 0.0;
      for (y = 0; y < _darkFrame->_h; y++)
      {
         for (x = 0; x < _darkFrame->_w; x++)
         {
            for (c = 0; c < _darkFrame->_nPlanes; c++)
            {
               double w = stdColorValue(_darkFrame->_weight, x, y, c);
               double v = w*stdColorValue(_darkFrame, x, y, c);
               s += v;
               s2 += v*v;
               n += w;
            }
         }
      }
      double mean = s/n;
      double sigma = sqrt(s2/n - mean*mean);
      for (y = 0; y < _darkFrame->_h; y++)
      {
         for (x = 0; x < _darkFrame->_w; x++)
         {
            for (c = 0; c < _darkFrame->_nPlanes; c++)
            {
               if (stdColorValue(_darkFrame->_weight, x, y, c) <= 0.0
                   || (stdColorValue(_darkFrame, x, y, c) - mean) < 3.0*sigma)
                  // Correct pixel, weight shall not change in calibrated image
                  stdColorValue(_darkFrame->_weight, x, y, c) = 0.0;
               // Otherwise, it is a dead pixel, keep the weight, in order to null it in the  calibrated image
               //               else
               //                  NSLog(@"Dead pixel at %d,%d in plane %d, value %f weight %f",
               //                        x, y, c, stdColorValue(_darkFrame, x, y, c),
               //                        stdColorValue(_darkFrame->_weight, x, y, c));
            }
         }
      }
   }
}

- (void) setFlatField:(LynkeosImageBuffer*)flat
{
   if (_flatField != nil)
   {
      [_flatField release];
      _flatField = nil;
   }
   NSAssert(_flatField == nil || _mode == ImageMode || _mode == UnsetListMode,
            @"Bad reader mode to set flat field %d", _mode);
   _flatField = [flat retain];
}

- (BOOL) canBeCalibratedBy:(id <LynkeosFileReader>)reader asMode:(ListMode_t)mode
{
   u_short w, h;
   [reader imageWidth:&w height:&h];
   return w == _width && h == _height
          && (mode == FlatFieldMode
              || (mode == DarkFrameMode && [reader isKindOfClass:[self class]]));
}

- (void) getImageSample:(REAL * const * const)sample
             withPlanes:(u_short)nPlanes
                    atX:(short)x Y:(short)y W:(u_short)w H:(u_short)h
              lineWidth:(u_short)lineW
{
   NSAssert( nPlanes == 1, @"Try to read multiplane FITS" );
   const NSAffineTransformStruct ident = {1.0, 0.0, 0.0, 1.0, 0.0, 0.0};
   const NSPoint still[3] = {{0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}};
   LynkeosImageBuffer* customImage = [self getCustomImageSampleAtX:x Y:y W:w H:h
                                                     withTransform:ident withOffsets:still];
   [customImage convertToPlanar:sample withPlanes:nPlanes lineWidth:lineW];
}

- (LynkeosImageBuffer*) getCustomImageSampleAtX:(short)x Y:(short)y
                                              W:(u_short)w H:(u_short)h
                                  withTransform:(NSAffineTransformStruct)transform
                                    withOffsets:(const NSPoint*)offsets
{
   // Read the data
   REAL *imageData = (REAL*)malloc(_width*_height*sizeof(REAL));
   LynkeosImageBuffer* image = nil;
   int err = 0;

   fits_set_bscale( _fits, _imageScale, _imageZero, &err );

   if ( err == 0 )
   {
#warning Optimize by reading only the data in the square (the offsets must be communicated to the color selector)
      for( u_short yl = 0; yl < _height; yl++ )
      {
         long first[2] = {1,yl+1};
         void *buf;
         u_short ly = (_bottomUpRows ? _height-yl-1 : yl);

         buf = &(imageData[ly*_width]);

         fits_read_pix(_fits,
                       PROCESSING_PRECISION == DOUBLE_PRECISION ? TDOUBLE : TFLOAT,
                       first, _width, NULL, buf, NULL, &err );
      }
   }

   if ( err != 0 )
      fits_report_error( stderr, err );

   image = [[[BayerImageBuffer alloc] initWithData:imageData
                                            opaque:_bayerArray getPlane:colorSelector
                                             width:_width lineW:_width height:_height
                                           xoffset:0 yoffset:0
                                               atX:x Y:y W:w H:h
                                     withTransform:transform withColorAlign:offsets
                                          withDark: _darkFrame withFlat:_flatField] autorelease];
   free(imageData);

   return image;
}


- (NSDictionary*) getMetaData
{
   return( nil );
}

@end