/*
 * Copyright (c) 2003 Michael Niedermayer <michaelni@gmx.at>
 *
 * This file is part of FFmpeg.
 *
 * FFmpeg is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * FFmpeg 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
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with FFmpeg; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <inttypes.h>
#include <assert.h>

void vp8_ssim_parms_8x8_c(unsigned char *s,
                          int sp,
                          unsigned char *r,
                          int rp,
                          unsigned long *sum_s,
                          unsigned long *sum_r,
                          unsigned long *sum_sq_s,
                          unsigned long *sum_sq_r,
                          unsigned long *sum_sxr)
{
    int i,j;
    for(i=0;i<8;i++,s+=sp,r+=rp)
    {
        for(j=0;j<8;j++)
        {
            *sum_s += s[j];
            *sum_r += r[j];
            *sum_sq_s += s[j] * s[j];
            *sum_sq_r += r[j] * r[j];
            *sum_sxr += s[j] * r[j];
        }
    }
}

const static int64_t cc1 =  26634; // (64^2*(.01*255)^2
const static int64_t cc2 = 239708; // (64^2*(.03*255)^2

static double similarity(unsigned long sum_s,
                         unsigned long sum_r,
                         unsigned long sum_sq_s,
                         unsigned long sum_sq_r,
                         unsigned long sum_sxr,
                         int count)
{
    int64_t ssim_n, ssim_d;
    int64_t c1, c2;
    
    //scale the constants by number of pixels
    c1 = (cc1*count*count)>>12;
    c2 = (cc2*count*count)>>12;
    
    ssim_n = (2*sum_s*sum_r+ c1)*((int64_t) 2*count*sum_sxr-
                                  (int64_t) 2*sum_s*sum_r+c2);
    
    ssim_d = (sum_s*sum_s +sum_r*sum_r+c1)*
    ((int64_t)count*sum_sq_s-(int64_t)sum_s*sum_s +
     (int64_t)count*sum_sq_r-(int64_t) sum_r*sum_r +c2) ;
    
    return ssim_n * 1.0 / ssim_d;
}

static double ssim_8x8(unsigned char *s,int sp, unsigned char *r,int rp)
{
    unsigned long sum_s=0,sum_r=0,sum_sq_s=0,sum_sq_r=0,sum_sxr=0;
    vp8_ssim_parms_8x8_c(s, sp, r, rp, &sum_s, &sum_r, &sum_sq_s, &sum_sq_r, &sum_sxr);
    return similarity(sum_s, sum_r, sum_sq_s, sum_sq_r, sum_sxr, 64);
}

// We are using a 8x8 moving window with starting location of each 8x8 window
// on the 4x4 pixel grid. Such arrangement allows the windows to overlap
// block boundaries to penalize blocking artifacts.
double vp8_ssim2(unsigned char *img1,
                 unsigned char *img2,
                 int stride_img1,
                 int stride_img2,
                 int width,
                 int height)
{
    int i,j;
    int samples =0;
    double ssim_total=0;
    
    // sample point start with each 4x4 location
    for(i=0; i < height-8; i+=4, img1 += stride_img1*4, img2 += stride_img2*4)
    {
        for(j=0; j < width-8; j+=4 )
        {
            double v = ssim_8x8(img1+j, stride_img1, img2+j, stride_img2);
            ssim_total += v;
            samples++;
        }
    }
    ssim_total /= samples;
    return ssim_total;
}

int main(int argc, char *argv[])
{
    FILE *f[2];
    uint8_t *buf[2];
    int size[2];
    int w, h;
    double ssim;

    f[0] = fopen(argv[1], "rb");
    f[1] = fopen(argv[2], "rb");
    sscanf(argv[3], "%dx%d", &w, &h);
    if (!f[0] || !f[1]) {
        fprintf(stderr, "Could not open input files.\n");
        return 1;
    }
    fseek(f[0], 0, SEEK_END);
    fseek(f[1], 0, SEEK_END);
    size[0] = ftell(f[0]);
    size[1] = ftell(f[1]);
    assert(size[0] == w * h);
    assert(size[1] == w * h);
    buf[0] = malloc(size[0]);
    buf[1] = malloc(size[1]);
    fseek(f[0], 0, SEEK_SET);
    fseek(f[1], 0, SEEK_SET);
    fread(buf[0], size[0], 1, f[0]);
    fread(buf[1], size[1], 1, f[1]);
    ssim = vp8_ssim2(buf[0], buf[1], w, w, w, h);
    printf("SSIM: %lf\n", ssim);

    return 0;
}