/* Copyright (c) 2012 MJ Rutter 
 * 
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License version 2
 * as published by the Free Software Foundation.
 * 
 * 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., 51 Franklin Street,
 * Fifth Floor, Boston, MA  02110-1301, USA.
 */ 

/* Apply PNG-style predictor to image, row by row.
 * So far the "Sub", "Up" and "Av" predictors are supported
 * No predictor used if found to be of no benefit
 *
 * Strictly on should compress the whole image with every combination
 * of rows predicted or unpredicted, i.e. 2^image->height tests.
 * This is impractical. This code decides to predict or not based on
 * whether it makes the row smaller when compressed in isolation.
 * This is wrong, as the previous data will affect the compression
 * ratio. It is arguably slower, and less wrong, than what libpng
 * does...
 */

#if defined(USE_ZLIB) || defined (USE_LZW)
#include <stdio.h>
#include <stdlib.h>  /* malloc */
#include <string.h>  /* memcpy */

#include "bmp2eps.h"

#ifdef USE_ZLIB
#include <zlib.h>
#endif
#ifdef USE_LZW
unsigned int to_lzw(unsigned char *buff, unsigned char *out,
		    unsigned int buff_len, int len);
#endif

static unsigned int squish(unsigned char *buff, unsigned char *out,
			   unsigned int buff_len, int len, int method);

static void predict(int method, unsigned char *source, unsigned char *dest,
		    unsigned int plen, int bytes_per_pixel);
static void p_sub(unsigned char *source, unsigned char *dest,
                  unsigned int plen, int bytes_per_pixel);
static void p_up(unsigned char *source, unsigned char *dest, unsigned int plen,
		 int bytes_per_pixel);
static void p_av(unsigned char *source, unsigned char *dest, unsigned int plen,
		 int bytes_per_pixel);
static void p_paeth(unsigned char *source, unsigned char *dest,
                    unsigned int plen, int bytes_per_pixel);

void lzw_flate(struct bitmap *image){
  unsigned char *new_data,*p,*d,*buffer,*row_buffer;
  int i,row_len,test1,test2,compress,debug,method;
  int *gains;
  unsigned int base,best;
  int predictor[5];
  unsigned long zlen;

  test1=0; /* Avoid spurious uninitialised use warning */

  debug=image->flags&DEBUG_MASK;
  compress=image->compression;
  if ((compress!=FLATE)&&(compress!=LZW)){
    fprintf(stderr,"Invalid call to lzw_flate\n");
    exit(1);
  }

#ifndef USE_ZLIB
  if (compress==FLATE){
    fprintf(stderr,"Invalid call to lzw_flate\n"
            "FLATE requested but compiled without FLATE support\n");
    exit(1);
  }
#endif

#ifndef USE_LZW
  if (compress==LZW){
    fprintf(stderr,"Invalid call to lzw_flate\n"
            "LZW requested but compiled without LZW support\n");
    exit(1);
  }
#endif

  row_len=(image->width*image->depth+7)/8;

  if (row_len*image->height!=image->data_len){
    fprintf(stderr,"Inconsistency in LZW/Flate\n");
  }

  if (compress==FLATE)
    zlen=1.01*(image->data_len+image->height)+14;
  else
    zlen=(image->data_len+image->height)*3/2+3;

  new_data=malloc(zlen);
  buffer=NULL;
  test2=0; /* Avoid spurious warning about uninitialised use */
  if (!new_data){
    fprintf(stderr,"Malloc error in LZW/Flate\n");
    exit(1);
  }

  if (debug) fprintf(stderr,"Trying %s compression",
		     compress==FLATE?"Flate":"LZW");

  /* Should we try prediction? */

  /* General view is yes for 8 bit greyscale or 24 bit colour, 
   * no for paletted */

  if ((image->flags&PREDICT)&&(image->depth>=8)&&(image->type==BMP)&&
      (image->height>1)){

    if (debug) fprintf(stderr," with prediction\n");

    buffer=malloc(zlen);
    gains=malloc(image->height*sizeof(int));
    row_buffer=malloc(row_len+1);
    if ((!buffer)||(!gains)||(!row_buffer)){
      fprintf(stderr,"Malloc error in LZW/Flate\n");
      exit(1);
    }

    p=new_data;

    /* First row can use only sub prediction */

    if (image->flags&(1<<P_SHIFT)){
      base=squish(image->data,buffer,row_len,zlen,compress);
      predict(1,image->data,new_data,image->width,image->depth/8);
      best=squish(new_data+1,buffer,row_len,zlen,compress);
      gains[0]=base-best;
    }
    else{
      predict(0,image->data,new_data,image->width,image->depth/8);
      gains[0]=0;
    }

    for(i=1;i<image->height;i++){
      p=new_data+i*(row_len+1);
      d=image->data+i*row_len;
      /* Test by compressing original row */
      base=squish(d,buffer,row_len,zlen,compress);

      if (debug>3) fprintf(stderr,"Row %5d: %9d ",i,base);

      best=1e9;
      for(method=1;method<=4;method++){

	if ((image->flags&(1<<(P_SHIFT+method-1)))==0) continue;
      
	predict(method,d,row_buffer,image->width,image->depth/8);

      /* Test by compressing row */
        test2=squish(row_buffer+1,buffer,row_len,zlen,compress);
	if (debug>3) fprintf(stderr,"%9d ",test2);
	if (test2<best){
	  memcpy(p,row_buffer,row_len+1);
	  best=test2;
	}
      }
      gains[i]=base-best;
      if (debug>3) fprintf(stderr," gain %9d\n",base-best);
    }

    /* Don't use prediction if it wastes a byte */

    for(i=0;i<5;i++)
      predictor[i]=0;
    for(i=0;i<image->height;i++){
      p=new_data+i*(row_len+1);
      if (gains[i]<0){
	*p=0;
	memcpy(p+1,image->data+i*row_len,row_len);
      }	
      predictor[(int)(*p)]++;
    }
    
    if (debug>1) {
      fprintf(stderr,"Predictors used: ");
      for(i=0;i<5;i++) fprintf(stderr,"%dx%d ",i,predictor[i]);
      fprintf(stderr,"\n");
    }
    
    /* See if we really won */
    
    /* new_data is uncompressed predicted image */
    
    test2=squish(new_data,buffer,image->data_len+image->height,zlen,compress);

  /* buffer is compressed predicted image, length test2 
   * unless we have skipped prediction, in which case buffer=NULL */
  }
  else{ /* No prediction */
    if (debug) fprintf(stderr,"\n");
  }

  /* We don't need the uncompressed predicted data any more,
   * so new_data is free to be (re-)used */

  test1=squish(image->data,new_data,image->data_len,zlen,compress);

  /* buffer is compressed predicted image, length test2
   * new_data is compressed unpredicted image, length test1 
   * test1 & test2 can be zero -- this means LZW was used,
   * and compression ineffective */

  if (debug>2){
    fprintf(stderr,"Unpredicted compressed length: %d\n",test1);
    if (buffer)
      fprintf(stderr,"Predicted compressed length:   %d\n",test2);
}

  if ((buffer)&&(debug))
    fprintf(stderr,"Prediction saves %d bytes in compressed data\n",
		     test1-test2);

  if (((test1==0)&&(test2==0))||
      ((image->data_len-test1<40)&&
       ((buffer==NULL)||(image->data_len-test2<140)))){
    if (debug) fprintf(stderr,"Neither prediction nor compression used\n");
    free(new_data);
    if (buffer) free(buffer);
    image->flags&=(~PREDICT);
    image->compression=0;
    return;
  }

  if ((buffer)&&(test1-test2>100)&&(test2>0)){
    free(image->data);
    image->data=buffer;
    image->data_len=test2;
    free(new_data);
    if (debug) fprintf(stderr,"Prediction used\n");
  }
  else{
    free(image->data);
    image->data=new_data;
    image->data_len=test1;
    if (buffer) free(buffer);
    image->flags&=(~PREDICT);
    if ((debug)&&(buffer)) fprintf(stderr,"Prediction not used\n");
  }

  if (debug) fprintf(stderr,"After compression %d bytes\n",image->data_len);

  return;
}

static void predict(int method, unsigned char *source, unsigned char *dest,
		    unsigned int plen, int bytes_per_pixel){
  switch (method){
  case 0:
    *dest++=0;
    memcpy(dest,source,plen*bytes_per_pixel);
    return;
  case 1:
    p_sub(source,dest,plen,bytes_per_pixel);
    return;
  case 2:
    p_up(source,dest,plen,bytes_per_pixel);
    return;
  case 3:
    p_av(source,dest,plen,bytes_per_pixel);
    return;
  case 4:
    p_paeth(source,dest,plen,bytes_per_pixel);
    return;
  default:
    fprintf(stderr,"Invalid predictor in call to predict\n");
    exit(1);
  }
}

static void p_up(unsigned char *source, unsigned char *dest, unsigned int plen,
	  int bytes_per_pixel){
  unsigned char *s2;
  int i,blen;

  blen=plen*bytes_per_pixel;

  *dest++=2;  /* Code for "Up" predictor */
  s2=source-blen;

  for(i=0;i<blen;i++) *dest++=*source++-*s2++;

}

static void p_sub(unsigned char *source, unsigned char *dest, unsigned int plen,
	  int bytes_per_pixel){
  int i,blen;


  if ((bytes_per_pixel!=1)&&(bytes_per_pixel!=3)){
    fprintf(stderr,"Invalid call to p_sub, bytes_per_pixel=%d\n",
	    bytes_per_pixel);
    exit(1);
  }

  *dest++=1;  /* Code for "Sub" predictor */
  blen=plen*bytes_per_pixel;
  if (bytes_per_pixel==1){
    *dest++=*source;
    for(i=1;i<plen;i++) {
      *dest++=*(source+1)-*source;
      source++;
    }
  }
  else if (bytes_per_pixel==3){
    *dest++=*source;
    *dest++=*(source+1);
    *dest++=*(source+2);
    for(i=1;i<blen;i++) {
      *dest++=*(source+3)-*source;
      source++;
    }
  }

}

static void p_av(unsigned char *source, unsigned char *dest, unsigned int plen,
	  int bytes_per_pixel){
  int i,blen,av;
  unsigned char *s2;

  if ((bytes_per_pixel!=1)&&(bytes_per_pixel!=3)){
    fprintf(stderr,"Invalid call to p_av, bytes_per_pixel=%d\n",
	    bytes_per_pixel);
    exit(1);
  }

  *dest++=3;  /* Code for "Av" predictor */
  blen=plen*bytes_per_pixel;
  s2=source-blen;

  if (bytes_per_pixel==1){
    *dest++=*source-*s2++/2;
    for(i=1;i<plen;i++) {
      av=((int)(*source)+(int)(*s2++))>>1;
      *dest++=*(source+1)-av;
      source++;
    }
  }
  else if (bytes_per_pixel==3){
    *dest++=*source-*s2++/2;
    *dest++=*(source+1)-*s2++/2;
    *dest++=*(source+2)-*s2++/2;
    for(i=1;i<blen;i++) {
      av=((int)(*source)+(int)(*s2++))>>1;
      *dest++=*(source+3)-av;
      source++;
    }
  }

}

static void p_paeth(unsigned char *source, unsigned char *dest, unsigned int plen,
	  int bytes_per_pixel){
  int i,blen,a,b,c,p,x;
  unsigned char *s2;

  if ((bytes_per_pixel!=1)&&(bytes_per_pixel!=3)){
    fprintf(stderr,"Invalid call to p_av, bytes_per_pixel=%d\n",
	    bytes_per_pixel);
    exit(1);
  }

  *dest++=4;  /* Code for "Paeth" predictor */
  blen=plen*bytes_per_pixel;
  s2=source-blen;

  if (bytes_per_pixel==1){
    *dest++=*source-*s2++;
    for(i=1;i<plen;i++) {
      a=*source;
      b=*s2;
      c=*(s2-1);
      p=a+b-c;
      if ((abs(p-a)<=abs(p-b))&&(abs(p-a)<=abs(p-c))) x=a;
      else if (abs(p-b)<=abs(p-c)) x=b;
      else x=c;
      *dest++=*(source+1)-x;
      source++;
      s2++;
    }
  }
  else if (bytes_per_pixel==3){
    *dest++=*source-*s2++;
    *dest++=*(source+1)-*s2++;
    *dest++=*(source+2)-*s2++;
    for(i=1;i<blen;i++) {
      a=*source;
      b=*s2;
      c=*(s2-3);
      p=a+b-c;
      if ((abs(p-a)<=abs(p-b))&&(abs(p-a)<=abs(p-c))) x=a;
      else if (abs(p-b)<=abs(p-c)) x=b;
      else x=c;
      *dest++=*(source+3)-x;
      source++;
      s2++;
    }
  }

}

static unsigned int squish(unsigned char *data, unsigned char *out,
			   unsigned int data_len, int out_len, int method){

  int j;

#ifdef USE_ZLIB
  uLongf zlen;
#endif

    if (method==FLATE){
#ifdef USE_ZLIB
      zlen=out_len;
      j=compress2(out,&zlen,data,data_len,9);
      if(j){
	fprintf(stderr,"ZLIB returned error code %d\n",j);
	exit(1);
      }
      return zlen;
#else
      fprintf(stderr,"Flate compression requested, "
	      "but compiled without ZLIB\n");
      exit(1);
#endif
    }
    else{
#ifdef USE_LZW
      return to_lzw(data,out,out_len,data_len);
#else
      fprintf(stderr,"LZW compression requested, "
	      "but compiled without LZW\n");
      exit(1);
#endif
    }

}


#endif