//
//  Lynkeos
//  $Id: $
//
//  Created by Jean-Etienne LAMIAUD on Thu Oct 18 2018.
//
//  Copyright (c) 2018. 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 <objc/objc-class.h>

#include <LynkeosCore/LynkeosInterpolator.h>

#include "SER_ImageBuffer.h"

@interface SER_ImageBuffer(Private)
- (LynkeosStandardImageBuffer *) planarImage;
@end

@implementation SER_ImageBuffer(Private)
- (LynkeosStandardImageBuffer *) planarImage
{
   LynkeosStandardImageBuffer *img
      = [LynkeosStandardImageBuffer imageBufferWithNumberOfPlanes:3
                                                            width:[_image width]
                                                           height:[_image height]];
   [self convertToPlanar:[img colorPlanes] withPlanes:img->_nPlanes lineWidth:img->_padw];

   return img;
}
@end

@implementation SER_ImageBuffer

- (id) init
{
   if ( (self = [super init]) != nil)
   {
      _image = nil;
      _weight = nil;
      _accumulations = 0;
      _substractive = NO;
   }

   return self;
}

- (id) initWithData:(REAL*)data format:(ColorID_t)format
              width:(u_short)width lineW:(u_short)lineW height:(u_short)height
                atX:(u_short)x Y:(u_short)y W:(u_short)w H:(u_short)h
      withTransform:(NSAffineTransformStruct)transform
        withOffsets:(const NSPoint*)offsets
{
   if ((self = [super init]) != nil)
   {
      u_short bayerPlanes[2][2];

      // Allocate an RGB image buffer
      _image = [[LynkeosStandardImageBuffer imageBufferWithNumberOfPlanes:3 width:w height:h] retain];
      // And weight planes
      _weight = [[LynkeosStandardImageBuffer imageBufferWithNumberOfPlanes:3 width:w height:h] retain];
      _accumulations = 1.0;

      switch (format)
      {
         case SER_BAYER_CYYM:
         case SER_BAYER_YCMY:
         case SER_BAYER_YMCY:
         case SER_BAYER_MYYC:
            _substractive = YES;
            break;
         default:
            _substractive = NO;
            break;
      }

      // Fill in the pixels in their respective planes
      switch (format)
      {
         case SER_BAYER_CYYM: // Planes in CYM order
         case SER_BAYER_RGGB:
            bayerPlanes[0][0] = 0; bayerPlanes[0][1] = 1;
            bayerPlanes[1][0] = 1; bayerPlanes[1][1] = 2;
            break;
         case SER_BAYER_YCMY: // Planes in CYM order
         case SER_BAYER_GRBG:
            bayerPlanes[0][0] = 1; bayerPlanes[0][1] = 0;
            bayerPlanes[1][0] = 2; bayerPlanes[1][1] = 1;
            break;
         case SER_BAYER_YMCY: // Planes in CYM order
         case SER_BAYER_GBRG:
            bayerPlanes[0][0] = 1; bayerPlanes[0][1] = 2;
            bayerPlanes[1][0] = 0; bayerPlanes[1][1] = 1;
            break;
         case SER_BAYER_MYYC: // Planes in CYM order
         case SER_BAYER_BGGR:
            bayerPlanes[0][0] = 2; bayerPlanes[0][1] = 1;
            bayerPlanes[1][0] = 1; bayerPlanes[1][1] = 0;
            break;
         default: // Non bayer format
            NSAssert(NO, @"Non Bayer format in SER_ImageBuffer");
            break;
      }

      // Allocate temporary images before transformation
      LynkeosStandardImageBuffer *originalImage
         = [LynkeosStandardImageBuffer imageBufferWithNumberOfPlanes:3 width:lineW height:height];
      LynkeosStandardImageBuffer *originalWeight
         = [LynkeosStandardImageBuffer imageBufferWithNumberOfPlanes:3 width:lineW height:height];

      u_short xl, yl, p;
      for (yl = 0; yl < height; yl++)
      {
         for (xl = 0; xl < width; xl++)
         {
            for (p = 0; p < 3; p++)
            {
               if (bayerPlanes[yl%2][xl%2] == p)
               {
                  stdColorValue(originalImage, xl, yl, p) = data[xl + yl*lineW];
                  stdColorValue(originalWeight, xl, yl, p) = 1.0;
               }
               else
               {
                  stdColorValue(originalImage, xl, yl, p) = 0.0;
                  stdColorValue(originalWeight, xl, yl, p) = 0.0;
               }
            }
         }
      }

      // And interpolate to fill the image and weight
      const LynkeosIntegerRect r = {{x, y}, {w, h}};
      Class interpolatorClass = [LynkeosInterpolatorManager interpolatorWithScaling:UseTransform
                                                                          transform:transform];
      id <LynkeosInterpolator> imageInterpolator
         = [[[interpolatorClass alloc] initWithImage:originalImage
                                              inRect:r
                                  withNumberOfPlanes:3
                                        withTranform:transform
                                         withOffsets:offsets
                                      withParameters:nil] autorelease];
      id <LynkeosInterpolator> weightInterpolator
         = [[[interpolatorClass alloc] initWithImage:originalWeight
                                              inRect:r
                                  withNumberOfPlanes:3
                                        withTranform:transform
                                         withOffsets:offsets
                                      withParameters:nil] autorelease];

      for (p = 0; p < 3; p++)
      {
         for (yl = 0; yl < h; yl++)
         {
            for (xl = 0; xl < w; xl++)
            {
               REAL v;
               v = [imageInterpolator interpolateInPLane:p atX:xl atY:yl];
               if (isnan(v))
                  NSLog(@"NaN pixel value at %d %d, in plane %d", xl, yl, p);
               else
                  stdColorValue(_image, xl, yl, p) = v;
               v = [weightInterpolator interpolateInPLane:p atX:xl atY:yl];
               if (isnan(v))
                  NSLog(@"NaN pixel weight at %d %d, in plane %d", xl, yl, p);
               else
                  stdColorValue(_weight, xl, yl, p) = v;
            }
         }
      }
   }

   return(self);
}

- (nullable instancetype) initWithCoder:(nonnull NSCoder *)aDecoder
{
   [self doesNotRecognizeSelector:_cmd];
   return nil;
}

- (nonnull id) copyWithZone:(nullable NSZone *)zone
{
   SER_ImageBuffer *other = [[SER_ImageBuffer allocWithZone:zone] init];
   other->_image = [_image copyWithZone:zone];
   other->_weight = [_weight copyWithZone:zone];
   other->_accumulations = _accumulations;
   other->_substractive = _substractive;

   return other;
}

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

- (size_t) memorySize
{
   return(class_getInstanceSize([self class])
          + [_image memorySize]
          + [_weight memorySize]);
}

- (u_short) width {return [_image width];}

- (u_short) height {return [_image height];}

- (u_short) numberOfPlanes {return 3;}

- (NSBitmapImageRep*) getNSImageWithBlack:(double*)black white:(double*)white
                                    gamma:(double*)gamma
{
   LynkeosStandardImageBuffer *img = [self planarImage];

   return [img getNSImageWithBlack:black white:white gamma:gamma];
}

- (void) add :(id <LynkeosImageBuffer>)image
{
   NSAssert([image isKindOfClass:[self class]], @"SER_ImageBuffer can only add with itself");
   SER_ImageBuffer *other = (SER_ImageBuffer*)image;
   [_image add:other->_image];
   [_weight add:other->_weight];
   _accumulations += other->_accumulations;
}

- (void) calibrateWithDarkFrame:(id <LynkeosImageBuffer>)darkFrame
                      flatField:(id <LynkeosImageBuffer>)flatField
                            atX:(u_short)ox Y:(u_short)oy
{
   // Nothing to do. Calibration occurs in the reader
}

- (void) normalizeWithFactor:(double)factor mono:(BOOL)mono
{
   // To do
}

- (void) convertToPlanar:(REAL * const * const)planes
              withPlanes:(u_short)nPlanes
               lineWidth:(u_short)lineW
{
   const u_short w = [_image width], h = [_image height];
   u_short x, y, p;

   for (y = 0; y < h; y++)
   {
      for (x = 0; x < w; x++)
      {
         REAL pixel[3];

         for (p = 0; p < 3; p++)
         {
            double localWeight = stdColorValue(_weight, x, y, p) / _accumulations;
            // Above a wheight threshold, keep the pixels unchanged
            if ( localWeight >= 0.25)
               pixel[p] = stdColorValue(_image, x, y, p) / localWeight;

            // Otherwise interpolate with neighbours having a better weight
            else
            {
               const u_short mxl = (x < w - 1 ? x + 1 : w - 1);
               const u_short myl = (y < h - 1 ? y + 1 : h - 1);
               const u_short sxl = (x < 1 ? 0 : x - 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++)
                  {
                     double pixelValue = stdColorValue(_image, xl, yl, p);
                     if (xl == x && yl == y)
                     {
                        sum += pixelValue;
                        weight += localWeight;
                     }
                     else
                     {
                        double otherWeight = stdColorValue(_weight, xl, yl, p)/_accumulations;
                        double interpolationWeight = otherWeight - localWeight;

                        if (interpolationWeight > 0.0)
                        {
                           sum += (pixelValue / otherWeight) * interpolationWeight;
                           weight += interpolationWeight;
                        }
                     }
                  }
               }
               pixel[p] = (weight != 0.0 ? sum / weight : 0.0);
            }
         }
         if (nPlanes == 1)
            // Convert to monochrome
            planes[0][x+lineW*y] = (pixel[0]+pixel[1]+pixel[2])/3.0;
         else
         {
#warning Convert CMY to RGB when needed
            for (p = 0; p < nPlanes; p++)
               planes[p][x+lineW*y] = pixel[p];
         }
      }
   }
}

- (void) clear
{
   [_image clear];
   [_weight clear];
}


- (void) encodeWithCoder:(nonnull NSCoder *)aCoder
{
   // No coding, this class is only used during processing
   [self doesNotRecognizeSelector:_cmd];
}
@end
