// Todo
// ====
//
// Versions
// ========
// 9. A serious bug in previous versions has been removed, leading to much 
//    better segmentations.  The probabilities were being computed in such 
//    a way that they could be noticably greater than 1.0, so 
//    the Metropolis update step was producing incorrect results.
//    Interestingly, this bug was originally hidden by a bit of defensive
//    programming.  
//    Pixels at the boundaries are now assumed to be in the background.
// 8. Random scans can now be enabled with the -r option.  
// 7. RGB probabilities are now computed with the Laplace Rule of Succession 
//    by default.  This is producing good results so far, but there is
//    an option to disable it (-R) just in case.
// 6. A bug in Metropolis-Hastings was removed.  The final result is saved
//    as a png file: ppm-seg-result.png
// 5. Metropolis-Hastings updating has been added.  This helps to avoid 
//    islands of spurious background within the brain.
// 4. Gibbs sampling from an MRF has been added.  
// 3. 3D Gaussians have been replaced by RGB histograms.
// 2. -t threshold option has been added.
// 1. Basic functionality

// #define HAVE_GDK

#ifdef HAVE_GDK
#include <gdk/gdk.h>
#include <gdk-pixbuf/gdk-pixbuf.h>
#endif

#include <ppm.h>
#include <stdlib.h>
#include <assert.h>
#include <math.h>
#include <sys/types.h>
#include <unistd.h>
#include <getopt.h>
#include <string.h>

#ifndef min
#define min(a,b) ((a)<(b)?(a):(b))
#endif

#ifndef max
#define max(a,b) ((a)>(b)?(a):(b))
#endif

typedef struct { int r,g,b; } Color;
typedef struct { double mu,sigma; } Gaussian1D;

double p_fg=0.5; 
double p_bg=0.5; 

double p_rgb_given_fg[256][256][256]; 
double p_rgb_given_bg[256][256][256]; 
double p_neis_given_fg[16]; // probs for sixteen possible neighbor configurations
double p_neis_given_bg[16]; 
double p_neis[16];
char *p_filename="p.pgm";
char *rgb_filename="foreground.ppm";
char *color_truth_filename="ground-truth.ppm";
char *class_truth_filename="ground-truth.pbm";

// options
char *filename=NULL; /* of image to segment */
double threshold=1e-2;
int num_gibbs_samples=3;
int passes_per_gibbs_sample=4;
int using_metropolis_hastings=1;
int using_ros=1; /* ros means the Laplace Rule of Succession */
int using_random_scans=0;

static void print_usage(void);
static void do_gibbs(void);
static void init_distros_from_ground_truth_images(void);
static void output_simple_seg_and_p_fg_given_colors(void);

int main(int argc, char *argv[])
{

    ppm_init(&argc, argv);
#ifdef HAVE_GDK
    gdk_init(&argc, &argv);
#endif

    /* Read commands from standard input, to load example images 
       and set up the probabilities. */
    {
        int ch;
        while ((ch=getopt(argc,argv,"hMp:rs:t:")) != -1) {
            switch(ch) {
                case 'h':
                    print_usage();
                    exit(0);
                    break;
                case 'M':
                    using_metropolis_hastings=0;
                    fprintf(stderr,"Metropolis-Hastings mode disabled.\n");
                    break;
                case 'p':
                    passes_per_gibbs_sample=atoi(optarg);
                    fprintf(stderr,"passes_per_gibbs_sample: %i\n",
                            passes_per_gibbs_sample);
                    if(passes_per_gibbs_sample<=0||passes_per_gibbs_sample>50) {
                        fprintf(stderr,
"Unreasonable value given for -p option: %s\n",optarg);
                        print_usage();
                        exit(10);
                    }
                    break;
                case 'R':
                    using_ros=0;
                    fprintf(stderr,"Disabling Rule of Succession for RGBs.\n");
                    break;
                case 'r':
                    using_random_scans=1;
                    fprintf(stderr,"Enabling random scans for Gibbs.\n");
                    break;
                case 's':
                    num_gibbs_samples=atoi(optarg);
                    fprintf(stderr,"num_gibbs_samples: %i\n",num_gibbs_samples);
                    if(num_gibbs_samples<0||num_gibbs_samples>50) {
                        fprintf(stderr,
"Unreasonable value given for -s option: %s\n",optarg);
                        print_usage();
                        exit(11);
                    }
                    break;
                case 't':
                    threshold=atof(optarg);
                    fprintf(stderr,"threshold: %f\n",threshold);
                    if(threshold<=0.0||threshold>=1.0) {
                        fprintf(stderr,"Bad threshold value: %s\n",optarg);
                        print_usage();
                        exit(8);
                    }
                    break;
            }
        }
        argc -= optind;
        argv += optind;
        filename=argv[0];
        fprintf(stderr,"filename: %s\n",filename);
    }

    if(argc==0) {
        print_usage();
        exit(9);
    }
    filename=argv[0]; argv++; argc--;

    init_distros_from_ground_truth_images();

    output_simple_seg_and_p_fg_given_colors();

    if(num_gibbs_samples>0) { 
        do_gibbs();
    }

    return 0;
}

static void print_usage(void)
{
    fprintf(stderr,
"\n"
"Welcome to ppm-blockface-seg\n"
"written in 2005 by Issac Trotts, with support from the NIH Human Brain Project.\n"
"\n"
"This program segments a color blockface image in PPM format.\n"
"The general form for how to invoke this program is\n"
"\n"
"  ppm-seg [options] image-to-segment.ppm\n"
"\n"
"For example:"
"\n"
"  ppm-seg image-to-segment.ppm\n"
"\n"
"Options:\n"
"=======\n"
" -M   : Disable Metropolis-Hastings update for Gibbs sampling.\n"
" -p # : Set # passes per Gibbs sample.  1 to 50.  Default is 2.\n"
" -R   : Disable the Laplace rule of succession for RGB probabilities.\n"
" -r   : Enable random scans for Gibbs sampling.\n"
" -s # : Set # Gibbs samples.  1 to 50.  Default is 0.\n"
" -t # : Set threshold.  0 to 1, exclusive.  Default is 1e-2."
"\n");
}

/* Load the image to segment, one line at a time.  
    Each time a line is loaded, classify each pixel in it and 
    write the results out.  This way, we can process huge images. */
static void output_simple_seg_and_p_fg_given_colors(void)
{
    int xi,yi,nx,ny,format;
    pixval maxval;
    pixel* pixelrow;
    FILE* file;
    double dr,dg,db;

    dr=1./255.;
    dg=1./255.;
    db=1./255.;
    file=fopen(filename,"rb");
    if(!file) {
        fprintf(stderr,"Could not open %s for reading.\n",filename);
        exit(1);
    }
    ppm_readppminit(file, &nx, &ny, &maxval, &format);
    pixelrow=ppm_allocrow(nx);

    FILE *p_out,*rgb_out;
    p_out=fopen(p_filename,"wb");
    rgb_out=fopen(rgb_filename,"wb");
    if(!p_out) {
        fprintf(stderr,"Could not open %s for writing.\n",p_filename);
        exit(6);
    }
    if(!rgb_out) { 
        fprintf(stderr,"Could not open %s for writing.\n",rgb_filename);
        exit(7);
    }
    fprintf(p_out,"P5\n%i %i\n255\n",nx,ny);
    fprintf(rgb_out,"P6\n%i %i\n255\n",nx,ny);

    fprintf(stderr,"Generating %s and %s\n",p_filename,rgb_filename);
    for(yi=0;yi<ny;yi++) {
        fprintf(stderr,"% 6i of % 6i rows\r", yi,ny);
        ppm_readppmrow(file, pixelrow, nx, maxval, format);
        for(xi=0;xi<nx;xi++) {
            int r,g,b;
            double p_fg_given_rgb,p_rgb_bg,p_rgb_fg,p_rgb;
            pixel pxl = pixelrow[xi];
            r=pxl.r; g=pxl.g; b=pxl.b;
            p_rgb_bg = p_bg*p_rgb_given_bg[r][g][b];
            p_rgb_fg = p_fg*p_rgb_given_fg[r][g][b];
            p_rgb = p_rgb_fg + p_rgb_bg;
            p_fg_given_rgb = p_rgb_fg / p_rgb;

            /* Output results */
            fputc(max(0,min(255,p_fg_given_rgb*255.0)), p_out);

            if(p_fg_given_rgb>=threshold) {
                fputc(pxl.r, rgb_out);
                fputc(pxl.g, rgb_out);
                fputc(pxl.b, rgb_out);
            }
            else {
                fputc(0, rgb_out);
                fputc(0, rgb_out);
                fputc(0, rgb_out);
            }
        }
    }
    fprintf(stderr,"                                               \r");
    fclose(p_out);
    fclose(rgb_out);

    fclose(file);
    ppm_freerow(pixelrow);
}

static void do_gibbs(void)
{
    int s,pass,xi,yi,nx,ny;
    pixval maxval;
    bit** bitmap;
    pixel** pixels;
    gray** opacities;

    /* Load the file to be segmented. */
    {
        FILE* file=fopen(filename,"rb");
        if(!file) {
            fprintf(stderr,"Could not open %s for reading.\n",filename);
            exit(12);
        }
        pixels=ppm_readppm(file,&nx,&ny,&maxval);
        fclose(file);
    }

    /* Allocate the bitmap and initialize it to noise. */
    {
        bitmap=pbm_allocarray(nx,ny);
        fprintf(stderr,"Initializing segmentation bitmap to noise\n");
        for(yi=0;yi<ny;yi++) {
            fprintf(stderr,"  % 6i of % 6i rows\r",yi,ny);
            for(xi=0;xi<nx;xi++) {
                /* Assume that the outer boundary of pixels is all background.
                 * All other pixels get a random initial assignment of fg 
                 * or bg. */
                bitmap[yi][xi]=
                    (yi==0||yi==ny-1||xi==0||xi==nx-1)
                    ?0
                    :(drand48()>0.5?1:0);
            }
        }
        fprintf(stderr,"                                        \r");
    }

    /* Allocate the opacity map and initialize it to zeros. */
    {
        opacities=pgm_allocarray(nx,ny);
        for(yi=0;yi<ny;yi++) {
            for(xi=0;xi<nx;xi++) {
                opacities[yi][xi]=0;
            }
        }
    }

    /* Gibbs sample. */
    {
        for(s=0;s<num_gibbs_samples;s++) {
            fprintf(stderr,"Gibbs sample %i of %i\n",s,num_gibbs_samples);
            for(pass=0;pass<passes_per_gibbs_sample;pass++) {
                int rows[ny];
                fprintf(stderr," Pass %i of %i\n",pass,passes_per_gibbs_sample);
                for(yi=0;yi<ny;yi++) { rows[yi]=yi; }
                if(using_random_scans) {
                    /* Shuffle the row order */
                    for(yi=0;yi<ny;yi++) { 
                        int i1,i2,tmp;
                        i1=1+drand48()*(ny-2); i1=min(ny-2,i1);
                        i2=1+drand48()*(ny-2); i2=min(ny-2,i2);
                        tmp=rows[i1];
                        rows[i1]=rows[i2];
                        rows[i2]=tmp; 
                    }
                }
                for(yi=1;yi<ny-1;yi++) {
                    fprintf(stderr,"  % 6i of % 6i rows\r",yi,ny);
                    int cols[nx];
                    for(xi=0;xi<nx;xi++) { cols[xi]=xi; }
                    if(using_random_scans) {
                        /* Shuffle the column order */
                        for(xi=0;xi<nx;xi++) { 
                            int i1,i2,tmp;
                            i1=1+drand48()*(nx-2); i1=min(nx-2,i1);
                            i2=1+drand48()*(nx-2); i2=min(nx-2,i2);
                            tmp=cols[i1];
                            cols[i1]=cols[i2];
                            cols[i2]=tmp; 
                        }
                    }
                    for(xi=1;xi<nx-1;xi++) {
                        int xi2,yi2;
                        pixel p; 
                        int r,g,b;
                        xi2=cols[xi];
                        yi2=rows[yi];
                        p=pixels[yi2][xi2];
                        r=p.r;
                        g=p.g;
                        b=p.b;
                        int neis=
                             (bitmap[yi2-1][xi2])
                            +(bitmap[yi2+1][xi2]<<1)
                            +(bitmap[yi2][xi2-1]<<2)
                            +(bitmap[yi2][xi2+1]<<3);
                        double p_fg_rgb_neis = 
                            p_fg*p_rgb_given_fg[r][g][b]*p_neis_given_fg[neis];
                        double p_bg = 1.0-p_fg;
                        double p_bg_rgb_neis = 
                            p_bg*p_rgb_given_bg[r][g][b]*p_neis_given_bg[neis];
                        double p_rgb_neis = p_fg_rgb_neis + p_bg_rgb_neis;
                        double p_fg_given_rgb_neis = p_fg_rgb_neis/p_rgb_neis;
                        assert(p_fg_given_rgb_neis>=0.0);
                        assert(p_fg_given_rgb_neis<=1.0);
                        if(using_metropolis_hastings) {
                            double p;
                            p=p_fg_given_rgb_neis;
                            p=bitmap[yi2][xi2]?p:1.0-p;
                            if(drand48()<(p==0.0?1.0:min(1.0,(1.0-p)/p))) {
                                /* flip */
                                bitmap[yi2][xi2]=bitmap[yi2][xi2]?0:1;
                            }
                        }
                        else {
                            bitmap[yi2][xi2]=(drand48()<p_fg_given_rgb_neis)?1:0;
                        }
                    }
                }
                fprintf(stderr,"                                         \r");
            }

            /* Output bitmap. */
            {
                char filename[80];
                FILE* f;
                snprintf(filename,sizeof(filename),"seg-%i.pbm",s);
                f=fopen(filename,"wb");
                if(!f) {
                    fprintf(stderr,"Could not open %s for writing.\n",filename);
                    continue;
                }
                pbm_writepbm(f,bitmap,nx,ny,/*forceplain*/ 0);
                fclose(f);
                fprintf(stderr,"Wrote %s\n",filename);
            }

            /* Update the opacity map */
            {
                for(yi=0;yi<ny;yi++) {
                    for(xi=0;xi<nx;xi++) {
                        opacities[yi][xi]+=bitmap[yi][xi];
                    }
                }
            }
        }
    }

    /* Normalize the opacity map */
    {
        for(yi=0;yi<ny;yi++) {
            for(xi=0;xi<nx;xi++) {
                opacities[yi][xi]=
                    (gray)((double)opacities[yi][xi]*255./(double)num_gibbs_samples);
            }
        }
    }

    /* Output a PAM image with the original image plus the alpha channel 
     * we just computed. */
    {

#ifdef HAVE_GDK
        char filename[]="ppm-seg-result.png";
        GdkPixbuf* pixbuf=gdk_pixbuf_new(GDK_COLORSPACE_RGB,/*has_alpha*/ 1, 
                /*bits_per_sample*/ 8, nx, ny);
        guchar* pbpix=gdk_pixbuf_get_pixels(pixbuf);
        for(yi=0;yi<ny;yi++) {
            fprintf(stderr,"  % 6i of % 6i rows\r",yi,ny);
            for(xi=0;xi<nx;xi++) {
                pixel p=pixels[yi][xi];
                gray a=opacities[yi][xi];
                pbpix[0+4*(xi+nx*yi)]=p.r;
                pbpix[1+4*(xi+nx*yi)]=p.g;
                pbpix[2+4*(xi+nx*yi)]=p.b;
                pbpix[3+4*(xi+nx*yi)]=  a;
            }
        }
        GError *error=NULL;
        if(!gdk_pixbuf_save(pixbuf, filename, "png", &error, NULL)) {
            fprintf(stderr,"Save failed. %s\n",error->message);
            exit(15);
        }
        fprintf(stderr,"Wrote %s\n",filename);
#else
        char filename[]="ppm-seg-result.pam";
        FILE* file=fopen(filename,"wb");
        if(!file) {
            fprintf(stderr,"Could not open %s for writing.\n",filename);
            exit(14);
        }
        fprintf(stderr,"Writing %s\n",filename);
        fprintf(file,"P7\nWIDTH %i\nHEIGHT %i\nDEPTH 4\nMAXVAL 255\n",nx,ny);
        fprintf(file,"TUPLTYPE RGB_ALPHA\nENDHDR\n");
        for(yi=0;yi<ny;yi++) {
            fprintf(stderr,"  % 6i of % 6i rows\r",yi,ny);
            for(xi=0;xi<nx;xi++) {
                pixel p=pixels[yi][xi];
                gray a=opacities[yi][xi];
                fputc(p.r,file);
                fputc(p.g,file);
                fputc(p.b,file);
                fputc(  a,file);
            }
        }
        fclose(file);
#endif
    }

    pgm_freearray(opacities,ny);
    ppm_freearray(pixels,ny);
    pbm_freearray(bitmap,ny);    
}

static void init_distros_from_ground_truth_images(void)
{
    int nfg=0,nbg=0;

    fprintf(stderr,"Initializing color histograms.\n");
    {
        int r,g,b;
        for(r=0;r<256;r++) {
            fprintf(stderr,"% 3i of 255\r",r);
            for(g=0;g<256;g++) {
                for(b=0;b<256;b++) {
                    p_rgb_given_fg[r][g][b]=using_ros?1.0:0.0;
                    p_rgb_given_bg[r][g][b]=using_ros?1.0:0.0;
                }
            }
        }
        fprintf(stderr,"             \r");
    }

    fprintf(stderr,"Initializing spatial histograms.\n");
    {
        int neis;
        for(neis=0;neis<16;neis++) {
            p_neis_given_fg[neis]=using_ros?1.0:0.0;
            p_neis_given_bg[neis]=using_ros?1.0:0.0;
        }
    }

    fprintf(stderr,"Filling color histograms from %s.\n",color_truth_filename);
    {
        pixval maxval;
        int xi,yi,nx,ny,nxb,nyb;
        FILE *colorfile,*classfile;
        colorfile=fopen(color_truth_filename,"rb");
        if(!colorfile) { 
            fprintf(stderr,"File not found: %s\n",color_truth_filename);
            exit(10);
        }
        classfile=fopen(class_truth_filename,"rb");
        if(!classfile) {
            fprintf(stderr,"Could not open %s for reading.\n",
                    class_truth_filename);
            exit(17);
        }
        pixel** expixels=ppm_readppm(colorfile,&nx,&ny,&maxval);
        bit** bitmap=pbm_readpbm(classfile,&nxb,&nyb);
        if(nxb!=nx||ny!=ny) {
            fprintf(stderr,
"Dimensions of the segmentation bitmap do not match those of the color image.\n");
            exit(16);
        }
        fclose(colorfile);
        fclose(classfile);
        for(yi=1;yi<ny-1;yi++) {
            fprintf(stderr,"% 6i of % 6i\r",yi,ny-1);
            for(xi=1;xi<nx-1;xi++) {
                pixel p=expixels[yi][xi];
                if(bitmap[yi][xi]) {
                    p_rgb_given_fg[p.r][p.g][p.b]+=1.0;
                    nfg++;
                }
                else {
                    p_rgb_given_bg[p.r][p.g][p.b]+=1.0;
                    nbg++;
                }
            }
        }
        fprintf(stderr,"                                        \r");

        for(yi=1;yi<ny-1;yi++) {
            for(xi=1;xi<nx-1;xi++) {
                int neis=
                    bitmap[yi-1][xi]
                    +(bitmap[yi+1][xi]<<1)
                    +(bitmap[yi][xi-1]<<2)
                    +(bitmap[yi][xi+1]<<3);
                p_neis[neis]+=1.;
                if(bitmap[yi][xi]) {
                    p_neis_given_fg[neis]+=1.;
                }
                else {
                    p_neis_given_bg[neis]+=1.;
                }
            }
        }

        ppm_freearray(expixels,nx);
        pbm_freearray(bitmap,nx);
    }

    p_fg = (double)(nfg+1)/(double)(nfg+nbg+2);
    p_bg = 1.0-p_fg;
    fprintf(stderr,"p(fg)=%g, p(bg)=%g\n",p_fg,p_bg);

    fprintf(stderr,"Normalizing color histograms\n");
    {
        int r,g,b;
        for(r=0;r<256;r++) {
            fprintf(stderr,"% 3i of 255\r",r);
            for(g=0;g<256;g++) {
                for(b=0;b<256;b++) {
                    p_rgb_given_fg[r][g][b] /= nfg+(using_ros?256*256*256:0);
                    p_rgb_given_bg[r][g][b] /= nbg+(using_ros?256*256*256:0);
                }
            }
        }
        fprintf(stderr,"             \r");
    }

    fprintf(stderr,"Normalizing spatial histograms.\n");
    {
        int neis;
        for(neis=0;neis<1<<4;neis++) {
            p_neis_given_fg[neis]/=nfg+(using_ros?16:0);
            p_neis_given_bg[neis]/=nbg+(using_ros?16:0);
            fprintf(stderr,"p(neis:%i%i%i%i|fg) = %1.8f\n",
                    (neis>>3)&0x0001, 
                    (neis>>2)&0x0001, 
                    (neis>>1)&0x0001, 
                    (neis>>0)&0x0001,
                    p_neis_given_fg[neis]);
            fprintf(stderr,"p(neis:%i%i%i%i|bg) = %1.8f\n",
                    (neis>>3)&0x0001, 
                    (neis>>2)&0x0001, 
                    (neis>>1)&0x0001, 
                    (neis>>0)&0x0001,
                    p_neis_given_bg[neis]);
        }
    }
}