#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fftw3.h>
#include <math.h>

int main(int argc, char **argv) {
    fftw_complex *in, *out;
    fftw_plan p;
    FILE *fin, *fout;
    int pnm_ver, x_dim, y_dim, max_val;
    int i, j;
    double max_out;

    if (argc != 1 && argc != 2 && argc != 3) {
        fprintf(stderr, "Usage: pnm-to-fft [IN.PNM [OUT.PNM]]\n");
        exit(1);
    }

    if (argc > 1 && strcmp(argv[1], "-")) {
        fin = fopen(argv[1], "rb");
        if (fin == 0) {
            fprintf(stderr, "Can't open %s\n", argv[1]);
            exit(1);
        }
    } else {
        fin = stdin;
    }
    if (fscanf(fin, "P%d%d%d", &pnm_ver, &x_dim, &y_dim) != 3) {
        fprintf(stderr, "Bad header\n");
        exit(1);
    }
    if (pnm_ver < 1 || pnm_ver > 6 || x_dim <= 0 || x_dim > 1024 || y_dim <= 0 || y_dim > 1024) {
        fprintf(stderr, "Header parameters are out of range\n");
        exit(1);
    }
    if (pnm_ver == 2 || pnm_ver == 3 || pnm_ver == 5 || pnm_ver == 6) {
        if (fscanf(fin, "%d", &max_val) != 1 ||  max_val <= 0 || max_val > 255) {
            fprintf(stderr, "Bad maximum value\n");
            exit(1);
        }
    }
    in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * x_dim * y_dim);
    if (in == 0) {
        fprintf(stderr, "Cannot allocate %lu bytes for the source data\n", (unsigned long)(sizeof(fftw_complex) * x_dim * y_dim));
        exit(1);
    }
    out = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * x_dim * y_dim);
    if (out == 0) {
        fprintf(stderr, "Cannot allocate %lu bytes for the destination data\n", (unsigned long)(sizeof(fftw_complex) * x_dim * y_dim));
        exit(1);
    }
    p = fftw_plan_dft_2d(x_dim, y_dim, in, out, FFTW_FORWARD, FFTW_MEASURE);
    if (pnm_ver == 1) {
        for(j=0; j<y_dim; j++) {
            for(i=0; i<x_dim; i++) {
                int c;
                if (fscanf(fin, "%d", &c) != 1) {
                    fprintf(stderr, "Error reading input file at i=%d, j=%d\n", i, j);
                    exit(1);
                }
                in[j*x_dim + i][0] = c;
                in[j*x_dim + i][1] = 0;
            }
        }
    } else if (pnm_ver == 2) {
        for(j = 0; j<y_dim; j++) {
            for(i = 0; i< x_dim; i++) {
                int c;
                if (fscanf(fin, "%d", &c) != 1) {
                    fprintf(stderr, "Error reading input file\n");
                     exit(1);
                }
                in[j*x_dim + i][0] = (double)c / max_val;
                in[j*x_dim + i][1] = 0;
            }
        }
    } else if (pnm_ver == 4) {
        int pack_len = (x_dim + 7)/8;
        char *buf = (char *)malloc(pack_len);
        if (!buf) {
            fprintf(stderr,"malloc(%d) failed\n", pack_len);
            exit(1);
        }
        for(j = 0; j<y_dim; j++) {
            if (fread(buf, 1, pack_len, fin) != pack_len) {
                fprintf(stderr, "Error reading input file\n");
                exit(1);
            }
            for(i = 0; i< x_dim; i++) {
                in[j*x_dim + i][0] = buf[i/8] & (0x80 >> (i%8)) ? 1 : 0;
                in[j*x_dim + i][1] = 0;
            }
        }
        free(buf);
    } else if (pnm_ver == 5 ) {
        for(j = 0; j<y_dim; j++) {
            for(i = 0; i< x_dim; i++) {
                int c = fgetc(fin);
                if (c == EOF) {
                    fprintf(stderr, "Error reading input file\n");
                     exit(1);
                }
                in[j*x_dim + i][0] = (double)c / max_val;
                in[j*x_dim + i][1] = 0;
            }
        }
    } else {
        fprintf(stderr,"Unsupported PNM type = %d\n", pnm_ver);
        exit(1);
    }
    if (fin != stdin)
        fclose(fin);

    fftw_execute(p);
    fftw_destroy_plan(p);

    if (argc > 2 && strcmp(argv[2], "-")) {
        fout = fopen(argv[2], "wb");
        if (fout == 0) {
            fprintf(stderr, "Can't open %s\n", argv[2]);
            exit(1);
        }
    } else {
        fout = stdout;
    }
    fprintf(fout, "P2\n%d %d\n255\n", x_dim, y_dim);

    
    out[0][0] = 0;
    out[0][1] = 0;
    max_out = 0;
    for(j=0; j<y_dim; j++) {
        for(i=0; i<x_dim; i++) {
            double h = out[j*x_dim + i][0] * out[j*x_dim + i][0] + 
                       out[j*x_dim + i][1] * out[j*x_dim + i][1];
            if (max_out < h)
                max_out = h;
        }
    }
    /* Shuffle */
    for(j=0; j<y_dim/2; j++) {
        for(i=0; i<x_dim/2; i++) {
           int p00 = j*x_dim + i;
           int p01 = j*x_dim + i + x_dim/2;
           int p10 = (j + y_dim/2)*x_dim + i;
           int p11 = (j + y_dim/2)*x_dim + i + x_dim/2;

           fftw_complex u00 = {out[p00][0], out[p00][1]}; 
           fftw_complex u01 = {out[p01][0], out[p01][1]}; 
           fftw_complex u10 = {out[p10][0], out[p10][1]}; 
           fftw_complex u11 = {out[p11][0], out[p11][1]}; 
           out[p00][0] = u11[0];
           out[p00][1] = u11[1];
           out[p01][0] = u10[0];
           out[p01][1] = u10[1];
           out[p10][0] = u01[0];
           out[p10][1] = u01[1];
           out[p11][0] = u00[0];
           out[p11][1] = u00[1];
        }
    }

    for(j=0; j<y_dim; j++) {
        for(i=0; i<x_dim; i++) {
            int c = (int)((out[j*x_dim + i][0]*out[j*x_dim + i][0] +
                           out[j*x_dim + i][1]*out[j*x_dim + i][1]) * 255 / max_out);
            fprintf(fout, "%d ", c);
        }
        fprintf(fout, "\n");
    }
    if (fout != stdout)
        fclose(fout);

    fftw_free(in); 
    fftw_free(out);
    return 0;
}