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

#define max(a,b) ((a)>(b)?(a):(b))
#define min(x,y) ((x)<(y)?(x):(y))
#define MARGIN 32

typedef enum {
    kNone, kOnce, kSearch
} Mode;

typedef struct tag_Env {
    fftw_complex *in, *out;
    fftw_plan plan;
    float *buf, *error1, *error2;
    float *mask;
    float mask_energy;
    float total_energy;
    float max_energy;
    int x_dim, y_dim;
    int level, from, to;
    int buf_size, buf_width;
    int batch; // suppress tracing    
} Env;

static int gcf(int a, int b) {
    while(b) {
       int c = a % b;
       a = b;
       b = c;
    }
    return a;
}

static int gcf4(int i, int j, int k, int l) {
    return gcf(gcf(gcf(i,j),k),l);
}

void halftone(float *buf, float *err1, float *err2, float level, int x, int y, int k0, int k1, int k2, int k3) {
	float n=(float)(k0 + k1 + k2 + k3), err;
	int i, j;
	for(i=0;i < y; i += 2) {
		err2[0] = 0;
		for(j = 0; j < x; j++) {
			buf[i*x + j]=(float)(level + err1[j] > 0.5);
			err=level - buf[i*x + j] + err1[j];
			err1[j+1] +=(k0/n)*err;
			err2[j-1] +=(k1/n)*err; 
			err2[j  ] +=(k2/n)*err;
			err2[j+1] = (k3/n)*err;
		}
		err1[x-1] = 0;
		for(j = x-1; j >= 0; j--) {
			buf[(i+1)*x + j]=(float)(level + err2[j] > 0.5);
			err=level - buf[(i+1)*x + j] + err2[j];
			err2[j-1] +=(k0/n)*err;
			err1[j+1] +=(k1/n)*err;
			err1[j  ] +=(k2/n)*err;
			err1[j-1]  =(k3/n)*err;
		}
	}
}

static void malloc_failed(size_t size, int line) {
    fprintf(stderr, "Cannot allocate %lu bytes at the line %d\n", (unsigned long)size, line);
    exit(1);
}        

static void write_data(const char *fname, float *p, int x, int y, int w) {
    char buf[120];
    int i, j;
    FILE *f;
    
    strcpy(buf, fname);
    strcat(buf, "-data.pnm");
    f = fopen(buf, "wb");
    if (!f) {
        fprintf(stderr, "Cannot open %s\n", buf);
        exit(1);      
    }
    fprintf(f,"P1\n%d %d\n",x,y);
    for (j = 0; j < y; j++) {
        for (i = 0; i < x; i++)
            fprintf(f, p[i + j*w] < 0.5 ? "0 " : "1 ");
        fprintf(f, "\n");
    }
    fclose(f);
}

static void copy_data(fftw_complex *in, float *buf, int x, int y, int w) {
    int i, j;
    for (j = 0; j < y; j++) {
        for (i = 0; i < x; i++) {
            in[i + x*j][0] = buf[i + w*j];
            in[i + x*j][1] = 0;
        }
    }
}

static void calc_energy(Env *e) {
    int i, j;
    double t = 0, m = 0;
    for (j = 0; j < e->y_dim; j++) {
        for (i = 0; i < e->x_dim; i++) {
            double v = e->out[i + e->x_dim*j][0] = e->out[i + e->x_dim*j][0] * e->out[i + e->x_dim*j][0] 
                                                 + e->out[i + e->x_dim*j][1] * e->out[i + e->x_dim*j][1];
            if (m < v)
                m = v;
            t += v;
        }
    }
    e->max_energy = m;
    e->total_energy = t;
}

static void write_power_spectrum(char *fname, fftw_complex *p, float max_power, int x, int y) {
    char buf[120];
    int i, j;
    FILE *f;
    
    strcpy(buf, fname);
    strcat(buf, "-fft.pnm");
    f = fopen(buf, "wb");
    if (!f) {
        fprintf(stderr, "Cannot open %s\n", buf);
        exit(1);      
    }
    fprintf(f,"P2\n%d %d\n255\n",x,y);
    for (j = 0; j < y; j++) {
        for (i = 0; i < x; i++) {
            fprintf(f, "%d ", 255-min(255,3*(int)(p[i + x*j][0]*255./max_power + 0.5)) );
        }
        fprintf(f, "\n");
    }
    fclose(f);
}

static void shuffle(fftw_complex *out, int x, int y) {
    int i, j;
    for(j=0; j<y/2; j++) {
        for(i=0; i<x/2; i++) {
           int p00 = j*x + i;
           int p01 = j*x + i + x/2;
           int p10 = (j + y/2)*x + i;
           int p11 = (j + y/2)*x + i + x/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];
        }
    }
}

static void usage(int rc) {
    fprintf(rc ? stdout : stderr, 
        "Usage: search-halftone search|batch level from-sum to-sum\n"
        "                       once level i j k l\n");
    exit(rc);
}

static double check_center(Env *e) {
    int rr,i,j;
    const int hx = e->x_dim / 2;
    const int hy = e->y_dim / 2;
    double f_sq = e->level/255.; // square of characteristic frequency
    double acc = 0, tot=0;
    if (f_sq > 0.5)
        f_sq = 1 - f_sq;
    rr = (int)(f_sq * hx * hx * 0.8);
    for (j=0; j < e->y_dim; j++) {
        for(i = 0; i < e->x_dim; i++) {
            if ((i - hx)*(i - hx) + (j - hy)*(j - hy) < rr)
                acc += e->out[i + e->x_dim*j][0];
            tot += e->out[i + e->x_dim*j][0];
        }
    }
    printf("Center = %lf ", tot/e->x_dim/e->y_dim/e->x_dim/e->y_dim);
    return acc/tot;
}

static float check_symmetry(Env *e) {
    const int ARGSTEP = 360;
    int r_min, r_max; // width of the band
    int i,j;
    const int hx = e->x_dim / 2;
    const int hy = e->y_dim / 2;
    double f_sq = e->level/255.; // square of characteristic frequency
    double acc = 0;
    int in_count[ARGSTEP], out_count[ARGSTEP]; // in-band and out-of-band pixel count
    float in_sum[ARGSTEP], out_sum[ARGSTEP];   // in-band and out-of-band pixel sum
    int in_count_total, out_count_total
    float in_sum_total, out_sum_total
    float in_mean, out_mean;
    

    memset(in_count, 0, sizeof(in_count));    
    memset(out_count, 0, sizeof(out_count));    
    memset(in_sum, 0, sizeof(in_sum));    
    memset(out_sum, 0, sizeof(out_sum));    
    
    if (f_sq > 0.5)
        f_sq = 1 - f_sq;
    r_min = (int)(f_sq * hx * hx * 0.9);
    r_max = (int)(f_sq * hx * hx * 1.1);
    
    for (j = 0; j < e->y_dim; j++) {
        for(i = 0; i < e->x_dim; i++) {
            int vi = i - hx;
            int vj = j - hy;
            int r_sq = vi*vi + vj*vj;
            if (r_sq >= r_min) {
                int arg = ARGSTEP/2 + (int)(atan2(vj, vi)*(ARGSTEP/2)/M_PI);
                if (r_sq <= r_max) { // in the band
                    in_count[arg] += 1;
                    in_sum[arg] += e->out[i + e->x_dim*j][0];
                } else {             // out of the band
                    out_count[arg] += 1;
                    out_sum[arg] += e->out[i + e->x_dim*j][0];
                }
            }
        }
    } 
    for (i=0; i < ARGSTEP; i++) {
            in_sum_total += in_sum[i];
            out_sum_total += out_sum[i];
            in_count_total += in_count[i];
            out_count_total += out_count[i];
    }





    printf("Symmetry = %lf ", acc);
    return acc;
}

static float check_mask(Env *e) {
    float acc = 0;
    int i, j;
    float k = e->mask_energy / e-> total_energy;
    
    for (j = 0; j < e->y_dim; j++) {
        for(i = 0; i < e->x_dim; i++) {
            float diff = e->mask[i + e->x_dim*j] - k * e->out[i + e->x_dim*j][0];
            acc += diff * diff;
        }
    }
    return sqrt(acc/e->x_dim/e->y_dim);
}

static float process(Env *e, int i, int j, int k, int l) {
    char fname[80];
    int baselen;
    float quality = 0;
 
    sprintf(fname, "%03d-%02d-%02d-%02d-%02d",e->level,i,j,k,l);
    baselen = strlen(fname);
    memset(e->error1, 0, sizeof(float)*(e->buf_width + 2));
    memset(e->error2, 0, sizeof(float)*(e->buf_width + 2));
    memset(e->buf, 0, sizeof(float)*e->buf_size);
    halftone(e->buf, e->error1 + 1, e->error2 + 1, e->level/255., e->buf_width, e->y_dim, i,j,k,l);
    halftone(e->buf, e->error1 + 1, e->error2 + 1, e->level/255., e->buf_width, e->y_dim, i,j,k,l);
    strcpy(fname + baselen, ".pnm");
    if (!e->batch)
        write_data(fname, e->buf + MARGIN, e->x_dim, e->y_dim, e->buf_width);

    // fft
    copy_data(e->in, e->buf + MARGIN, e->x_dim, e->y_dim, e->buf_width); 
    fftw_execute(e->plan);
    e->out[0][0] = e->out[0][1] = 0;
    calc_energy(e);
    shuffle(e->out, e->x_dim, e->y_dim);
    if (!e->batch)
        write_power_spectrum(fname, e->out, e->max_energy, e->x_dim, e->y_dim);

    // check
    quality += check_mask(e);
    quality += check_center(e);
    quality += check_symmetry(e);
    return quality;
}

float foo(float x, float f0){
	f0 = sqrt((f0 < 0.5) ? f0 : (1 - f0));
	if (x <= f0) {
		return exp(-200*(x-f0)*(x-f0));
	} else {
		return exp(-200*(x-f0)*(x-f0))*(1-0.03) + 0.03;
	}
}

int main(int argc, char **argv) {
    Env e;
    int i, j, k, sum;
    int kernel[4];
    Mode mode = kNone;
    e.x_dim = e.y_dim = 1024;
    e.batch = 0;
    
    if (argc < 2) 
        usage(1);
    if (!strcmp(argv[1], "search") || !strcmp(argv[1], "batch")) {
        if (argc != 5)
            usage(1);
        if (sscanf(argv[2],"%d", &e.level) != 1 || e.level < 0 || e.level > 255) {
            fprintf(stderr, "Bad argument 2: %s\n", argv[2]);
            exit(1);
        }
        if (sscanf(argv[3],"%d", &e.from) != 1 || e.from <= 0) {
            fprintf(stderr, "Bad argument 3: %s\n", argv[3]);
            exit(1);
        }
        if (sscanf(argv[4],"%d", &e.to) != 1 || e.to <= 0) {
            fprintf(stderr, "Bad argument 4: %s\n", argv[4]);
            exit(1);
        }
        if(!strcmp(argv[1], "batch"))
            e.batch = 1;
        mode = kSearch;
    } else if (!strcmp(argv[1], "once")) {
        if (argc != 7)
            usage(1);
        if (sscanf(argv[2], "%d", &e.level) != 1 || e.level < 0 || e.level > 255) {
            fprintf(stderr, "Bad argument 2: %s\n", argv[2]);
            exit(1);
        }
        for(i=0; i < 4; i++) {
            if (sscanf(argv[i + 3],"%d", kernel + i) != 1 || kernel[i] < 0) {
                fprintf(stderr, "Bad argument %d: %s\n", i + 3, argv[i + 3]);
                exit(1);
            }
        }
        mode = kOnce;
    } else {
        usage(1);
    }    
   
    e.in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * e.x_dim * e.y_dim);
    if (e.in == 0) 
        malloc_failed(sizeof(fftw_complex) * e.x_dim * e.y_dim, __LINE__ - 2);
 
    e.out = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * e.x_dim * e.y_dim);
    if (e.out == 0)
        malloc_failed(sizeof(fftw_complex) * e.x_dim * e.y_dim, __LINE__ - 2);
  
    e.plan = fftw_plan_dft_2d(e.x_dim, e.y_dim, e.in, e.out, FFTW_FORWARD, FFTW_MEASURE);

    e.buf_width = 2*MARGIN + e.x_dim;
    e.buf_size = e.buf_width * e.y_dim;
    e.buf = (float *)malloc(sizeof(float)*e.buf_size);
    if (!e.buf) 
        malloc_failed(sizeof(float)*e.buf_size, __LINE__ - 2);
  
    e.error1 = (float *)malloc(sizeof(float)*(e.buf_width + 2));
    if (!e.error1)
        malloc_failed(sizeof(float)*(e.buf_width + 2), __LINE__ - 2);
    
    e.error2 = (float *)malloc(sizeof(float)*(e.buf_width + 2));
    if (!e.error2)
        malloc_failed(sizeof(float)*(e.buf_width + 2), __LINE__ - 2);
 
     // Mask
	e.mask = (float *)malloc(sizeof(float) * e.x_dim * e.y_dim);
	e.mask_energy = 0;
    for(j = 0; j < e.y_dim; j++) {
        for(i = 0; i < e.x_dim; i++) {
            e.mask_energy += e.mask[j*e.x_dim + i] = foo(hypot((float)i/e.x_dim-0.5, (float)j/e.y_dim-0.5), e.level/255.);
        } 
    }
    if (!e.batch) {
        char buf[120];
        sprintf(buf,"%03d-mask.pnm",e.level);
        FILE *f = fopen(buf, "wb");
        if (!f) {
            fprintf(stderr, "Can't open: %s\n", argv[2]);
            exit(1);
        }
        fprintf(f, "P2\n%d %d\n255\n", e.x_dim, e.y_dim);
        for(j = 0; j < e.y_dim; j++) {
            for(i = 0; i < e.x_dim; i++) {
                fprintf(f, "%d ",(int)(255*e.mask[j*e.x_dim + i]));
            }
            fprintf(f, "\n"); 
        }
        fclose(f);
    }

    if (mode == kSearch) {
        double cost = 1e200;
        char buf[120];
        FILE *cum;
        sprintf(buf,"%03d-%02d-%02d.cum", e.level, e.from, e.to);
        
        cum = fopen(buf, "wb");
        if (!cum) {
            fprintf(stderr, "Can't open: %s\n", buf);
            exit(1);
        }
        
        for(sum = e.from; sum <= e.to; sum++) {
            for(i = 0; i <= sum; i++) {
                for(j = 0; j <= sum-i; j++) {
                    for(k = 0; k <= sum-i-j; k++) {
	                    int l = sum - i - j - k;
                        int g = gcf4(i,j,k,l);
                        if (g == 1) {
                            double c = process(&e, i, j, k, l);
                            if (c < cost) {
                                cost = c;
                                kernel[0] = i;
                                kernel[1] = j;
                                kernel[2] = k;
                                kernel[3] = l;
                            }
                            if (!e.batch)
                                printf("%d %d %d %d   %lg %lg\n", i, j, k, l, c, cost);
                            fprintf(cum, "%lf %d %d %d %d\n", c, i, j, k, l);
                        }
                    }
                }
            }
        }
        fclose(cum);
        
        printf("Cost(%d,%d,%d,%d)=%lg\n", kernel[0], kernel[1], kernel[2], kernel[3], cost);
    } else if (mode == kOnce) {
        double cost;
        cost = process(&e, kernel[0], kernel[1], kernel[2], kernel[3]);
        printf("Cost(%d,%d,%d,%d)=%lg\n", kernel[0], kernel[1], kernel[2], kernel[3], cost);
    } else {
        fprintf(stderr, "Unknown mode: %d\n", (int)mode);
        exit(1);
    }
 
    fftw_destroy_plan(e.plan);
    free(e.error2);
    free(e.error1);
    free(e.buf);
    free(e.mask);
    fftw_free(e.in); 
    fftw_free(e.out);
 
    return 0;
}