//
//  Lynkeos
//  $Id: $
//
//  Created by Jean-Etienne LAMIAUD on Thu 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/LynkeosInterpolator.h>

#include "BayerImageBuffer.h"

typedef REAL WIDE_REAL_VECT_t __attribute__ ((vector_size (32)));
#define wide_real_vect_length (sizeof(WIDE_REAL_VECT_t)/sizeof(REAL))

const NSString *BayerWeightKey = @"weight";
const NSString *BayerProcessedDarktKey = @"procDark";

@interface BayerImageBuffer(Private)
- (void) convertToPlanar:(REAL * const * const)planes
              withPlanes:(u_short)nPlanes lineWidth:(u_short)lineW
                     atX:(u_short)x Y:(u_short)y
                   width:(u_short)w height:(u_short)h;
@end

@implementation BayerImageBuffer(Private)
- (void) convertToPlanar:(REAL * const * const)planes
              withPlanes:(u_short)nPlanes lineWidth:(u_short)lineW
                     atX:(u_short)x Y:(u_short)y
                   width:(u_short)w height:(u_short)h
{
   u_short bx, by, p;
   u_short startx, endx, endy;

   // Process a rectangle one pixel larger, to have interpolation data available
   if (x > 0 && (x % wide_real_vect_length) != 0)
      startx = (x / wide_real_vect_length)*wide_real_vect_length; // start on a vector boundary
   else
      startx = x;
   endx = x + w;
   if (endx > _w)
      endx = _w;
   endy = y + h;
   if (endy > _h)
      endy = _h;

#warning Take into account the parallel strategy
   for (bx = startx; bx < endx; bx += wide_real_vect_length)
   {
      WIDE_REAL_VECT_t imageSegment[3][3], weightSegment[3][3]; // previous line, current line, next line, for each color
      REAL imageLeft[3][3], imageRight[3][3], weightLeft[3][3], weightRight[3][3];

      // Read the initial segments
      for (p = 0; p < 3; p++)
      {
         if (y > 0)  // Previous segment, if any
         {
            if (bx > 0 && bx >= x)
            {
               imageLeft[0][p] = stdColorValue(self,bx-1,y-1,p);
               weightLeft[0][p] = stdColorValue(_weight,bx-1,y-1,p);
            }
            imageSegment[0][p] = *(WIDE_REAL_VECT_t*)&stdColorValue(self,bx,y-1,p);
            weightSegment[0][p] = *(WIDE_REAL_VECT_t*)&stdColorValue(_weight,bx,y-1,p);
            if (bx < (_w - wide_real_vect_length) && bx <= (endx - wide_real_vect_length))
            {
               imageRight[0][p] = stdColorValue(self,bx+wide_real_vect_length,y-1,p);
               weightRight[0][p] = stdColorValue(_weight,bx+wide_real_vect_length,y-1,p);
            }
         }
         // Current segment
         if (bx > 0 && bx >= x)
         {
            imageLeft[1][p] = stdColorValue(self,bx-1,y,p);
            weightLeft[1][p] = stdColorValue(_weight,bx-1,y,p);
         }
         imageSegment[1][p] = *(WIDE_REAL_VECT_t*)&stdColorValue(self,bx,y,p);
         weightSegment[1][p] = *(WIDE_REAL_VECT_t*)&stdColorValue(_weight,bx,y,p);
         if (bx < (_w - wide_real_vect_length) && bx <= (endx - wide_real_vect_length))
         {
            imageRight[1][p] = stdColorValue(self,bx+wide_real_vect_length,y,p);
            weightRight[1][p] = stdColorValue(_weight,bx+wide_real_vect_length,y,p);
         }
      }

      for (by = y; by < endy; by++)
      {
         REAL pixel[3];

         for (p = 0; p < 3; p++)
         {
            // Read next line, for interpolation pupose
            if ( by < (_h -1) && by <= (endy - 1) )
            {
               if (bx > 0 && bx >= x)
               {
                  imageLeft[2][p] = stdColorValue(self,bx-1,by+1,p);
                  weightLeft[2][p] = stdColorValue(_weight,bx-1,by+1,p);
               }
               imageSegment[2][p] = *(WIDE_REAL_VECT_t*)&stdColorValue(self,bx,by+1,p);
               weightSegment[2][p] = *(WIDE_REAL_VECT_t*)&stdColorValue(_weight,bx,by+1,p);
               if (bx < (_w - wide_real_vect_length) && bx <= (endx - wide_real_vect_length))
               {
                  imageRight[2][p] = stdColorValue(self,bx+wide_real_vect_length,by+1,p);
                  weightRight[2][p] = stdColorValue(_weight,bx+wide_real_vect_length,by+1,p);
               }
            }
         }

         u_short dxend = endx - bx;
         if (dxend > wide_real_vect_length)
            dxend = wide_real_vect_length;
         for ( u_short dx = (bx >= x ? 0 : x-bx); dx < dxend; dx++)
         {
            for (p = 0; p < 3; p++)
            {
               double localWeight = weightSegment[1][p][dx] / _accumulations;

               // Dark frames processed for dead pixels are displayed without de-bayering
               if (_isProcessedDark)
                  pixel[p] = imageSegment[1][p][dx];

               // Above a wheight threshold, keep the pixels unchanged
               else if ( localWeight >= 0.1)
                  pixel[p] = imageSegment[1][p][dx] / localWeight;
               
               // Otherwise interpolate with neighbours having a better weight
               else
               {
                  const u_short mxl = ((bx+dx) < _w - 1 ? bx + dx + 1 : _w - 1);
                  const u_short myl = (by < _h - 1 ? by + 1 : _h - 1);
                  const u_short sxl = ((bx+dx) < 1 ? 0 : bx + dx - 1), syl = (by < 1 ? 0 : by - 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 (xl == (bx+dx) && yl == by)
                        {
                           sum += imageSegment[1][p][dx];
                           weight += localWeight;
                        }
                        else
                        {
                           double pixelValue, otherWeight;
                           if (xl < bx)
                           {
                              pixelValue = (double)imageLeft[yl-by+1][p];
                              otherWeight = (double)weightLeft[yl-by+1][p]/_accumulations;
                           }
                           else if (xl >= (bx + wide_real_vect_length))
                           {
                              pixelValue = (double)imageRight[yl-by+1][p];
                              otherWeight = (double)weightRight[yl-by+1][p]/_accumulations;
                           }
                           else
                           {
                              pixelValue = (double)(imageSegment[yl-by+1][p][xl-bx]);
                              otherWeight = (double)weightSegment[yl-by+1][p][xl-bx]/_accumulations;
                           }

                           if (otherWeight > 0.0)
                           {
                              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][bx+dx-x+lineW*(by-y)] = (pixel[0]+pixel[1]+pixel[2])/3.0;
            else
            {
               for (p = 0; p < nPlanes; p++)
                  planes[p][bx+dx-x+lineW*(by-y)] = pixel[p];
            }
         }

         // Prepare next line
         for (p = 0; p < 3; p++)
         {
            imageLeft[0][p] = imageLeft[1][p];
            imageLeft[1][p] = imageLeft[2][p];
            weightLeft[0][p] = weightLeft[1][p];
            weightLeft[1][p] = weightLeft[2][p];
            imageSegment[0][p] = imageSegment[1][p];
            imageSegment[1][p] = imageSegment[2][p];
            weightSegment[0][p] = weightSegment[1][p];
            weightSegment[1][p] = weightSegment[2][p];
            imageRight[0][p] = imageRight[1][p];
            imageRight[1][p] = imageRight[2][p];
            weightRight[0][p] = weightRight[1][p];
            weightRight[1][p] = weightRight[2][p];
         }
      }
   }
}
@end

@implementation BayerImageBuffer

- (BOOL) hasCustomFormat
{
   return YES;
}

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

   return self;
}

- (id) initWithData:(REAL*)data
             opaque:(nullable void*)opaque getPlane:(int (*)(void *opaque, u_short x, u_short y))getPlane
              width:(u_short)width lineW:(u_short)lineW height:(u_short)height
            xoffset:(u_short)xoffset yoffset:(u_short)yoffset
                atX:(u_short)x Y:(u_short)y W:(u_short)w H:(u_short)h
      withTransform:(NSAffineTransformStruct)transform
     withColorAlign:(const NSPoint*)colorAlign
           withDark:(BayerImageBuffer*)dark withFlat:(LynkeosImageBuffer*)flat
{
   // Allocate an RGB image buffer
   if ((self = [super initWithNumberOfPlanes:3 width:w height:h]) != nil)
   {
      BOOL identity = (x == 0 && y == 0 && w == width && h == height
                       && transform.m11 == 1.0 && transform.m12 == 0.0 && transform.m21 == 0.0 && transform.m22 == 1.0
                       && colorAlign[0].x == 0.0 && colorAlign[0].y == 0.0
                       && colorAlign[1].x == 0.0 && colorAlign[1].y == 0.0
                       && colorAlign[2].x == 0.0 && colorAlign[2].y == 0.0);

      // And weight planes
      _weight = [[LynkeosImageBuffer imageBufferWithNumberOfPlanes:3 width:w height:h] retain];
      _accumulations = 1.0;

      // Allocate temporary images before transformation
      LynkeosImageBuffer *originalImage;
      LynkeosImageBuffer *originalWeight;
      LynkeosImageBuffer *darkImage = nil;
      LynkeosImageBuffer *darkWeight = nil;
      LynkeosImageBuffer *flatImage = nil;

      if (identity)
      {
         originalImage = self;
         originalWeight = _weight;
      }
      else
      {
         originalImage = [LynkeosImageBuffer imageBufferWithNumberOfPlanes:3
                                             width:width height:height];
         originalWeight = [LynkeosImageBuffer imageBufferWithNumberOfPlanes:3
                                              width:width height:height];
      }
      if (dark != nil)
         darkImage = [LynkeosImageBuffer subImageBufferWithImage:dark
                                                               x:xoffset y:yoffset width:width height:height];
      if (dark != nil)
         darkWeight = [LynkeosImageBuffer subImageBufferWithImage:dark->_weight
                                                                x:xoffset y:yoffset width:width height:height];
      if (flat != nil)
         flatImage = [LynkeosImageBuffer subImageBufferWithImage:flat
                                                               x:xoffset y:yoffset width:width 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 (getPlane(opaque, xl+xoffset, yl+yoffset) == 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;
               }
            }
         }
      }

      // Apply calibration frames, if any
      if (dark != nil)
      {
         [originalImage substract:darkImage];
         [originalWeight substract:darkWeight];
      }

      if (flat != nil)
      {
         [originalImage divideBy:flatImage result:originalImage];
      }

      if (!identity)
      {
         // 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:colorAlign
                                         withParameters:nil] autorelease];
         id <LynkeosInterpolator> weightInterpolator
            = [[[interpolatorClass alloc] initWithImage:originalWeight
                                                 inRect:r
                                     withNumberOfPlanes:3
                                           withTranform:transform
                                            withOffsets:colorAlign
                                         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(self, 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
{
   if ((self = [super initWithCoder:aDecoder]) != nil)
   {
      _weight = [[aDecoder decodeObjectForKey:(NSString*)BayerWeightKey] retain];
      _accumulations = 1.0;
      _isProcessedDark = [aDecoder decodeBoolForKey:(NSString*)BayerProcessedDarktKey];
   }

   return self;
}

- (nonnull id) copyWithZone:(nullable NSZone *)zone
{
   BayerImageBuffer *other = [[BayerImageBuffer allocWithZone:zone] initWithData:_data
                                                                          copy:YES
                                                                  freeWhenDone:YES
                                                                numberOfPlanes:_nPlanes
                                                                         width:_w
                                                                   paddedWidth:_padw
                                                                        height:_h];
   if (other != nil)
   {
      other->_weight = [_weight copyWithZone:zone];
      other->_accumulations = _accumulations;
      other->_isProcessedDark = _isProcessedDark;
   }

   return other;
}

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

- (size_t) memorySize
{
   return([super memorySize] + [_weight memorySize]);
}

- (u_short) width {return _w;}

- (u_short) height {return _h;}

- (u_short) numberOfPlanes {return 3;}

#if !GNUSTEP
- (CGImageRef) getImageInRect:(LynkeosIntegerRect)r
                    withBlack:(double*)black white:(double*)white gamma:(double*)gamma
{
   LynkeosImageBuffer *img = [LynkeosImageBuffer imageBufferWithNumberOfPlanes:3
                                                                         width:r.size.width
                                                                        height:r.size.height];

   [self convertToPlanar:[img colorPlanes] withPlanes:3 lineWidth:img->_padw
                     atX:r.origin.x Y:r.origin.y width:r.size.width height:r.size.height];
   [self getMinLevel:&img->_min[3] maxLevel:&img->_max[3]];
   for( int c = 0; c < 3; c++)
      [self getMinLevel:&img->_min[c] maxLevel:&img->_max[c] forPlane:c];

   LynkeosIntegerRect extract = {{0,0},r.size};
   return [img getImageInRect:extract withBlack:black white:white gamma:gamma];
}
#endif

- (void) getMinLevel:(double *)vmin maxLevel:(double *)vmax
{
   if ( _min[_nPlanes] >= _max[_nPlanes] )
   {
      int x, y, c;
      double n, ntotal;

      ntotal = 0.0;
      for( c = 0; c <= _nPlanes; c++ )
      {
         _min[c] = HUGE_VAL;
         _max[c] = -HUGE_VAL;
         _mean[c] = 0.0;
      }
      for( c = 0; c < _nPlanes; c++ )
      {
         n = 0.0;
         for( y = 0; y < _h; y++ )
         {
            for( x = 0; x < _w; x++ )
            {
               double w = stdColorValue(_weight, x, y, c);
               // Only take non null weight into account
               if (w != 0.0)
               {
                  double v = stdColorValue(self,x,y,c)/w;
                  if ( _min[c] > v )
                     _min[c] = v;
                  if ( _max[c] < v )
                     _max[c] = v;
                  _mean[c] += v;
                  _mean[_nPlanes] += v;
                  n += 1.0;
               }
            }
         }
         if ( _min[_nPlanes] > _min[c] )
            _min[_nPlanes] = _min[c];
         if ( _max[_nPlanes] < _max[c] )
            _max[_nPlanes] = _max[c];
         _mean[c] /= n;
         ntotal += n;
      }

      _mean[_nPlanes] /= ntotal;
   }

   *vmin = _min[_nPlanes];
   *vmax = _max[_nPlanes];
}

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

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

// This is used only on callibration frames
- (void) normalizeWithFactor:(double)factor mono:(BOOL)mono
{
   if (factor == 0 || mono)
   {
      // We need to debayer first
      LynkeosImageBuffer *buf
         = [[LynkeosImageBuffer alloc] initWithNumberOfPlanes:_nPlanes width:_w height:_h];
      [self convertToPlanar:[buf colorPlanes] withPlanes:_nPlanes lineWidth:buf->_padw];
      [buf extractSample:_planes atX:0 Y:0 withWidth:_w height:_h withPlanes:_nPlanes lineWidth:_padw];
      // Set a constant one weight
      _accumulations = 1.0;
      REALVECT one = {1.0, 1.0, 1.0, 1.0};
      for (u_short p = 0; p < _nPlanes; p++)
      {
         for (u_short y = 0; y < _h; y++)
         {
            for (u_short x = 0; x < _w; x+= sizeof(REALVECT)/sizeof(REAL))
               *((REALVECT*)&_weight->_planes[p][y*_padw+x]) = one;
         }
      }
      // And normalize
      [super normalizeWithFactor:factor mono:mono];
   }
   else
   {
      // Apply the normalization on the image and weights
      [super normalizeWithFactor:factor mono:mono];
      [_weight normalizeWithFactor:factor mono:mono];
      _accumulations *= factor;
   }
}

- (void) convertToPlanar:(REAL * const * const)planes
              withPlanes:(u_short)nPlanes
               lineWidth:(u_short)lineW
{
   [self convertToPlanar:planes withPlanes:nPlanes lineWidth:lineW atX:0 Y:0 width:_w height:_h];
}

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

- (void)setOperatorsStrategy:(ImageOperatorsStrategy_t)strategy
{
   [super setOperatorsStrategy:strategy];
   [_weight setOperatorsStrategy:strategy];
}

- (void) encodeWithCoder:(nonnull NSCoder *)aCoder
{
   // Set the accumulations number to one
   const double factor = 1.0 / _accumulations;
   [self multiplyWithScalar:factor];
   [_weight multiplyWithScalar:factor];
   // And save the object
   [super encodeWithCoder:aCoder];
   [aCoder encodeObject:_weight forKey:(NSString*)BayerWeightKey];
   [aCoder encodeBool:_isProcessedDark forKey:(NSString*)BayerProcessedDarktKey];
}
@end