#include <stdio.h>
#include <math.h>
#include "libjpegrest.h"

#define PI 3.141592657

/* This opens a 8x8 PGM file from argv[1], and stores
 * a 64+7x64+7 PGM file with each DCT-IDCT result, considering
 * only the coefficients from top-left until the concerning
 * DCT-IDCT result. */

static void mydct(float *in, float *out)
{
    int x,y,u,v;

    float Cu, Cv;

    for (u=0; u < 8; ++u)
        for (v=0; v < 8; ++v)
        {
            float *optr;
            float sum;
            optr = out+v*8+u;

            sum = 0.;
            for (x=0; x < 8; ++x)
                for (y=0; y < 8; ++y) {
                    float *iptr;
                    iptr = in+y*8+x;

                    sum += *iptr
                        * cos(((float) ((2*x+1)*u) )*PI/16.)
                        * cos(((float) ((2*y+1)*v) )*PI/16.);
                }

            if (u == 0) Cu = 1. / sqrt(2); else Cu = 1.;
            if (v == 0) Cv = 1. / sqrt(2); else Cv = 1.;
            *optr = 0.25 * Cu * Cv * sum;
        }
}

static void myidct(float *in, float *out)
{
    int x,y,u,v;

    float Cu, Cv;

    for (x=0; x < 8; ++x)
        for (y=0; y < 8; ++y)
        {
            float *optr;
            float sum;
            optr = out+y*8+x;

            sum = 0.;
            for (v=0; v < 8; ++v)
                for (u=0; u < 8; ++u)
                {
                    float *iptr;
                    iptr = in+v*8+u;

                    if (u == 0) Cu = 1. / sqrt(2); else Cu = 1.;
                    if (v == 0) Cv = 1. / sqrt(2); else Cv = 1.;

                    sum += Cu * Cv * *iptr
                        * cos(((float) ((2*x+1)*u) )*PI/16.)
                        * cos(((float) ((2*y+1)*v) )*PI/16.);
                }

            *optr = 0.25 * sum;
        }
}

void zero_block(float *block)
{
    int i;
    for (i=0; i < 64; ++i)
        block[i] = 0.;
}

void prepare_subblock(float *block, int x, int y, const float *input)
{
    /* Fill the subsquare from 0..x and 0..y */
    int i,j;

    zero_block(block);

    for (j=0; j <= y; ++j)
        for (i=0; i <= x; ++i)
        {
            block[j*8+i] = input[j*8+i];
        }
}

void write_block_to_big(struct bitmap *bdest, const float *sub, int x, int y)
{
    int i,j;

    for (j=0; j < 8; ++j)
        for (i=0; i < 8; ++i)
        {
            gsl_matrix *dest = bdest->m[0];
            dest->data[(j+y*9)*dest->tda+x*9+i] = sub[j*8+i];
        }

    if (x < 7)
        for (j=0; j < 8; ++j)
        {
            gsl_matrix *dest = bdest->m[0];
            dest->data[(j+y*9)*dest->tda+x*9+8] = 255;
        }
    if (y < 7)
        for (i=0; i < 8; ++i)
        {
            gsl_matrix *dest = bdest->m[0];
            dest->data[(8+y*9)*dest->tda+x*9+i] = 255;
        }

    if (y < 7 && x < 7)
    {
        gsl_matrix *dest = bdest->m[0];
        dest->data[(8+y*9)*dest->tda+x*9+8] = 255;
    }
}

int main(int argn, char **argv)
{
    struct bitmap bmap;
    float input[64];
    float input_test[64];
    float input_dct[64];
    float output_dct[64];
    float output[64];
    float quantize[64] = {
        20, 14, 13, 20, 30, 50, 64, 76, 
        15, 15, 18, 24, 33, 73, 75, 69, 
        18, 16, 20, 30, 50, 71, 86, 70, 
        18, 21, 28, 36, 64, 109, 100, 78, 
        23, 28, 46, 70, 85, 136, 129, 96, 
        30, 44, 69, 80, 101, 130, 141, 115, 
        61, 80, 98, 109, 129, 151, 150, 126, 
        90, 115, 119, 123, 140, 125, 129, 124 };


    struct bitmap obmap;
    int i,j,k;

    bitmap_pam_init(&argn, argv);

    bitmap_filename_pam_get(argv[1], &bmap);

    if (bmap.nplanes != 1 || bmap.width != 8 || bmap.height != 8)
    {
        fprintf(stderr, "Error: The file should be 8x8 PGM\n");
        return -1;
    }

    k=0;
    for(j=0;j < 8; ++j)
        for(i=0;i < 8; ++i)
        {
            gsl_matrix *mat = bmap.m[0];
            input[k++] = mat->data[j*mat->tda + i];
        }

    mydct(input, input_dct);

    {
        FILE *odct, *ifile, *qtable, *qfile;
        odct = fopen("dct.txt", "wt");
        ifile = fopen("ifile.txt", "wt");
        qtable = fopen("qtable.txt", "wt");
        qfile = fopen("qfile.txt", "wt");
        for(j=0;j < 8; ++j)
        {
            fprintf(odct, "<tr>\n");
            fprintf(ifile, "<tr>\n");
            fprintf(qtable, "<tr>\n");
            fprintf(qfile, "<tr>\n");
            for(i=0;i < 8; ++i)
            {
                fprintf(odct, "<td>%.2f\n", input_dct[j*8+i]);
                fprintf(ifile, "<td>%.0f\n", input[j*8+i]);
                fprintf(qtable, "<td>%.0f\n", quantize[j*8+i]);
                fprintf(qfile, "<td>%i\n", lrint(input_dct[j*8+i]/quantize[j*8+i]));
            }
        }
        fclose(odct);
        fclose(ifile);
        fclose(qtable);
        fclose(qfile);
    }

    myidct(input_dct, input_test);

    obmap = bitmap_alloc(1, 64+7, 64+7);

    for(j=0;j < 8; ++j)
        for(i=0;i < 8; ++i)
        {
            prepare_subblock(output_dct,i,j,input_dct);
            myidct(output_dct, output);
            write_block_to_big(&obmap, output, i, j);
        }

    bitmap_filename_pam_put(argv[2], &obmap);

    bitmap_free(obmap);
    return 0;
}