/*
    blendjpg - average multiple jpg images with 16 bit buffer
    Copyright (C) 2005  Scott Draves <source@electricsheep.org>

    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., 675 Mass Ave, Cambridge, MA 02139, USA.
*/

#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <jpeglib.h>
#include <png.h>

int width, height;
unsigned char *jpix, *rpix;
unsigned short *apix;

int comparing = 0;
double threshold;

void accum() {
    int i, j;
    unsigned short *a;
    unsigned char *s;

    for (i = 0; i < height; i++) {
	a = apix + 3 * i * width;
	s = jpix + 4 * i * width;
	for (j = 0; j < width; j++) {
	    a[0] += s[0];
	    a[1] += s[1];
	    a[2] += s[2];
	    a += 3;
	    s += 4;
	}
    }
}

void blend(int n) {
    int i, j;
    unsigned short *a;
    unsigned char *s;

    for (i = 0; i < height; i++) {
	a = apix + 3 * i * width;
	s = jpix + 3 * i * width;
	for (j = 0; j < width; j++) {
	    /* trancate not round ? */
	    s[0] = a[0] / n;
	    s[1] = a[1] / n;
	    s[2] = a[2] / n;
	    a += 3;
	    s += 3;
	}
    }
}

void
write_jpeg(FILE *file, unsigned char *image, int width, int height) {
    struct jpeg_compress_struct info;
    struct jpeg_error_mgr jerr;
    int i;

    info.err = jpeg_std_error(&jerr);
    jpeg_create_compress(&info);
    jpeg_stdio_dest(&info, file);
    info.in_color_space = JCS_RGB;
    info.input_components = 3;
    info.image_width = width;
    info.image_height = height;
    jpeg_set_defaults(&info);
    if (getenv("jpeg")) {
	int quality = atoi(getenv("jpeg"));
	jpeg_set_quality(&info, quality, TRUE);
    }
    jpeg_start_compress(&info, TRUE);
    for (i = 0; i < height; i++) {
	JSAMPROW row_pointer[1];
	row_pointer[0] = (unsigned char *) image + (3 * width * i);
	jpeg_write_scanlines(&info, row_pointer, 1);
    }
    jpeg_finish_compress(&info);
    jpeg_destroy_compress(&info);
}

unsigned char *read_jpeg(FILE *ifp, int *width, int *height) {
    struct jpeg_decompress_struct cinfo;
    struct jpeg_error_mgr jerr;
    int num_scanlines;
    unsigned char *p, *q, *t;

    cinfo.err = jpeg_std_error(&jerr);
    jpeg_create_decompress(&cinfo);
    jpeg_stdio_src(&cinfo, ifp);
    (void) jpeg_read_header(&cinfo, TRUE);
    (void) jpeg_start_decompress(&cinfo);

    /* or image_width? */
    *width = cinfo.output_width;
    *height = cinfo.output_height;

    if (3 != cinfo.output_components) {
	printf("can only read RGB JPEG files, not %d components.\n",
	       cinfo.output_components);
	return 0;
    }
    p = q = malloc(4 * *width * *height);
    t = malloc(3 * *width);

    while (cinfo.output_scanline < cinfo.output_height) {
	unsigned char *s = t;
	int i;
	num_scanlines = jpeg_read_scanlines(&cinfo, &t, 1);
	for (i = 0; i < *width; i++) {
	    p[0] = s[0];
	    p[1] = s[1];
	    p[2] = s[2];
	    p[3] = 255;
	    s += 3;
	    p += 4;
	}
    }

    (void) jpeg_finish_decompress(&cinfo);
    jpeg_destroy_decompress(&cinfo);

    free(t);
    return q;
}

#define SIG_CHECK_SIZE 8

unsigned char *read_png(FILE *ifp, int *width, int *height) {
  unsigned char sig_buf [SIG_CHECK_SIZE];
  png_struct *png_ptr;
  png_info *info_ptr;
  png_byte **png_image = NULL;
  int linesize, x, y;
  unsigned char *p, *q;

  if (fread (sig_buf, 1, SIG_CHECK_SIZE, ifp) != SIG_CHECK_SIZE) {
    fprintf (stderr, "input file empty or too short\n");
    return 0;
  }
  if (png_sig_cmp (sig_buf, (png_size_t) 0, (png_size_t) SIG_CHECK_SIZE) != 0) {
    fprintf (stderr, "input file not a PNG file\n");
    return 0;
  }

  png_ptr = png_create_read_struct (PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  if (png_ptr == NULL) {
    fprintf (stderr, "cannot allocate LIBPNG structure\n");
    return 0;
  }
  if (setjmp(png_jmpbuf(png_ptr))) {
     if (png_image) {
	 for (y = 0 ; y < info_ptr->height ; y++)
	     free (png_image[y]);
	 free (png_image);
     }
     png_destroy_read_struct(&png_ptr, &info_ptr, (png_infopp)NULL);
     perror("reading file");
     return 0;
  }
  info_ptr = png_create_info_struct (png_ptr);
  if (info_ptr == NULL) {
    png_destroy_read_struct (&png_ptr, (png_infopp)NULL, (png_infopp)NULL);
    fprintf (stderr, "cannot allocate LIBPNG structures\n");
    return 0;
  }

  png_init_io (png_ptr, ifp);
  png_set_sig_bytes (png_ptr, SIG_CHECK_SIZE);
  png_read_info (png_ptr, info_ptr);

  if (8 != info_ptr->bit_depth) {
    fprintf(stderr, "bit depth type must be 8, not %d.\n",
	    info_ptr->bit_depth);
    return 0;
  }

  *width = info_ptr->width;
  *height = info_ptr->height;
  p = q = malloc(4 * *width * *height);
  png_image = (png_byte **)malloc (info_ptr->height * sizeof (png_byte*));

  linesize = info_ptr->width;
  switch (info_ptr->color_type) {
    case PNG_COLOR_TYPE_RGB:
      linesize *= 3;
      break;
    case PNG_COLOR_TYPE_RGBA:
      linesize *= 4;
      break;
  default:
    fprintf(stderr, "color type must be RGB or RGBA not %d.\n",
	    info_ptr->color_type);
    return 0;
  }

  for (y = 0 ; y < info_ptr->height ; y++) {
    png_image[y] = malloc (linesize);
  }
  png_read_image (png_ptr, png_image);
  png_read_end (png_ptr, info_ptr);

  for (y = 0 ; y < info_ptr->height ; y++) {
    unsigned char *s = png_image[y];
    for (x = 0 ; x < info_ptr->width ; x++) {

      switch (info_ptr->color_type) {
      case PNG_COLOR_TYPE_RGB:
	p[0] = s[0];
	p[1] = s[1];
	p[2] = s[2];
	p[3] = 255;
	s += 3;
	p += 4;
	break;
      case PNG_COLOR_TYPE_RGBA:
	p[0] = s[0];
	p[1] = s[1];
	p[2] = s[2];
	p[3] = s[3];
	s += 4;
	p += 4;
	break;
      }
    }
  }

  for (y = 0 ; y < info_ptr->height ; y++)
    free (png_image[y]);
  free (png_image);
  png_destroy_read_struct (&png_ptr, &info_ptr, (png_infopp)NULL);  

  return q;
}

void write_png(FILE *file, unsigned char *image, int width, int height) {
  png_structp  png_ptr;
  png_infop    info_ptr;
  int          i;
  unsigned char **rows = malloc(sizeof(unsigned char *) * height);

  for (i = 0; i < height; i++)
    rows[i] = image + i * width * 3;
  
  png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING,
				    NULL, NULL, NULL);
  info_ptr = png_create_info_struct(png_ptr);

  if (setjmp(png_jmpbuf(png_ptr))) {
     fclose(file);
     png_destroy_write_struct(&png_ptr, &info_ptr);
     perror("writing file");
     return;
  }
  png_init_io(png_ptr, file);

  png_set_IHDR(png_ptr, info_ptr, width, height, 8,
	       PNG_COLOR_TYPE_RGB,
	       PNG_INTERLACE_NONE,
	       PNG_COMPRESSION_TYPE_BASE,
	       PNG_FILTER_TYPE_BASE);

  png_write_info(png_ptr, info_ptr);
  png_write_image(png_ptr, (png_bytepp) rows);
  png_write_end(png_ptr, info_ptr);
  png_destroy_write_struct(&png_ptr, &info_ptr);
  free(rows);
}

void compare(char *name) {
    int i, j, k;
    double d=0, ad;
    unsigned char *s, *r;
    for (i = 0; i < height; i++) {
	s = jpix + 4 * i * width;
	r = rpix + 3 * i * width;
	for (j = 0; j < width; j++) {
	    for (k = 0; k < 3; k++) {
		int dd = (s[k] - r[k]);
		d += dd * dd;
	    }
	    s += 4;
	    r += 3;
	}
    }
    ad = d/((double)width*height);
    if (ad > threshold) {
	fprintf(stderr, "%s %g\n", name, ad);
	exit(1);
    }
}

unsigned char *open_file(char *name, int *w, int *h) {
    FILE *f = fopen(name, "rb");
    unsigned char *r;
    if (NULL == f) {
	perror(name);
	exit(1);
    }
    if (strlen(name) > 4 && !strcmp(".png", &name[strlen(name)-4])) {
      r = read_png(f, w, h);
      if (NULL == r) {
	fprintf(stderr, "could not read jpeg from %s.\n", name);
	exit(1);
      }
    } else {
      
      r = read_jpeg(f, w, h);
      if (NULL == r) {
	fprintf(stderr, "could not read jpeg from %s.\n", name);
	exit(1);
      }
    }
    return r;
}

int is_compare() {
    if (getenv("threshold")) {
	threshold = atof(getenv("threshold"));
	return 1;
    }
    return 0;
}

int main(int argc, char **argv) {
    int i;
    if (argc < 2) {
	fprintf(stderr, "usage: blendjpg source.jpg ... > result.jpg\n");
	exit(0);
    }

 again:
    jpix = open_file(argv[1], &width, &height);
    if (comparing) {
	compare(argv[1]);
    } else {
	apix = malloc(width * height * 3 * sizeof(unsigned short));
	memset(apix, 0, width * height * 3 * sizeof(unsigned short));
	accum();
    }
    for (i = 2; i < argc; i++) {
	int w2, h2;
	free(jpix);
	jpix = open_file(argv[i], &w2, &h2);
	if (w2 != width || h2 != height) {
	    fprintf(stderr, "input images sizes differ: %dx%d vs %dx%d.\n",
		    w2, h2, width, height);
	    exit(1);
	}
	if (comparing)
	    compare(argv[i]);
	else
	    accum();
    }
    if (!comparing) {
      blend(argc - 1);
      free(apix);
    }
    if (is_compare()) {
	if (comparing) exit(0);
	comparing = 1;
	rpix = jpix;
	goto again;
    } else {
      if (getenv("write_png")) {
	write_png(stdout, jpix, width, height);
      } else {
	write_jpeg(stdout, jpix, width, height);
      }
    }
    return 0;
}
