/************************************************************************/
/*                                                                      */
/*    vspline - a set of generic tools for creation and evaluation      */
/*              of uniform b-splines                                    */
/*                                                                      */
/*            Copyright 2017 - 2023 by Kay F. Jahnke                     */
/*                                                                      */
/*    Permission is hereby granted, free of charge, to any person       */
/*    obtaining a copy of this software and associated documentation    */
/*    files (the "Software"), to deal in the Software without           */
/*    restriction, including without limitation the rights to use,      */
/*    copy, modify, merge, publish, distribute, sublicense, and/or      */
/*    sell copies of the Software, and to permit persons to whom the    */
/*    Software is furnished to do so, subject to the following          */
/*    conditions:                                                       */
/*                                                                      */
/*    The above copyright notice and this permission notice shall be    */
/*    included in all copies or substantial portions of the             */
/*    Software.                                                         */
/*                                                                      */
/*    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND    */
/*    EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES   */
/*    OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND          */
/*    NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT       */
/*    HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,      */
/*    WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING      */
/*    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR     */
/*    OTHER DEALINGS IN THE SOFTWARE.                                   */
/*                                                                      */
/************************************************************************/

/// ca_correct.cc
///
/// Perform correction of chromatic aberration using a cubic polynomial
/// This uses panotools-compatible parameters. Currently only processes
/// 8bit RGB, processing is done on the sRGB data without conversion
/// to linear and back, images with alpha channel won't compute.
/// To see how the panotools lens correction model functions,
/// please refer to https://wiki.panotools.org/Lens_correction_model
///
/// compile with:
/// clang++ -std=c++11 -march=native -o ca_correct -O3 -pthread -DUSE_VC ca_correct.cc -lvigraimpex -lVc
///
/// you can also use g++ instead of clang++. If you don't have Vc on
/// your system, omit '-DUSE_VC' and '-lVc'. To use highway instead of
/// Vc, use -DUSE_HWY and -lhwy, and to use std::simd, use -DUSE_STSIMD,
/// and compile with c++17 standard. Newer versions of gcc provide a std::simd
/// implementation, if that is missing you need to install it separately.
///
/// invoke with ca_correct <image> ar br cr dr ag bg cg dg ab bb cb db d e
/// where 'ar' stands for 'parameter a for red channel' etc., and the
/// trailing d and e are center shift in x and y direction in pixels.
///
/// The purpose here is more to demonstrate how to implement the maths
/// using vspline, but the program could easily be fleshed out to
/// become really usable:
///
/// - add alpha channel treatment
/// - add processing of 16bit data
/// - add colour space management and internal linear processing

// TODO: handle incoming alpha channel
// TODO: differentiate between 8/16 bit
// TODO: operate in linear RGB
// TODO: may produce transparent output where coordinate is out-of-range
// TODO might allow to parametrize normalization (currently using PT's way)
// TODO: while emulating the panotools way of ca correction is nice
// to have, try implementing shift-curve based correction: with
// pixel's radial distance r to origin (normalized to [0,1])
// access a 1D b-spline over M data points, so that
// res = ev ( r * ( M - 1 ) )
// and picking up from source at res instead of r.
// the coefficients of the shift curve can be generated by sampling the
// 'normal' method or any other model, it's methodically neutral, within
// the fidelity of the spline to the approximated signal, which can be
// taken arbitrarily high.
// TODO: consider two more shift curves Cx and Cy defined over [-1,1]
// then, if the incoming coordinate is(x,y), let
// x' = Cx ( x ) and y' = Cy ( y ). This would introduce a tensor-
// based component, usable for sensor tilt compensation (check!)

#include <iostream>

// <vspline/vspline.h> pulls in all of vspline's functionality

#include <vspline/vspline.h>

// in addition to <vigra/tinyvector.hxx> and <vigra/multi_array.hxx>,
// which are necessarily included by vspline, we want to use vigra's
// import/export facilities to load and store images:

#include <vigra/stdimage.hxx>
#include <vigra/imageinfo.hxx>
#include <vigra/impex.hxx>

// we'll be working with float data, so we set up a program-wide
// constant for the vector width appropriate for float data

const int VSIZE = vspline::vector_traits < float > :: size ;

// we silently assume we have a colour image

typedef vigra::RGBValue<float,0,1,2> pixel_type; 

// coordinate_type is a 2D coordinate

typedef vigra::TinyVector < float , 2 > coordinate_type ;

// type of b-spline object used as interpolator

typedef vspline::bspline < pixel_type , 2 > spline_type ;

// target_type is a 2D array of pixels 

typedef vigra::MultiArray < 2 , pixel_type > target_type ;

// type of b-spline evaluator producing single floats

typedef vspline::evaluator < coordinate_type , // incoming coordinate's type
                             float             // processing single channel data
                           > ev_type ;

// gate type to force singular coordinates to a range
// we are using mirror boundary conditions. Note that this may
// produce artifacts in the margins.
                           
typedef vspline::mirror_gate < float > gate_type ;

// mapper uses two gates, for x and y component

typedef vspline::map_functor < coordinate_type , VSIZE ,
                               gate_type , gate_type > mapper_type ;

// we start out by coding the functor implementing
// the coordinate transformation for a single channel.
// we inherit from vspline::unary_functor so that our coordinate
// transformation fits well into vspline's functional processing scheme

struct ev_radial_correction
: public vspline::unary_functor < float , float >
{
  // incoming coordinates are shifted by dx and dy. These values are expected
  // in image coordinates. The shift values should be so that a pixel which
  // is located at the intersection of the sensor with the optical axis comes
  // out as (0,0). If the optical system is perfectly centered, dx and dy will
  // be the coordinates of the image center, so if the image is size X * Y,
  // dx = ( X - 1 ) / 2
  // dy = ( Y - 1 ) / 2
  
  const float dx , dy ;
  
  // next we have a scaling factor. Once the coordinates are shifted to
  // coordinates relative to the optical axis, we apply a scaling factor
  // which 'normalizes' the pixel's distance from the optical axis.
  // A typical scaling factor would be the distance of the image center
  // from the top left corner at (-0.5,-0.5). With dx and dy as above
  // dxc = dx - -0.5 ;
  // dyc = dy - -0.5 ;
  // scale = 1 / sqrt ( dx * dx + dy * dy ) ;
  // Here we use a different choice to be compatible with panotools:
  // we use the vertical distance from image center to top/bottom
  // margin.
  // Since we'll be using a polynomial over the radial distance, picking
  // scale values larger or smaller than this 'typical' value can be used
  // to affect the precise effect of the radial function.
  // rscale is simply the reciprocal value for faster computation.
  
  const float scale ;
  const float rscale ;
  
  // After applying the scale, we have a normalized coordinate. The functor will
  // use this normalized coordinate to calculate the normalized distance from
  // the optical axis. The resulting distance is the argument to the radial
  // correction function. For the radial correction function, we use a cubic
  // polynomial, which needs four coefficients:
  
  const float a , b , c , d ;
  
  // finally we have the PT d and e values, which we label x_shift and y_shift
  // to avoid confusion with the fourth coefficient of the polynomial.
  
  const float x_shift , y_shift ;
  
  // we use two static functions to concisely initialize some of the
  // constant values above
  
  static double d_from_extent ( double d )
  {
    return ( d - 1.0 ) / 2.0 ;
  }
  
  static double rscale_from_wh ( double w , double h )
  {
    double dx = d_from_extent ( w ) ;
    double dy = d_from_extent ( h ) ;
    // I'd normalize to the corner, but to be compatible with panotools,
    // I use normalization to top margin center instead.
    //     return sqrt ( dx * dx + dy * dy ) ;
    return sqrt ( dy * dy ) ;
  }
  
  // here's the constructor for the radial correction functor, taking all
  // the values passed from main() and initializing the constants

  ev_radial_correction ( const double & _width ,
                         const double & _height ,
                         const double & _x_shift ,
                         const double & _y_shift ,
                         const double & _a ,
                         const double & _b ,
                         const double & _c ,
                         const double & _d )
  : dx ( d_from_extent ( _width ) ) ,
    dy ( d_from_extent ( _height ) ) ,
    x_shift ( _x_shift ) ,
    y_shift ( _y_shift ) ,
    rscale ( rscale_from_wh ( _width , _height ) ) ,
    scale ( 1.0 / rscale_from_wh ( _width , _height ) ) ,
    a ( _a ) ,
    b ( _b ) ,
    c ( _c ) ,
    d ( _d )
  {
    // we echo the internal state
    std::cout << "dx:     " << dx << std::endl ;
    std::cout << "dy:     " << dy << std::endl ;
    std::cout << "scale:  " << scale << std::endl ;
    std::cout << "rscale: " << rscale << std::endl ;
    std::cout << "a:      " << a << std::endl ;
    std::cout << "b:      " << b << std::endl ;
    std::cout << "c:      " << c << std::endl ;
    std::cout << "d:      " << d << std::endl ;
  } ;
 
  // now we provide evaluation code for the functor.
  // since the code is the same for vectorized and unvectorized
  // operation, we can write a template, In words:
  // eval is a function template with the coordinate type as it's
  // template argument. eval receives it's argument as a const
  // reference to a coordinate and deposits it's result to a reference
  // to a coordinate. This function will not change the state of the
  // functor (hence the const) - the functor does not have mutable state
  // anyway. Note how CRD can be a single coordinate_type, or it's
  // vectorized equivalent.
  
  template < class CRD >
  void eval ( const CRD & in ,
                    CRD & result ) const
  {
    // set up coordinate-type variable to work on, copy input to it.
    
    CRD cc ( in ) ;
    
    // shift and scale
    // TODO: is it right to add the shift here, or should I subtract

    cc[0] -= ( dx + x_shift ) ;
    cc[0] *= scale ;
    cc[1] -= ( dy + y_shift ) ;
    cc[1] *= scale ;
    
    // calculate distance from center (this is normalized due to scaled cc)
  
    auto r = sqrt ( cc[0] * cc[0] + cc[1] * cc[1] ) ;
    
    // apply polynomial to obtain the scaling factor.

    auto rr = a * r * r * r + b * r * r + c * r + d ;
    
    // use rr to scale cc - this is the radial correction

    cc[0] *= rr ;
    cc[1] *= rr ;
    
    // apply rscale to revert to centered image coordinates

    cc[0] *= rscale ;
    cc[1] *= rscale ;
    
    // reverse initial shift to arrive at UL-based image coordinates

    cc[0] += ( dx + x_shift ) ;
    cc[1] += ( dy + y_shift ) ;
    
    // assign to result

    result = cc ;
  }  
} ;

// next we set up the functor processing all three channels. This functor
// receives three ev_radial_correction functors and three channel views:

struct ev_ca_correct
: public vspline::unary_functor < coordinate_type , pixel_type >
{
  // these three functors hold the radial corrections for the three
  // colour channels

  ev_radial_correction rc_red ;
  ev_radial_correction rc_green ;
  ev_radial_correction rc_blue ;
  
  // these three functors hold interpolators for the colour channels
  
  ev_type ev_red ;
  ev_type ev_green ;
  ev_type ev_blue ;
  
  // and this object deals with out-of-bounds coordinates
  
  mapper_type m ;
  
  // the constructor receives all the functors we'll use. Note how we
  // can simply copy-construct the functors.
  
  ev_ca_correct ( const ev_radial_correction & _rc_red ,
                  const ev_radial_correction & _rc_green ,
                  const ev_radial_correction & _rc_blue ,
                  const ev_type & _ev_red ,
                  const ev_type & _ev_green ,
                  const ev_type & _ev_blue ,
                  const mapper_type & _m )
  : rc_red ( _rc_red ) ,
    rc_green ( _rc_green ) ,
    rc_blue ( _rc_blue ) ,
    ev_red ( _ev_red ) ,
    ev_green ( _ev_green ) ,
    ev_blue ( _ev_blue ) ,
    m ( _m )
    { } ;
    
  // the eval routine is simple, it simply applies the coordinate
  // transformation, applies the mapper to force the transformed
  // coordinate into the range, an then picks the interpolated value
  // using the interpolator for the channel. This is done for all
  // channels in turn.

  // since the code is the same for vectorized and unvectorized
  // operation, we can again write a template:
  
  template < class IN , class OUT >
  void eval ( const IN & c ,
                    OUT & result ) const
  {
    // work variable containing a (possibly vectorized) 2D coordinate

    IN cc ;
    
    // apply the radial correction to the incoming coordinate in c,
    // storing result to cc. Note that c contains the 'target' coordinate:
    // The coordinate of the pixel in the target which we want to compute
    
    rc_red.eval   ( c , cc ) ;
    
    // force coordinate into the defined range (here we use mirroring)
    
    m.eval ( cc , cc ) ;
    
    // evaluate channel view at corrected coordinate, storing result
    // to the red channel of 'result'
    
    ev_red.eval   ( cc , result[0] ) ;
    
    // ditto, for the remaining channels

    rc_green.eval ( c , cc ) ;     
    m.eval ( cc , cc ) ;
    ev_green.eval ( cc , result[1] ) ;

    rc_blue.eval  ( c , cc ) ;     
    m.eval ( cc , cc ) ;
    ev_blue.eval  ( cc , result[2] ) ;
  }
    
} ;

int main ( int argc , char * argv[] )
{
  if ( argc < 11 )
  {
    std::cerr << "pass a colour image file as first argument" << std::endl ;
    std::cerr << "followed by a, b, c for red, green, blue" << std::endl ;
    std::cerr << "and the horizontal and vertical shift" << std::endl ;
    std::cerr << "like ca_correct xx.jpg 0.0001411 -0.0005236 0.0008456 1.0002093 0 0 0 1 0.0002334 -0.0007607 0.0011446 0.9996757 176 116"
              << std::endl ;
    exit( -1 ) ;
  }

  double ar = atof ( argv[2] ) ;
  double br = atof ( argv[3] ) ;
  double cr = atof ( argv[4] ) ;
  double dr = atof ( argv[5] ) ;
  double ag = atof ( argv[6] ) ;
  double bg = atof ( argv[7] ) ;
  double cg = atof ( argv[8] ) ;
  double dg = atof ( argv[9] ) ;
  double ab = atof ( argv[10] ) ;
  double bb = atof ( argv[11] ) ;
  double cb = atof ( argv[12] ) ;
  double db = atof ( argv[13] ) ;
  double x_shift = atof ( argv[14] ) ; // aka panotools 'd'
  double y_shift = atof ( argv[15] ) ; // aka panotools 'e'
  
  // get the image file name
  
  vigra::ImageImportInfo imageInfo ( argv[1] ) ;

  // we want a b-spline with natural boundary conditions
  
  vigra::TinyVector < vspline::bc_code , 2 > bcv ( vspline::NATURAL ) ;
  
  // create cubic 2D b-spline object containing the image data
  // TODO allow passing in arbitrary spline order
  
  spline_type bspl ( imageInfo.shape() , // the shape of the data for the spline
                     5 ,                 // degree 5 == quintic spline
                     bcv                 // specifies natural BCs along both axes
                   ) ;
  
  // load the image data into the b-spline's core. This is a common idiom:
  // the spline's 'core' is a MultiArrayView to that part of the spline's
  // data container which precisely corresponds with the input data.
  // This saves loading the image to some memory first and then transferring
  // the data into the spline. Since the core is a vigra::MultiarrayView,
  // we can pass it to importImage as the desired target for loading the
  // image from disk.
  
  std::cout << "reading image " << argv[1] << " from disk" << std::endl ;
  
  vigra::importImage ( imageInfo , bspl.core ) ;
  
  // prefilter the b-spline

  std::cout << "setting up b-spline interpolator for image data" << std::endl ;
  
  bspl.prefilter() ;
  
  // this is where the result should go:
  
  target_type target ( imageInfo.shape() ) ;
  
  // process the image metrics

  float width = imageInfo.width() ;
  float height = imageInfo.height() ;

  // set up the radial transformation functors
  
  std::cout << "setting up radial correction for red channel:" << std::endl ;
  
  ev_radial_correction ca_red
    ( width , height , x_shift , y_shift , ar , br , cr , dr ) ;

  std::cout << "setting up radial correction for green channel:" << std::endl ;
  
  ev_radial_correction ca_green
    ( width , height , x_shift , y_shift , ag , bg , cg , dg ) ;
    
  std::cout << "setting up radial correction for blue channel:" << std::endl ;
  
  ev_radial_correction ca_blue
    ( width , height , x_shift , y_shift , ab , bb , cb , db ) ;
  
  // here we create the channel views.

  auto red_channel = bspl.get_channel_view ( 0 ) ;
  auto green_channel = bspl.get_channel_view ( 1 ) ;
  auto blue_channel = bspl.get_channel_view ( 2 ) ;
  
  // and set up the per-channel interpolators

  ev_type red_ev ( red_channel ) ;
  ev_type green_ev ( green_channel ) ;
  ev_type blue_ev ( blue_channel ) ;
  
  // next we set up coordinate mapping to the defined range
  
  gate_type g_x ( 0.0 , width - 1.0 ) ;
  gate_type g_y ( 0.0 , height - 1.0 ) ;
  mapper_type m ( g_x , g_y ) ;

  // using vspline's factory functions to create the 'gates' and the
  // 'mapper' applying them, we could instead create m like this:
  // (Note how we have to be explicit about using 'float'
  // for the arguments to the gates - using double arguments would not
  // work here unless we'd also specify the vector width.)
  
//   auto m = vspline::mapper < coordinate_type >
//                ( vspline::mirror ( 0.0f , width - 1.0f ) ,
//                  vspline::mirror ( 0.0f , height - 1.0f ) ) ;                 

  // finally, we create the top-level functor, passing in the three
  // radial correction functors, the channel-wise evaluators and the
  // mapper object

  ev_ca_correct correct ( ca_red , ca_green , ca_blue ,
                          red_ev , green_ev , blue_ev ,
                          m ) ;

  // now we obtain the result by performing a vspline::transform. this transform
  // successively passes discrete coordinates into the target to the functor
  // it's invoked with, storing the result of the functor's evaluation
  // at the self-same coordinates in it's target, so for each coordinate
  // (X,Y), target[(X,Y)] = correct(X,Y)
  
  std::cout << "rendering the target image" << std::endl ;
  
  vspline::transform ( correct , target ) ;

  // store the result with vigra impex

  std::cout << "storing the target image as 'ca_correct.tif'" << std::endl ;
  
  vigra::ImageExportInfo eximageInfo ( "ca_correct.tif" );
  
  vigra::exportImage ( target ,
                       eximageInfo
                       .setPixelType("UINT8")
                       .setForcedRangeMapping ( 0 , 255 , 0 , 255 ) ) ;
  
  std::cout << "done" << std::endl ;
}