// NeuQuant Neural-Net Quantization Algorithm

// ------------------------------------------

//

// Copyright (c) 1994 Anthony Dekker

//

// NEUQUANT Neural-Net quantization algorithm by Anthony Dekker, 1994.

// See "Kohonen neural networks for optimal colour quantization"

// in "Network: Computation in Neural Systems" Vol. 5 (1994) pp 351-367.

// for a discussion of the algorithm.

//

// Any party obtaining a copy of these files from the author, directly or

// indirectly, is granted, free of charge, a full and unrestricted irrevocable,

// world-wide, paid up, royalty-free, nonexclusive right and license to deal

// in this software and documentation files (the "Software"), including without

// limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,

// and/or sell copies of the Software, and to permit persons who receive

// copies from any such party to do so, with the only requirement being

// that this copyright notice remain intact.

///////////////////////////////////////////////////////////////////////

// History

// -------

// January 2001: Adaptation of the Neural-Net Quantization Algorithm

//               for the FreeImage 2 library

//               Author: Hervé Drolon (drolon@infonie.fr)

// March 2004:   Adaptation for the FreeImage 3 library (port to big endian processors)

//               Author: Hervé Drolon (drolon@infonie.fr)

// April 2004:   Algorithm rewritten as a C++ class. 

//               Fixed a bug in the algorithm with handling of 4-byte boundary alignment.

//               Author: Hervé Drolon (drolon@infonie.fr)

///////////////////////////////////////////////////////////////////////

#include "Quantizers.h"

#include "FreeImage.h"

#include "Utilities.h"

// Four primes near 500 - assume no image has a length so large

// that it is divisible by all four primes

// ==========================================================

#define prime1    499

#define prime2    491

#define prime3    487

#define prime4    503

// ----------------------------------------------------------------

NNQuantizer::NNQuantizer(int PaletteSize)

{

  netsize = PaletteSize;

  maxnetpos = netsize - 1;

  initrad = netsize < 8 ? 1 : (netsize >> 3);

  initradius = (initrad * radiusbias);

  network = NULL;

  network = (pixel *)malloc(netsize * sizeof(pixel));

  bias = (int *)malloc(netsize * sizeof(int));

  freq = (int *)malloc(netsize * sizeof(int));

  radpower = (int *)malloc(initrad * sizeof(int));

  if( !network || !bias || !freq || !radpower ) {

    if(network) free(network);

    if(bias) free(bias);

    if(freq) free(freq);

    if(radpower) free(radpower);

    throw FI_MSG_ERROR_MEMORY;

  }

}

NNQuantizer::~NNQuantizer()

{

  if(network) free(network);

  if(bias) free(bias);

  if(freq) free(freq);

  if(radpower) free(radpower);

}

///////////////////////////////////////////////////////////////////////////

// Initialise network in range (0,0,0) to (255,255,255) and set parameters

// -----------------------------------------------------------------------

void NNQuantizer::initnet() {

  int i, *p;

  for (i = 0; i < netsize; i++) {

    p = network[i];

    p[FI_RGBA_BLUE] = p[FI_RGBA_GREEN] = p[FI_RGBA_RED] = (i << (netbiasshift+8))/netsize;

    freq[i] = intbias/netsize;  /* 1/netsize */

    bias[i] = 0;

  }

}

/////////////////////////////////////////////////////////////////////////////////////// 

// Unbias network to give byte values 0..255 and record position i to prepare for sort

// ------------------------------------------------------------------------------------

void NNQuantizer::unbiasnet() {

  int i, j, temp;

  for (i = 0; i < netsize; i++) {

    for (j = 0; j < 3; j++) {

      // OLD CODE: network[i][j] >>= netbiasshift; 

      // Fix based on bug report by Juergen Weigert jw@suse.de

      temp = (network[i][j] + (1 << (netbiasshift - 1))) >> netbiasshift;

      if (temp > 255) temp = 255;

      network[i][j] = temp;

    }

    network[i][3] = i;      // record colour no 

  }

}

//////////////////////////////////////////////////////////////////////////////////

// Insertion sort of network and building of netindex[0..255] (to do after unbias)

// -------------------------------------------------------------------------------

void NNQuantizer::inxbuild() {

  int i,j,smallpos,smallval;

  int *p,*q;

  int previouscol,startpos;

  previouscol = 0;

  startpos = 0;

  for (i = 0; i < netsize; i++) {

    p = network[i];

    smallpos = i;

    smallval = p[FI_RGBA_GREEN];      // index on g

    // find smallest in i..netsize-1

    for (j = i+1; j < netsize; j++) {

      q = network[j];

      if (q[FI_RGBA_GREEN] < smallval) { // index on g

        smallpos = j;

        smallval = q[FI_RGBA_GREEN];  // index on g

      }

    }

    q = network[smallpos];

    // swap p (i) and q (smallpos) entries

    if (i != smallpos) {

      j = q[FI_RGBA_BLUE];  q[FI_RGBA_BLUE]  = p[FI_RGBA_BLUE];  p[FI_RGBA_BLUE]  = j;

      j = q[FI_RGBA_GREEN]; q[FI_RGBA_GREEN] = p[FI_RGBA_GREEN]; p[FI_RGBA_GREEN] = j;

      j = q[FI_RGBA_RED];   q[FI_RGBA_RED]   = p[FI_RGBA_RED];   p[FI_RGBA_RED]   = j;

      j = q[3];   q[3] = p[3];   p[3] = j;

    }

    // smallval entry is now in position i

    if (smallval != previouscol) {

      netindex[previouscol] = (startpos+i)>>1;

      for (j = previouscol+1; j < smallval; j++)

        netindex[j] = i;

      previouscol = smallval;

      startpos = i;

    }

  }

  netindex[previouscol] = (startpos+maxnetpos)>>1;

  for (j = previouscol+1; j < 256; j++)

    netindex[j] = maxnetpos; // really 256

}

///////////////////////////////////////////////////////////////////////////////

// Search for BGR values 0..255 (after net is unbiased) and return colour index

// ----------------------------------------------------------------------------

int NNQuantizer::inxsearch(int b, int g, int r) {

  int i, j, dist, a, bestd;

  int *p;

  int best;

  bestd = 1000;   // biggest possible dist is 256*3

  best = -1;

  i = netindex[g];  // index on g

  j = i-1;      // start at netindex[g] and work outwards

  while ((i < netsize) || (j >= 0)) {

    if (i < netsize) {

      p = network[i];

      dist = p[FI_RGBA_GREEN] - g;        // inx key

      if (dist >= bestd)

        i = netsize; // stop iter

      else {

        i++;

        if (dist < 0)

          dist = -dist;

        a = p[FI_RGBA_BLUE] - b;

        if (a < 0)

          a = -a;

        dist += a;

        if (dist < bestd) {

          a = p[FI_RGBA_RED] - r;

          if (a<0)

            a = -a;

          dist += a;

          if (dist < bestd) {

            bestd = dist;

            best = p[3];

          }

        }

      }

    }

    if (j >= 0) {

      p = network[j];

      dist = g - p[FI_RGBA_GREEN];      // inx key - reverse dif

      if (dist >= bestd)

        j = -1; // stop iter

      else {

        j--;

        if (dist < 0)

          dist = -dist;

        a = p[FI_RGBA_BLUE] - b;

        if (a<0)

          a = -a;

        dist += a;

        if (dist < bestd) {

          a = p[FI_RGBA_RED] - r;

          if (a<0)

            a = -a;

          dist += a;

          if (dist < bestd) {

            bestd = dist;

            best = p[3];

          }

        }

      }

    }

  }

  return best;

}

///////////////////////////////

// Search for biased BGR values

// ----------------------------

int NNQuantizer::contest(int b, int g, int r) {

  // finds closest neuron (min dist) and updates freq

  // finds best neuron (min dist-bias) and returns position

  // for frequently chosen neurons, freq[i] is high and bias[i] is negative

  // bias[i] = gamma*((1/netsize)-freq[i])

  int i,dist,a,biasdist,betafreq;

  int bestpos,bestbiaspos,bestd,bestbiasd;

  int *p,*f, *n;

  bestd = ~(((int) 1)<<31);

  bestbiasd = bestd;

  bestpos = -1;

  bestbiaspos = bestpos;

  p = bias;

  f = freq;

  for (i = 0; i < netsize; i++) {

    n = network[i];

    dist = n[FI_RGBA_BLUE] - b;

    if (dist < 0)

      dist = -dist;

    a = n[FI_RGBA_GREEN] - g;

    if (a < 0)

      a = -a;

    dist += a;

    a = n[FI_RGBA_RED] - r;

    if (a < 0)

      a = -a;

    dist += a;

    if (dist < bestd) {

      bestd = dist;

      bestpos = i;

    }

    biasdist = dist - ((*p)>>(intbiasshift-netbiasshift));

    if (biasdist < bestbiasd) {

      bestbiasd = biasdist;

      bestbiaspos = i;

    }

    betafreq = (*f >> betashift);

    *f++ -= betafreq;

    *p++ += (betafreq << gammashift);

  }

  freq[bestpos] += beta;

  bias[bestpos] -= betagamma;

  return bestbiaspos;

}

///////////////////////////////////////////////////////

// Move neuron i towards biased (b,g,r) by factor alpha

// ---------------------------------------------------- 

void NNQuantizer::altersingle(int alpha, int i, int b, int g, int r) {

  int *n;

  n = network[i];       // alter hit neuron

  n[FI_RGBA_BLUE]  -= (alpha * (n[FI_RGBA_BLUE]  - b)) / initalpha;

  n[FI_RGBA_GREEN] -= (alpha * (n[FI_RGBA_GREEN] - g)) / initalpha;

  n[FI_RGBA_RED]   -= (alpha * (n[FI_RGBA_RED]   - r)) / initalpha;

}

////////////////////////////////////////////////////////////////////////////////////

// Move adjacent neurons by precomputed alpha*(1-((i-j)^2/[r]^2)) in radpower[|i-j|]

// ---------------------------------------------------------------------------------

void NNQuantizer::alterneigh(int rad, int i, int b, int g, int r) {

  int j, k, lo, hi, a;

  int *p, *q;

  lo = i - rad;   if (lo < -1) lo = -1;

  hi = i + rad;   if (hi > netsize) hi = netsize;

  j = i+1;

  k = i-1;

  q = radpower;

  while ((j < hi) || (k > lo)) {

    a = (*(++q));

    if (j < hi) {

      p = network[j];

      p[FI_RGBA_BLUE]  -= (a * (p[FI_RGBA_BLUE] - b)) / alpharadbias;

      p[FI_RGBA_GREEN] -= (a * (p[FI_RGBA_GREEN] - g)) / alpharadbias;

      p[FI_RGBA_RED]   -= (a * (p[FI_RGBA_RED] - r)) / alpharadbias;

      j++;

    }

    if (k > lo) {

      p = network[k];

      p[FI_RGBA_BLUE]  -= (a * (p[FI_RGBA_BLUE] - b)) / alpharadbias;

      p[FI_RGBA_GREEN] -= (a * (p[FI_RGBA_GREEN] - g)) / alpharadbias;

      p[FI_RGBA_RED]   -= (a * (p[FI_RGBA_RED] - r)) / alpharadbias;

      k--;

    }

  }

}

/////////////////////

// Main Learning Loop

// ------------------

/**

 Get a pixel sample at position pos. Handle 4-byte boundary alignment.

 @param pos pixel position in a WxHx3 pixel buffer

 @param b blue pixel component

 @param g green pixel component

 @param r red pixel component

*/

void NNQuantizer::getSample(long pos, int *b, int *g, int *r) {

  // get equivalent pixel coordinates 

  // - assume it's a 24-bit image -

  int x = pos % img_line;

  int y = pos / img_line;

  BYTE *bits = FreeImage_GetScanLine(dib_ptr, y) + x;

  *b = bits[FI_RGBA_BLUE] << netbiasshift;

  *g = bits[FI_RGBA_GREEN] << netbiasshift;

  *r = bits[FI_RGBA_RED] << netbiasshift;

}

void NNQuantizer::learn(int sampling_factor) {

  int i, j, b, g, r;

  int radius, rad, alpha, step, delta, samplepixels;

  int alphadec; // biased by 10 bits

  long pos, lengthcount;

  // image size as viewed by the scan algorithm

  lengthcount = img_width * img_height * 3;

  // number of samples used for the learning phase

  samplepixels = lengthcount / (3 * sampling_factor);

  // decrease learning rate after delta pixel presentations

  delta = samplepixels / ncycles;

  if(delta == 0) {

    // avoid a 'divide by zero' error with very small images

    delta = 1;

  }

  // initialize learning parameters

  alphadec = 30 + ((sampling_factor - 1) / 3);

  alpha = initalpha;

  radius = initradius;

  rad = radius >> radiusbiasshift;

  if (rad <= 1) rad = 0;

  for (i = 0; i < rad; i++) 

    radpower[i] = alpha*(((rad*rad - i*i)*radbias)/(rad*rad));

  // initialize pseudo-random scan

  if ((lengthcount % prime1) != 0)

    step = 3*prime1;

  else {

    if ((lengthcount % prime2) != 0)

      step = 3*prime2;

    else {

      if ((lengthcount % prime3) != 0) 

        step = 3*prime3;

      else

        step = 3*prime4;

    }

  }

  i = 0;    // iteration

  pos = 0;  // pixel position

  while (i < samplepixels) {

    // get next learning sample

    getSample(pos, &b, &g, &r);

    // find winning neuron

    j = contest(b, g, r);

    // alter winner

    altersingle(alpha, j, b, g, r);

    // alter neighbours 

    if (rad) alterneigh(rad, j, b, g, r);

    // next sample

    pos += step;

    while (pos >= lengthcount) pos -= lengthcount;

    i++;

    if (i % delta == 0) { 

      // decrease learning rate and also the neighborhood

      alpha -= alpha / alphadec;

      radius -= radius / radiusdec;

      rad = radius >> radiusbiasshift;

      if (rad <= 1) rad = 0;

      for (j = 0; j < rad; j++) 

        radpower[j] = alpha * (((rad*rad - j*j) * radbias) / (rad*rad));

    }

  }

}

//////////////

// Quantizer

// -----------

FIBITMAP* NNQuantizer::Quantize(FIBITMAP *dib, int ReserveSize, RGBQUAD *ReservePalette, int sampling) {

  if ((!dib) || (FreeImage_GetBPP(dib) != 24)) {

    return NULL;

  }

  // 1) Select a sampling factor in range 1..30 (input parameter 'sampling')

  //    1 => slower, 30 => faster. Default value is 1

  // 2) Get DIB parameters

  dib_ptr = dib;

  img_width  = FreeImage_GetWidth(dib); // DIB width

  img_height = FreeImage_GetHeight(dib);  // DIB height

  img_line   = FreeImage_GetLine(dib);  // DIB line length in bytes (should be equal to 3 x W)

  // For small images, adjust the sampling factor to avoid a 'divide by zero' error later 

  // (see delta in learn() routine)

  int adjust = (img_width * img_height) / ncycles;

  if(sampling >= adjust)

    sampling = 1;

  // 3) Initialize the network and apply the learning algorithm

  if( netsize > ReserveSize ) {

    netsize -= ReserveSize;

    initnet();

    learn(sampling);

    unbiasnet();

    netsize += ReserveSize;

  }

  // 3.5) Overwrite the last few palette entries with the reserved ones

    for (int i = 0; i < ReserveSize; i++) {

    network[netsize - ReserveSize + i][FI_RGBA_BLUE] = ReservePalette[i].rgbBlue;

    network[netsize - ReserveSize + i][FI_RGBA_GREEN] = ReservePalette[i].rgbGreen;

    network[netsize - ReserveSize + i][FI_RGBA_RED] = ReservePalette[i].rgbRed;

    network[netsize - ReserveSize + i][3] = netsize - ReserveSize + i;

  }

  // 4) Allocate a new 8-bit DIB

  FIBITMAP *new_dib = FreeImage_Allocate(img_width, img_height, 8);

  if (new_dib == NULL)

    return NULL;

  // 5) Write the quantized palette

  RGBQUAD *new_pal = FreeImage_GetPalette(new_dib);

    for (int j = 0; j < netsize; j++) {

    new_pal[j].rgbBlue  = (BYTE)network[j][FI_RGBA_BLUE];

    new_pal[j].rgbGreen = (BYTE)network[j][FI_RGBA_GREEN];

    new_pal[j].rgbRed = (BYTE)network[j][FI_RGBA_RED];

  }

  inxbuild();

  // 6) Write output image using inxsearch(b,g,r)

  for (WORD rows = 0; rows < img_height; rows++) {

    BYTE *new_bits = FreeImage_GetScanLine(new_dib, rows);      

    BYTE *bits = FreeImage_GetScanLine(dib_ptr, rows);

    for (WORD cols = 0; cols < img_width; cols++) {

      new_bits[cols] = (BYTE)inxsearch(bits[FI_RGBA_BLUE], bits[FI_RGBA_GREEN], bits[FI_RGBA_RED]);

      bits += 3;

    }

  }

  return (FIBITMAP*) new_dib;

}