#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "lodepng.h"

typedef double real;
typedef unsigned char byte;

real variance(int N, real* a) ;

real mean(int N, real* a) {
	int i; real s = 0;	
	for (i = 0; i < N; i++) s += a[i];
	return s/i;
}

real variance(int N, real* a) {
	int i; real ss = 0, s = 0;
	for (i = 0; i < N; i++) { ss += a[i]*a[i]; s+= a[i]; }
	return (ss/i) - (s/i)*(s/i);
}

real covariance(int N, real* a, real* b) {
	int i; real sab = 0, sa = 0, sb = 0;	
	for (i = 0; i < N; i++) { sab += a[i]*b[i]; sa += a[i]; sb += b[i]; }
	return (sab/i) - (sa/i)*(sb/i);
}

real ssim(int N, real* a, real* b) {
	real ma   = mean(N, a), mb = mean(N, b);
	real va   = variance(N, a), vb = variance(N, b);
	real cv   = covariance(N, a, b);
	real L    = 255.0; /* Pixel range */
	real c1   = (0.01*L)*(0.01*L), c2 = (0.03*L)*(0.03*L);
	
	real ssim =
		( (2*ma*mb + c1)*(2*cv + c2) ) /
		( (ma*ma + mb*mb + c1)*(va + vb + c2) )
	;
	return ssim;
}

real luma(byte r, byte g, byte b) {
	return 0.2126*r + 0.7152*g + 0.0722*b; /* R.709 */
}

void subarray(
	real* dest, real* src, int srcwidth,
	int xoff, int yoff, int xsize, int ysize
) {
	int b, x, y;
	for (b = 0, y = yoff; y < yoff + ysize; y++) {
		for (x = xoff; x < xoff + xsize; x++, b++) {
			dest[b] = src[y*srcwidth + x];
		}
	}
}

int main(int argc, char* argv[]) {
	byte *imgA, *imgB;
	unsigned int widthA, widthB, heightA, heightB;
	unsigned int i, x, y, b, area;
	
	// Load images
	
	if (argc != 3) {
		puts("USAGE: ssim original.png compressed.png");
		return 1;
	}	
	if ((i = LodePNG_decode32_file(&imgA, &widthA, &heightA, argv[1]))) {
		printf("FATAL: Error loading \"%s\": %s\n",
			argv[1], LodePNG_error_text(i));
		return 1;
	}
	if ((i = LodePNG_decode32_file(&imgB, &widthB, &heightB, argv[2]))) {
		printf("FATAL: Error loading \"%s\": %s\n",
			argv[2], LodePNG_error_text(i));
		return 1;
	}
	if (widthA != widthB || heightA != heightB) {
		puts("FATAL: Images must have the same dimension.");
		return 1;
	}
	if (widthA % 8 || heightA % 8) {
		puts("FATAL: Images must have dimensions divisible by 8.");
		return 1;
	}
	
	// Make space
	
	area = widthA * heightA;
	real *lumA   = malloc(sizeof(real) * area);
	real *lumB   = malloc(sizeof(real) * area);
	real *ssvals = malloc(sizeof(real) * (area/64) );
	
	// Take luminosity values
	
	for (b = 0; b < widthA*heightA; b++) {
		lumA[b] = luma(imgA[4*b + 0], imgA[4*b + 1], imgA[4*b + 2]);
		lumB[b] = luma(imgB[4*b + 0], imgB[4*b + 1], imgB[4*b + 2]);
	}
	
	// Take SSIM for each block
	
	for (b = 0, x = 0; x < widthA; x+= 8) {
		for (y = 0; y < heightA; y+=8, b++) {
			real blockA[64] = {0}, blockB[64] = {0};
			subarray(blockA, lumA, widthA, x, y, 8, 8);
			subarray(blockB, lumB, widthA, x, y, 8, 8);
			ssvals[b] = ssim(64, blockA, blockB);
		}
	}	
	
	// Average SSIM value
	
	printf( "%2.8f\n", mean(area / 64, ssvals) );
	
	// Tidy up
	
	free(lumA);
	free(lumB);
	free(ssvals);
	
	free(imgA);
	free(imgB);
	
	return 0;
}