/*
 * quantize.c -- image quantizer
 * (C)Copyright 1998, 98 by Hiroshi Takekawa
 * This file is part of Enfle.
 *
 * Last Modified: Wed Aug 25 00:16:08 1999.
 * $Id: quantize.c,v 1.2 1999/08/24 17:12:39 sian Exp $
 *
 * Enfle 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.
 *
 * Enfle 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 <stdio.h>
#include <stdlib.h>
#include "enfle.h"
#include "quantize.h"

/* minimize, determine the longest edge, calculate volume and population */
static void update_box(box *bp, int hist[R_INDEX][G_INDEX][B_INDEX])
{
  int r, g, b, n;
  int rl, gl, bl;

  /* minimize box */
  for (r = bp->rmin; r <= bp->rmax; r++)
    for (g = bp->gmin; g <= bp->gmax; g++)
      for (b = bp->bmin; b <= bp->bmax; b++)
	if (hist[r][g][b]) {
	  bp->rmin = r;
	  goto RMINEND;
	}
 RMINEND:

  for (r = bp->rmax; r >= bp->rmin; r--)
    for (g = bp->gmin; g <= bp->gmax; g++)
      for (b = bp->bmin; b <= bp->bmax; b++)
	if (hist[r][g][b]) {
	  bp->rmax = r;
	  goto RMAXEND;
	}
 RMAXEND:

  for (g = bp->gmin; g <= bp->gmax; g++)
    for (r = bp->rmin; r <= bp->rmax; r++)
      for (b = bp->bmin; b <= bp->bmax; b++)
	if (hist[r][g][b]) {
	  bp->gmin = g;
	  goto GMINEND;
	}
 GMINEND:

  for (g = bp->gmax; g >= bp->gmin; g--)
    for (r = bp->rmin; r <= bp->rmax; r++)
      for (b = bp->bmin; b <= bp->bmax; b++)
	if (hist[r][g][b]) {
	  bp->gmax = g;
	  goto GMAXEND;
	}
 GMAXEND:

  for (b = bp->bmin; b <= bp->bmax; b++)
    for (r = bp->rmin; r <= bp->rmax; r++)
      for (g = bp->gmin; g <= bp->gmax; g++)
	if (hist[r][g][b]) {
	  bp->bmin = b;
	  goto BMINEND;
	}
 BMINEND:

  for (b = bp->bmax; b >= bp->bmin; b--)
    for (r = bp->rmin; r <= bp->rmax; r++)
      for (g = bp->gmin; g <= bp->gmax; g++)
	if (hist[r][g][b]) {
	  bp->bmax = b;
	  goto BMAXEND;
	}
 BMAXEND:

  rl = bp->rmax - bp->rmin + 1;
  gl = bp->gmax - bp->gmin + 1;
  bl = bp->bmax - bp->bmin + 1;

  if (rl > gl)
    if (bl > rl)
      bp->longest = 3;
    else
      bp->longest = 1;
  else
    if (bl > gl)
      bp->longest = 3;
    else
      bp->longest = 2;

  bp->volume = rl * gl * bl;

  n = 0;
  for (r = bp->rmin; r <= bp->rmax; r++)
    for (g = bp->gmin; g <= bp->gmax; g++)
      for (b = bp->bmin; b <= bp->bmax; b++)
	if (hist[r][g][b])
	  n++;

  bp->population = n;
}

static int find_splitbox_by_population(box bo[256], int n)
{
  int i, m = 0, p = 0;

  for (i = 0; i < n; i++)
    if (p < bo[i].population) {
      p = bo[i].population;
      m = i;
    }

  return m;
}

static int find_splitbox_by_volume(box bo[256], int n)
{
  int i, m = 0, v = 0;

  for (i = 0; i < n; i++)
    if (v < bo[i].volume) {
      v = bo[i].volume;
      m = i;
    }

  return m;
}

static void split_box(box bo[256], int hist[R_INDEX][G_INDEX][B_INDEX], int cnum)
{
  int c = 1, m;
  int lmin, lmax, sum1, sum2;
  int r, g, b;

  while (c < cnum) {
    /* select which box to be splitted */
    if (c < cnum / 2)
      m = find_splitbox_by_population(bo, c);
    else
      m = find_splitbox_by_volume(bo, c);

    sum1 = sum2 = 0;
    /* find half distributed line and split box */
    switch (bo[m].longest) {
    case 1:
      lmin = bo[m].rmin;
      lmax = bo[m].rmax;
      while (lmin + 1 < lmax) {
	if (sum1 <= sum2) {
	  for (g = bo[m].gmin; g <= bo[m].gmax; g++)
	    for (b = bo[m].bmin; b <= bo[m].bmax; b++)
	      sum1 += hist[lmin][g][b];
	  lmin++;
	} else {
	  for (g = bo[m].gmin; g <= bo[m].gmax; g++)
	    for (b = bo[m].bmin; b <= bo[m].bmax; b++)
	      sum2 += hist[lmax][g][b];
	  lmax--;
	}
      }
      bo[c].rmin = lmax;
      bo[c].rmax = bo[m].rmax;
      bo[c].gmin = bo[m].gmin;
      bo[c].gmax = bo[m].gmax;
      bo[c].bmin = bo[m].bmin;
      bo[c].bmax = bo[m].bmax;
      bo[m].rmax = lmin;
      break;
    case 2:
      lmin = bo[m].gmin;
      lmax = bo[m].gmax;
      while (lmin + 1 < lmax) {
	if (sum1 <= sum2) {
	  for (r = bo[m].rmin; r <= bo[m].rmax; r++)
	    for (b = bo[m].bmin; b <= bo[m].bmax; b++)
	      sum1 += hist[r][lmin][b];
	  lmin++;
	} else {
	  for (r = bo[m].rmin; r <= bo[m].rmax; r++)
	    for (b = bo[m].bmin; b <= bo[m].bmax; b++)
	      sum2 += hist[r][lmax][b];
	  lmax--;
	}
      }
      bo[c].gmin = lmax;
      bo[c].gmax = bo[m].gmax;
      bo[c].rmin = bo[m].rmin;
      bo[c].rmax = bo[m].rmax;
      bo[c].bmin = bo[m].bmin;
      bo[c].bmax = bo[m].bmax;
      bo[m].gmax = lmin;
      break;
    case 3:
      lmin = bo[m].bmin;
      lmax = bo[m].bmax;
      while (lmin + 1 < lmax) {
	if (sum1 <= sum2) {
	  for (r = bo[m].rmin; r <= bo[m].rmax; r++)
	    for (g = bo[m].gmin; g <= bo[m].gmax; g++)
	      sum1 += hist[r][g][lmin];
	  lmin++;
	} else {
	  for (r = bo[m].rmin; r <= bo[m].rmax; r++)
	    for (g = bo[m].gmin; g <= bo[m].gmax; g++)
	      sum2 += hist[r][g][lmax];
	  lmax--;
	}
      }
      bo[c].bmin = lmax;
      bo[c].bmax = bo[m].bmax;
      bo[c].rmin = bo[m].rmin;
      bo[c].rmax = bo[m].rmax;
      bo[c].gmin = bo[m].gmin;
      bo[c].gmax = bo[m].gmax;
      bo[m].bmax = lmin;
      break;
    default:
      fprintf(stderr, "Illegal longest edge number %d\n", bo[m].longest);
      exit(1);
    }

    /* recalculate volume and longest */
    update_box(&bo[m], hist);
    update_box(&bo[c], hist);
    c++;
  }
}

static void calc_average(box bo[256], int hist[R_INDEX][G_INDEX][B_INDEX], int n)
{
  int i;

  for (i = 0; i < n; i++) {
    int r, g, b, rs, gs, bs, s;

    rs = 0;
    gs = 0;
    bs = 0;
    s = 0;
    for (r = bo[i].rmin; r <= bo[i].rmax; r++)
      for (g = bo[i].gmin; g <= bo[i].gmax; g++)
	for (b = bo[i].bmin; b <= bo[i].bmax; b++) {
	  s += hist[r][g][b];
	  rs += hist[r][g][b] * r;
	  gs += hist[r][g][b] * g;
	  bs += hist[r][g][b] * b;
	}
#if 0
    if (s == 0)
      s = 1;
#endif
    bo[i].r = (rs << PRECISION_R) / s;
    bo[i].g = (gs << PRECISION_G) / s;
    bo[i].b = (bs << PRECISION_B) / s;
  }
}

static unsigned char color_quantize(box bo[256], int hist[R_INDEX][G_INDEX][B_INDEX],
				    int cnum,
				    unsigned char or, unsigned char og, unsigned char ob)
{
  int i;
  unsigned char r, g, b;

  r = or >> PRECISION_R;
  g = og >> PRECISION_G;
  b = ob >> PRECISION_B;

  /* see if color code is cached... */
  if (hist[r][g][b])
    return (unsigned char)hist[r][g][b];

  /* search color */
  for (i = 0; i < cnum; i++)
    if (bo[i].rmin <= r && r <= bo[i].rmax &&
	bo[i].gmin <= g && g <= bo[i].gmax &&
	bo[i].bmin <= b && b <= bo[i].bmax) {
      hist[r][g][b] = i;
      return (unsigned char)i;
    }

  fprintf(stderr,
	  "Boxes are illegally splitted: (%d, %d, %d) not found in any box.\n",
	  (int)r, (int)g, (int)b);
  exit(1);
}

unsigned char *quantize(unsigned char *source, unsigned char *pal, int x, int y, int depth)
{
  int i;
  box bo[256];
  int hist[R_INDEX][G_INDEX][B_INDEX];
  unsigned char r, g, b, *dist, *d, *p;
  int cnum;

  /* depth must be from 1 to 8 */
  if (depth < 1 || depth > 8)
    return NULL;
  cnum = 1 << depth;

  /* make histgram */
  memset(hist, 0, sizeof(hist));
  p = source;
  for (i = 0; i < x * y; i++) {
    r = (*p++) >> PRECISION_R;
    g = (*p++) >> PRECISION_G;
    b = (*p++) >> PRECISION_B;
    hist[r][g][b]++;
  }

  /* create first entire box */
  bo[0].rmin = bo[0].gmin = bo[0].bmin = 0;
  bo[0].rmax = R_INDEX - 1;
  bo[0].gmax = G_INDEX - 1;
  bo[0].bmax = B_INDEX - 1;

  update_box(&bo[0], hist);
#ifdef DEBUG
  fprintf(stderr, "minimized box: (%d, %d, %d) - (%d, %d, %d) v = %d, l = %d\n",
	  bo[0].rmin, bo[0].gmin, bo[0].bmin,
	  bo[0].rmax, bo[0].gmax, bo[0].bmax,
	  bo[0].volume, bo[0].longest);
#endif

  /* split box into desired number of boxes */
  split_box(bo, hist, cnum);

#ifdef DEBUG
  fprintf(stderr, "splitted\n");
#endif

#if 0
  for (i = 0; i < cnum; i++)
    fprintf(stderr, "%d (%d, %d, %d) - (%d, %d, %d), v = %d, l = %d\n",
	    i,
	    bo[i].rmin, bo[i].gmin, bo[i].bmin,
	    bo[i].rmax, bo[i].gmax, bo[i].bmax,
	    bo[i].volume, bo[i].longest);
#endif

  /* calculate represented(averaged) value */
  calc_average(bo, hist, cnum);

#ifdef DEBUG
  fprintf(stderr, "averaged\n");
#endif

  /* clear histgram .. to use color selection cache */
  memset(hist, 0, sizeof(hist));

  /* allocate destination image memory */
  if ((d = dist = malloc(x * y)) == NULL)
    return NULL;

#ifdef DEBUG
  fprintf(stderr, "memory allocated %d bytes\n", cnum * 3 + x * y);
#endif

  /* set palette */
  for (i = 0; i < cnum; i++) {
    *pal++ = bo[i].r;
    *pal++ = bo[i].g;
    *pal++ = bo[i].b;
  }

  /* quantization */
  p = source;
  for (i = 0; i < x * y; i++) {
    r = *p++;
    g = *p++;
    b = *p++;
    *d++ = color_quantize(bo, hist, cnum, r, g, b);
  }

  return dist;
}