//
//  Lynkeos
//  $Id: $
//
//  Created by Jean-Etienne LAMIAUD on Tue Feb 14 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 <LynkeosCore/LynkeosProcessing.h>
#include <LynkeosCore/LynkeosMetadata.h>
#include <LynkeosCore/LynkeosInterpolator.h>

#include <libraw/libraw.h>

#include "DcrawReaderPrefs.h"
#include "LibRaw_Reader.h"

typedef uint16_t uint16_vect_t __attribute__ ((vector_size (32)));
#define uint16_vect_length (sizeof(uint16_vect_t)/sizeof(uint16_t))

static NSMutableArray *rawFilesTypes = nil;

static int colorSelector(void *opaque, u_short x, u_short y)
{
   libraw_data_t *data = (libraw_data_t*)opaque;
   int c = libraw_COLOR(data,
                        y + data->sizes.left_margin,
                        x + data->sizes.left_margin);
   return(c == 3 ? 1 : c);
}

@implementation LibRaw_Reader

+ (void) load  // Only to force the runtime to load the class
{
}

+ (void) lynkeosFileTypes:(NSArray**)fileTypes
{
   // Read the file extensions in the RAW configuration file
   // And give priority over the Dcraw reader
   if ( rawFilesTypes == nil )
   {
      NSNumber *pri = [NSNumber numberWithInt:2];
      NSString *cfgFile;
      NSArray *cfgFileTypes;
      cfgFile = [[NSBundle bundleForClass:[self class]] pathForResource:@"dcraw_file_extensions"
                                                                 ofType:@"plist"];
      if (cfgFile != nil)
      {
         NSData *plistData;
         NSError *error;
         NSPropertyListFormat format;
         NSMutableDictionary *dict;
         plistData = [NSData dataWithContentsOfFile:cfgFile];
         dict = [NSPropertyListSerialization propertyListWithData:plistData
                                                          options:NSPropertyListMutableContainers
                                                           format:&format
                                                            error:&error];
         NSAssert( dict != nil, @"Failed to read RAW files configuration" );
         cfgFileTypes = [dict objectForKey:@"extensions"];
         NSAssert( cfgFileTypes != nil,
                  @"Failed to access to RAW files extensions" );
         
         rawFilesTypes = [[NSMutableArray arrayWithCapacity:[cfgFileTypes count]*2] retain];
         
         NSEnumerator *list;
         NSString *fileType;
         for( list = [cfgFileTypes objectEnumerator];
             (fileType = [list nextObject]) != nil ; )
         {
            [rawFilesTypes addObject:pri];
            [rawFilesTypes addObject:fileType];
         }
      }
      else
         rawFilesTypes = [[NSMutableArray alloc] init];
   }
   *fileTypes = rawFilesTypes;
}

- (id) init
{
   self = [super init];
   if ( self != nil )
   {
      NSUserDefaults *prefs = [NSUserDefaults standardUserDefaults];
      _width = 0;
      _bytesPerRow = 0;
      _height = 0;
		_xOrigin = 0;
		_yOrigin = 0;
      _xLimit = 0;
      _yLimit = 0;
      _scale = 255.0/65535.0;
      if ([prefs boolForKey:K_MANUALWB_KEY])
      {
         double rw = [prefs doubleForKey:K_RED_KEY];
         _whiteBalance[RED_PLANE] = (rw > 0.0 ? rw : 1.0);
         double gw = [prefs doubleForKey:K_GREEN1_KEY];
         _whiteBalance[GREEN_PLANE] = (gw > 0.0 ? gw : 1.0);
         double bw = [prefs doubleForKey:K_BLUE_KEY];
         _whiteBalance[BLUE_PLANE] = (bw > 0.0 ? bw : 1.0);
      }
      else
      {
         for (u_short c = 0; c < 3; c++)
            _whiteBalance[c] = 1.0;
      }
      _mode = ImageMode;
      _darkFrame = nil;
      _flatField = nil;
      _metadata = [[NSMutableDictionary alloc] init];
   }
   return( self );
}

- (id) initWithURL:(NSURL*)url
{
   libraw_data_t *libraw_data;
   int          ret;

   self = [self init];

   if ( self != nil )
   {
      // Open image file
      _url = url;
      libraw_data = libraw_init(0);
      if (libraw_data == NULL)
      {
         NSLog( @"Could not initialize libraw" );
         [self release];
         return( nil );
      }
      ret = libraw_open_file(libraw_data, [[_url path] fileSystemRepresentation]);
      if (ret == 0)
         ret = libraw_unpack(libraw_data);

      if (ret != 0 || libraw_data->rawdata.raw_image == NULL)
      {
         NSLog( @"Could not open file %@", [url absoluteString] );
         libraw_close(libraw_data);
         [self release];
         return( nil );
      }

      // Retrieve image information
      _xOrigin = libraw_data->sizes.left_margin;
      _yOrigin = libraw_data->sizes.top_margin;
      _xLimit = libraw_data->sizes.iwidth;
      _yLimit = libraw_data->sizes.iheight;
      _width = _xLimit - _xOrigin;
      _bytesPerRow = libraw_data->sizes.raw_pitch;
      _height = _yLimit - _yOrigin;
      if (libraw_data->color.maximum > 0)
         _scale = 256.0 / (double)libraw_data->color.maximum;

      // Retrieve metadata
      if (strnlen(libraw_data->other.artist, 64) != 0)
         [_metadata setObject:[NSArray arrayWithObject:
                                          [NSString stringWithCString:libraw_data->other.artist
                                                             encoding:NSUTF8StringEncoding]]
                       forKey:LynkeosMD_Authors()];
      if (strnlen(libraw_data->idata.model, 64) != 0)
         [_metadata setObject:[NSString stringWithCString:libraw_data->idata.model
                     encoding:NSUTF8StringEncoding]
                       forKey:LynkeosMD_CameraModel()];
      if (strnlen(libraw_data->lens.Lens, 128) != 0)
         [_metadata setObject:[NSString stringWithCString:libraw_data->lens.Lens
                     encoding:NSUTF8StringEncoding]
                       forKey:LynkeosMD_Telescope()];
      if (libraw_data->other.timestamp > 0)
         [_metadata setObject:[NSDate dateWithTimeIntervalSince1970:
                                  (NSTimeInterval)libraw_data->other.timestamp]
                       forKey:LynkeosMD_CaptureDate()];
      if (libraw_data->other.shutter > 0.0)
         [_metadata setObject:[NSNumber numberWithFloat:libraw_data->other.shutter]
                       forKey:LynkeosMD_ExposureTime()];
      if (libraw_data->other.aperture > 0.0)
         [_metadata setObject:[NSNumber numberWithInt:libraw_data->other.aperture]
                       forKey:LynkeosMD_Aperture()];
      if (libraw_data->other.focal_len > 0.0)
         [_metadata setObject:[NSNumber numberWithInt:libraw_data->other.focal_len]
                       forKey:LynkeosMD_FocalLength()];
      if (libraw_data->other.iso_speed > 0.0)
         [_metadata setObject:[NSNumber numberWithInt:libraw_data->other.iso_speed]
                       forKey:LynkeosMD_ISOSpeed()];

      if (libraw_data->other.parsed_gps.gpsstatus == 'A')
      {
         float l = (libraw_data->other.parsed_gps.latitude[0]
                    + (libraw_data->other.parsed_gps.latitude[1]
                       + libraw_data->other.parsed_gps.latitude[2]/60.0)/60.0) / 180.0 * M_PI;
         if (libraw_data->other.parsed_gps.latref == 'S')
            l *= -1.0;
         [_metadata setObject:[NSNumber numberWithFloat:l] forKey:LynkeosMD_Latitude()];
         l = (libraw_data->other.parsed_gps.longitude[0]
              + (libraw_data->other.parsed_gps.longitude[1]
                 + libraw_data->other.parsed_gps.longitude[2]/60.0)/60.0) / 180.0 * M_PI;
         if (libraw_data->other.parsed_gps.longref == 'W')
            l *= -1;
         [_metadata setObject:[NSNumber numberWithFloat:l] forKey:LynkeosMD_Longitude()];
         [_metadata setObject:[NSNumber numberWithFloat:libraw_data->other.parsed_gps.altitude]
                       forKey:LynkeosMD_Altitude()];
      }

      // TODO: Keep data in cache
      libraw_close(libraw_data);
   }

   return( self );
}

- (void) dealloc
{
   if (_darkFrame != nil)
      [_darkFrame release];
   if (_flatField != nil)
      [_flatField release];
   [_metadata release];

   [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
{
   *vmin = 0.0;
   *vmax = 255.0;
}

- (NSImage*) getNSImage
{
   NSImage *image = nil;
   NSBitmapImageRep* bitmap;

   // Create a RGB bitmap
   bitmap = [[[NSBitmapImageRep alloc] initWithBitmapDataPlanes:NULL
                                                   pixelsWide:_width
                                                   pixelsHigh:_height
                                                  bitsPerSample:8
                                                samplesPerPixel:3
                                                       hasAlpha:NO
                                                       isPlanar:NO
                                                 colorSpaceName: NSCalibratedRGBColorSpace
                                                    bytesPerRow:0
                                                   bitsPerPixel:0]
             autorelease];

   if ( bitmap != nil )
   {
      u_char *pixels = (u_char*)[bitmap bitmapData];
      int bpr = (int)[bitmap bytesPerRow];
      int bpp = (int)[bitmap bitsPerPixel]/8;
      libraw_data_t *libraw_data;
      int ret;
      u_short x, y, p;

      libraw_data = libraw_init(0);
      ret = libraw_open_file(libraw_data, [[_url path] fileSystemRepresentation]);
      if (ret == 0)
         ret = libraw_unpack(libraw_data);

      if (ret == 0 && libraw_data->rawdata.raw_image != NULL)
      {
         for( x = _xOrigin; x < _xLimit; x+= uint16_vect_length )
         {
            void *colPtr = (void*)libraw_data->rawdata.raw_image + x*sizeof(uint16_t);
            uint16_vect_t segment[3]; // previous line, current line, next line
            uint16_t left[3], right[3];

            // Read the current segment (previous is not needed on first line)
            if (x > _xOrigin)
               left[1] = *(uint16_t*)(colPtr - sizeof(uint16_t));
            segment[1] = *(uint16_vect_t*)colPtr;
            if (x < (_xLimit - uint16_vect_length))
               right[1] = *(uint16_t*)(colPtr + sizeof(uint16_vect_t));

            for( y = _yOrigin; y < _yLimit; y++ )
            {
               // Read next line, for interpolation pupose
               if ( y < (_yLimit - 1) )
               {
                  if (x > _xOrigin)
                     left[2] = *(uint16_t*)(colPtr + (y+1)*_bytesPerRow - sizeof(uint16_t));
                  segment[2] = *(uint16_vect_t*)(colPtr + (y+1)*_bytesPerRow);
                  if (x < (_xLimit - uint16_vect_length))
                     right[2] = *(uint16_t*)(colPtr + (y+1)*_bytesPerRow + sizeof(uint16_vect_t));
               }
               for ( u_short dx = 0; dx < uint16_vect_length; dx++)
               {
                  u_short pixColor = colorSelector(libraw_data, x+dx - _xOrigin, y - _yOrigin);
                  for ( p = 0 ; p < 3; p++ )
                  {
                     double vf = 0.0;
                     if ( pixColor == p)
                     {
                        // Lucky, we fall in the correct bayer pixel
                        vf = (double)(segment[1][dx]) * _scale * _whiteBalance[p];
                     }
                     else
                     {
                        // Unlucky, perform a simple linear interpolation
                        const u_short mxl = ((x+dx) < (_xLimit - 1) ? x + dx + 1 : _xLimit - 1);
                        const u_short myl = (y < (_yLimit - 1) ? y + 1 : _yLimit - 1);
                        const u_short sxl = ((x+dx) < 1 ? 0 : x + dx - 1), syl = (y < 1 ? 0 : y - 1);
                        u_short xl, yl;
                        double sum = 0.0, weight = 0.0;
                        for ( yl = syl; yl <= myl; yl++)
                        {
                           for ( xl = sxl; xl <= mxl; xl++)
                           {
                              if (colorSelector(libraw_data, xl - _xOrigin, yl - _yOrigin) == p)
                              {
                                 if (xl < x)
                                    vf = (double)left[yl-y+1];
                                 else if (xl >= (x + uint16_vect_length))
                                    vf = (double)right[yl-y+1];
                                 else
                                    vf = (double)(segment[yl-y+1][xl-x]);
                                 
                                 sum += vf;
                                 weight += 1.0;
                              }
                           }
                        }
                        vf = sum/weight * _scale * _whiteBalance[p];
                     }
                     
                     pixels[(y-_yOrigin)*bpr+(x+dx-_xOrigin)*bpp+p] = (vf < 256.0 ? (u_char)vf : 255);
                  }
               }
               // Prepare next line
               left[0] = left[1];
               left[1] = left[2];
               right[0] = right[1];
               right[1] = right[2];
               segment[0] = segment[1];
               segment[1] = segment[2];
            }
         }
      }

      libraw_close(libraw_data);

      image = [[[NSImage alloc] initWithSize:NSMakeSize(_width, _height)] autorelease];

      if ( image != nil )
         [image addRepresentation:bitmap];
   }

   return( image );
}

- (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 dark frame is not in Bayer format");
   _darkFrame = (BayerImageBuffer*)[dark retain];

   if (_darkFrame != nil && !_darkFrame->_isProcessedDark)
   {
      // 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));
            }
         }
      }
      _darkFrame->_isProcessedDark = YES;
   }
}

- (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
{
   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
   libraw_data_t *libraw_data;
   int ret = -1;

   libraw_data = libraw_init(0);
   if (libraw_data != NULL)
      ret = libraw_open_file(libraw_data, [[_url path] fileSystemRepresentation]);
   if (ret == 0)
      ret = libraw_unpack(libraw_data);

   LynkeosImageBuffer* image = nil;
   u_short xl, yl;
   REAL v;

   // Compute the enclosing rectangle in which the transform will sample
   NSAffineTransform *affTr = [NSAffineTransform transform];
   affTr.transformStruct = transform;
   [affTr invert];
   LynkeosIntegerPoint bottomLeft = {_width - 1  , _height - 1}, topRight = {0, 0};
   CGPoint p, tp;
   for (p.y = y; p.y < (y + h); p.y += (h-1))
   {
      for (p.x = x; p.x < (x + w); p.x += (w-1))
      {
         tp = [affTr transformPoint:p];
         if (tp.x < 0.0)
            tp.x = 0.0;
         else if (tp.x > (_width - 1))
            tp.x = _width - 1;
         if (tp.y < 0.0)
            tp.y = 0.0;
         else if (tp.y > (_height - 1))
            tp.y = _height - 1;

         if (tp.x < bottomLeft.x)
            bottomLeft.x = floor(tp.x);
         if (tp.y < bottomLeft.y)
            bottomLeft.y = floor(tp.y);
         if (tp.x > topRight.x)
            topRight.x = ceil(tp.x);
         if (tp.y > topRight.y)
            topRight.y = ceil(tp.y);
      }
   }
   // Apply the origin change to the transform struct
   NSAffineTransformStruct newTransform = transform;
   newTransform.tX += bottomLeft.x - x;
   newTransform.tY += bottomLeft.y - y;
   // Extract only this enclosing rectangle
   const u_short newW = topRight.x - bottomLeft.x + 1;
   const u_short newH = topRight.y - bottomLeft.y + 1;
   REAL *imageData = (REAL*)malloc(newW*newH*sizeof(REAL));

   if (ret == 0 && libraw_data->rawdata.raw_image != NULL)
   {
      for( yl = bottomLeft.y; yl <= topRight.y; yl++ )
      {
         void *linePtr = (void*)libraw_data->rawdata.raw_image + (yl+_yOrigin)*_bytesPerRow;
         for( xl = bottomLeft.x; xl <= topRight.x; xl++ )
         {
            void *pixPtr = linePtr + (xl+_xOrigin)*sizeof(u_short);

            v = (REAL)(((uint16*)pixPtr)[0])*_scale;
            v *= _whiteBalance[colorSelector(libraw_data, xl, yl)];
            imageData[xl-bottomLeft.x+newW*(yl-bottomLeft.y)] = v;
         }
      }
   }

   image = [[[BayerImageBuffer alloc] initWithData:imageData
                                      opaque:libraw_data getPlane:colorSelector
                                            width:newW lineW:newW height:newH
                                           xoffset:bottomLeft.x yoffset:bottomLeft.y
                                              atX:0 Y:0 W:w H:h
                                    withTransform:newTransform withColorAlign:offsets
                                         withDark: _darkFrame withFlat:_flatField] autorelease];
   libraw_close(libraw_data);
   free(imageData);

   return image;
}

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

@end
