/* historgram.cpp
 * scott wehrwein
 * make a histogram of color intensities between
 * two color png images
 */


#include "imageLib.h"
#include <cmath>
#include <string.h>
#define VERBOSE 1
#define GAMMA 3.0
#define ABS(x) ((x) >= 0 ? (x) : (-(x)))

int w, h, nB;
int startx = 0, starty = 0;
int incr;

void polyfit(CByteImage *im1, CByteImage *im2, int degree, double X[][3]) {
    
    int n = degree;

    CShape sh = im1->Shape();

    double XtotheN[2*(n-1)];
    double Ys[n];
    for (int i = 0; i < 2*(n-1); i++) {
        XtotheN[i] = 0;
        if (i < n) {
            Ys[i] = 0;
        }
    }

    for (int b = 0; b < 3; b++) {
        for (int y = 0; y < h; y++) {
            for (int x = 0; x < w; x++) {
                int im1val = im1->Pixel(x,y,b);
                for (int i = 0; i < 2*n-1; i++) {
                    XtotheN[i] += pow((double)im1val,(double)i);
                    if (i < n) {
                        Ys[i] += im2->Pixel(x,y,b) * pow((double)im1val, (double)i);
                    }
                }
            }
        }

        /*printf("\n\n");
        for (int i =0; i < 2*n-1; i++) {
            printf("X^%d: %f\n", i, XtotheN[i]);
        }*/
    
        int m = n+1;
        double M[n][m];
        for (int i = 0; i < n; i++) {
            for (int j = i; j < n; j++) {
                M[i][j] = XtotheN[2*n-i-j-2];
                M[j][i] = XtotheN[2*n-i-j-2];
            }
            M[i][n] = Ys[n-i-1];
        }    
    
        /*printf("\n\n");
        for ( int i =0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                printf("%2.2f ", M[i][j]);
            }
            printf("\n");
        }*/

        int k,i,j,iPivot,jPivot;
        double SwapColTmp, SwapRowTmp, NullFactor;
        
        // REQUIRES A EXTRA STORAGE TO RECORD COLUMN SWAPS (FULL PIVOTING SWAPS
        // BOTH ROWS AND COLUMNS SINCE THE PIVOT ELEMENT IS CHOSEN FROM THE ENTIRE
        // SUBMATRIX, RATHER THAN JUST THE COL AS IN PARTIAL PIVOTING).
        // EACH ELT j IN SWAP STORES THE NEW LOCATION OF THE jth COLUMN.
        int *Swap = new int[n];         // DOES NOT INCLUDE LAST COLUMN
        for (j=0; j<n; j++) Swap[j]=j;  // INIT TO INDICATE NO SWAPPING HAS OCCURRED
        int tSwap;                      // TEMP SWAPPING VARIABLE FOR SWAP ELEMENTS
        
        //------- PART 1: CONVERSION OF MATRIX M INTO ROW-ECHELON FORM -------
        for (k=0; k<n; k++)  // FOR EACH SUBMATRIX SIZE (n-k) x (n-m)
        {
            // FIND THE ELEMENT IN THE KTH SUBMATRIX WITH THE LARGEST ABS VAL (PIVOT ELEMENT)
            // EXCLUDE THE LAST COLUMN (NOT A PART OF THE COEFFICIENT MATRIX).
            iPivot = jPivot = k;
            for (i=k; i<n; i++) {
                for (j=k; j<n; j++) {
                    if (ABS(M[i][j]) > ABS(M[iPivot][jPivot]))
                    { iPivot=i; jPivot=j; }
                }
            }
            
            // IF THERE IS NO NONZERO ELEMENT AS A PIVOT, QUIT!
            if (M[iPivot][jPivot] == 0)
                throw CError("GaussElim: singular matrix!");
            
            // SWAP THE PIVOT ROW WITH THE KTH ROW
            for (j=0; j<m; j++)
            {           
                SwapRowTmp = M[iPivot][j];
                M[iPivot][j] = M[k][j];
                M[k][j] = SwapRowTmp;
            }
            // SWAP THE PIVOT COL WITH THE KTH COL. BE SURE TO STORE THE SWAP.
            for (i=0; i<n; i++)
            {
                SwapColTmp = M[i][jPivot];
                M[i][jPivot] = M[i][k];
                M[i][k] = SwapColTmp;
            }
            tSwap = Swap[jPivot];
            Swap[jPivot] = Swap[k];
            Swap[k] = tSwap;
            
            // NULLIFY (TURN INTO ZERO) ALL ELEMENT IN COL BELOW TOP ROW BY
            // SUBTRACTING SOME FACTOR TIMES THE TOP ROW.
            for (i=k+1; i<n; i++)
            {
                NullFactor = -(M[i][k]/M[k][k]);
                M[i][k]=0;
                for (j=k+1; j<m; j++)
                    M[i][j] += (M[k][j]*NullFactor);
            }
            
            // IF THEIR IS A ROW CONTAINING ALL ZEROS (EXCEPT FOR LAST ELT), QUIT!
            if (M[n-1][n-1] == 0) 
                throw CError("GaussElim: singular matrix!");
        }
        
        //------- PART 2: BACKWARDS SUBSTITUTION -------
        for (k=(n-1); k>=0; k--)  // FOR EACH ROW STARTING AT THE BOTTOM
        {
            X[k][b] = M[k][m-1];      // SOLVE FOR CURRENT X[b][k]
            for (j=(k+1); j<(m-1); j++)
                X[k][b] -= (M[k][j]*X[j][b]);
            X[k][b] /= M[k][k];
        }
        
        // MUST REARRANGE THE SOLUTION VECTOR X BASED ON SWAP ARRAY
        double *Xcopy = new double[n];
        for (i=0; i<n; i++) Xcopy[i] = X[i][b];
        for (i=0; i<n; i++) X[ Swap[i] ][b] = Xcopy[i];
        delete[] Xcopy; // DELETE COPY OF X SOLUTION COLUMN VECTOR
        
        delete[] Swap;  // FREE COLUMN SWAP SPACE

/*        printf("\nBand %d:\n", b);
        for (int i = 0; i < n; i++) {
            printf("X[%d][%d]: %f\n", i, b, X[i][b]);
        }
*/
    }

    /* residual and corrected */
    CByteImage residual(im1->Shape());
    CByteImage corrected(im1->Shape());
    for (int y = 0; y < h; y++) {
        uchar *residualrow = &residual.Pixel(0, y, 0);
        uchar *correctedrow = &corrected.Pixel(0, y, 0);
        uchar *im1row = &im1->Pixel(0, y, 0);
        uchar *im2row = &im2->Pixel(0, y, 0);
        for (int x = 0; x < w; x++) {
            for (int b = 0; b < 3; b++) {
	            int ix = im1row[b];
                double iy = 0;
                for (int d = 0; d < degree; d++) {
                    iy += X[degree - d - 1][b] * pow((double)ix, (double)d);
                }
                int val = (int)iy;
                val = __max(0, __min(255, val));
                correctedrow[b] = val;

                int residscale = 3; // scale up residual error
                int resid = 128 + residscale * (int)(iy - im2row[b]);
                residualrow[b] = __max(0, __min(255, resid));
            }
            residualrow[3] = 255; // full alpha
            correctedrow[3] = 255;
            residualrow += 4;
            correctedrow += 4;
            im1row += 4;
            im2row += 4;
        }
    }
    WriteImageVerb(corrected, "im1-corrected.ppm", VERBOSE);
    WriteImageVerb(residual, "residual.ppm", VERBOSE);


}

void logfit (CByteImage *im1, CByteImage *im2, double X[3][3], int mode) {

    int swap[3] = {0,0,0}; // whether we've swapped im1 and im2 to get a better line

    // pre-crunch all the logs
    double logtab[257];
    for (int i = 1; i < 257; i++) {
        logtab[i] = log(i);
    }

    double la[3], lb[3]; // line parameters (in log space) per color band
    double s1[]={0,0,0}, sx[]={0,0,0}, sy[]={0,0,0}, sxx[]={0,0,0}, sxy[]={0,0,0}, syy[]={0,0,0}; // least squares stuffs


    for (int y = starty; y < starty + h; y++) {
        uchar *im1row = &im1->Pixel(starty, y, 0);
        uchar *im2row = &im2->Pixel(starty, y, 0);
        int jump = incr * 4;
        for (int x = startx; x < startx + w; x+= incr) {
            int i1[] = {im1row[0], im1row[1], im1row[2]};
            int i2[] = {im2row[0], im2row[1], im2row[2]};
            im1row += jump;
            im2row += jump;
            for (int b = 0; b < 3; b++) {
                // line fit to log values
                int sat = 240; // saturation threshold (avoid fitting to saturated colors)
                if (i1[b]>0 && i2[b]>0 // numbers must be positive
                    && i1[b] < sat && i2[b] < sat) {  // and below saturation threshold
                    double x = logtab[i1[b]];
                    double y = logtab[i2[b]];
                    s1[b] += 1;
                    sx[b] += x;
                    sy[b] += y;
                    sxx[b] += x*x;
                    sxy[b] += x*y;
                    syy[b] += y*y;
                }

                if (mode == 1) {
                  i1[b] = (int)((logtab[__max(10, i1[b])]-logtab[10])/(logtab[256]-logtab[10])*255.0);
                  i2[b] = (int)((logtab[__max(10, i2[b])]-logtab[10])/(logtab[256]-logtab[10])*255.0);
                }

            }
        }
    }
    
    for (int b = 0; b < 3; b++) {
    // compute parameters of line
        double detx = (sxx[b] * s1[b] - sx[b] * sx[b]);
        double dety = (syy[b] * s1[b] - sy[b] * sy[b]);
        if ((fabs(detx) < 1e-6) && (fabs(dety) < 1e-6)) { // unstable
            la[b] = 1.0;
            lb[b] = 0.0;
        } else {
            la[b] = (s1[b] * sxy[b] -  sx[b] * sy[b]) / detx;
            lb[b] = (-sx[b] * sxy[b] + sxx[b] * sy[b]) / detx;
            if (la[b] > 1.0) {
                la[b] = (s1[b] * sxy[b] -  sx[b] * sy[b]) / dety;
                lb[b] = (-sy[b] * sxy[b] + syy[b] * sx[b]) / dety;
                swap[b] = 1;
            }
        }
        X[0][b] = la[b];
        X[1][b] = lb[b];
        X[2][b] = swap[b]; // 3rd column keeps track of whether we wapped
        printf("line %d: a=%g, b=%g (%s)\n", b, la[b], lb[b], (swap[b] ? "swapped" : "not swapped"));
    }

    /* residual and corrected */

    CByteImage residual(im1->Shape());
    CByteImage corrected(im1->Shape());
    for (int y = 0; y < h; y++) {
        uchar *residualrow = &residual.Pixel(0, y, 0);
        uchar *correctedrow = &corrected.Pixel(0, y, 0);
        uchar *im1row = &im1->Pixel(0, y, 0);
        uchar *im2row = &im2->Pixel(0, y, 0);
        for (int x = 0; x < w; x++) {
            for (int b = 0; b < 3; b++) {
                double xa;
                double xb;
                if (X[3][b] != 0) {
                    xa = 1 / X[0][b];
                    xb = -X[1][b] / X[0][b];
                } else {
                    xa = X[0][b];
                    xb = X[1][b];
                }
	            int ix = im1row[b];
                int iy = (int)(pow(ix, xa) * exp(xb) + .5);
                iy = __max(0, __min(255, iy));
                correctedrow[b] = iy;

                int residscale = 3; // scale up residual error
                int resid = 128 + residscale * (iy - im2row[b]);
                residualrow[b] = __max(0, __min(255, resid));
            }
            residualrow[3] = 255; // full alpha
            correctedrow[3] = 255;
            residualrow += 4;
            correctedrow += 4;
            im1row += 4;
            im2row += 4;
        }
    }
    WriteImageVerb(corrected, "im1-corrected.ppm", VERBOSE);
    WriteImageVerb(residual, "residual.ppm", VERBOSE);

}


CByteImage normalize(CIntImage *count, int max[]) {
    CByteImage result;
    result.ReAllocate(count->Shape());
    /* scale each pixel linearly to a value between 0.0 and 1.0
       take it to the 1.0/GAMMA power
       then scale it back up to a value between 0 and 255*/
    for (int y = 0; y < 256; y++) {
        uchar *resultrow = &result.Pixel(0, y, 0);
        int *countrow = &(count->Pixel(0,y,0));
        for (int x = 0; x < 256; x++) {
            for (int b = 0; b < 3; b++) {
		double val = __min(1.0, (double)countrow[b] / max[b]);
                resultrow[b] += (uchar)(pow(val, (1.0/GAMMA)) * 255);
            }
            resultrow += 4;
            countrow += 4;
        }
    }
    return result;
}

int main(int argc, char *argv[])
{
try {
    int mode = 0;
    int twist = 0;
    w = 0;
    h = 0;
    startx = 0;
    starty = 0;
    incr = 1; 

    if (argc < 3) {
        throw CError("\n  usage: %s im1.png im2.png [-mode <m>] [<x> <y> <width> <height>] [-fast <px>] [-twist <model>]\n", argv[0]);
    }
    
    // parse args
    for (int i = 3; i < argc; i++) {
        if (strcmp(argv[i], "-mode") == 0) {
                mode = atoi(argv[i+1]);
                i++;
        } else if (strcmp(argv[i], "-c") == 0) {
                startx = atoi(argv[i+1]); 
                starty = atoi(argv[i+2]); 
                w = atoi(argv[i+3]); 
                h = atoi(argv[i+4]);
                i += 4;
        } else if (strcmp(argv[i], "-fast") == 0) {
                incr = atoi(argv[i+1]);
                i++;
        } else if (strcmp(argv[i], "-twist") == 0) {
                twist = atoi(argv[i+1]);
                i++;
        } else {
                throw CError("\n  usage: %s im1.png im2.png [-mode <m>] [<x> <y> <width> <height>] [-fast <px>] [-twist <model>]\n", argv[0]);
        }
    }

    // twist im1 if necessary, using linear (twist) or quadratic (bigtwist)
    char *im1name = argv[1];
    if (twist) {
        printf("Twisting.\n");
        im1name = "twisted.png";
        char twistcommand[1000];
        if (twist == 1) {
            sprintf(twistcommand, "/home/schar/public_html/colormatching/runtwist %s %s twisted.png", argv[1], argv[2]); 
        } else if (twist == 2) {
	    throw CError("not available");
            sprintf(twistcommand, "/home/swehrwei/public_html/summer2008/bigtwist %s %s twisted.png", argv[1], argv[2]);
        }
        system(twistcommand); 
    }



 
    CByteImage im1, im2;
    ReadImageVerb(im1, im1name, VERBOSE);
    ReadImageVerb(im2, argv[2], VERBOSE);

    CShape sh1 = im1.Shape();
    CShape sh2 = im2.Shape();
    
    // if no crop region is specified, use the whole width and height
    if (w == 0) {
        w = sh1.width;
        h = sh1.height;
        if (sh2.width != w || sh2.height != h)
            throw CError("\n  Images must have the same dimensions\n");

    }

    if (w + startx > sh2.width || w + startx > sh1.width
      ||h + starty > sh2.height || h + starty > sh1.height)
            throw CError("\n  At least one image is not large enough for the given crop area\n");
    
    nB = sh1.nBands;
    if (sh2.nBands != nB)
        throw CError("\n  Images do not have the same number of bands\n");


    CIntImage count; // the counter int image
    CShape sh = CShape(256,256,4);
    count.ReAllocate(sh);


    int max[3] = {0,0,0}; // max value for each band
//    int swap[3] = {0,0,0}; // whether we've swapped im1 and im2 to get a better line
    double logtab[257];
    for (int i = 1; i < 257; i++) {
        logtab[i] = log(i);
    }

//    double la[3], lb[3]; // line parameters (in log space) per color band
    double s1[]={0,0,0}, sx[]={0,0,0}, sy[]={0,0,0}, sxx[]={0,0,0}, sxy[]={0,0,0}, syy[]={0,0,0};
    int npix = 0;
    for (int y = starty; y < starty + h; y++) {
        uchar *im1row = &im1.Pixel(starty, y, 0);
        uchar *im2row = &im2.Pixel(starty, y, 0);
        int jump = incr * 4;
        for (int x = startx; x < startx + w; x+= incr) {
            int i1[] = {im1row[0], im1row[1], im1row[2]};
            int i2[] = {im2row[0], im2row[1], im2row[2]};
            im1row += jump;
            im2row += jump;
	    npix++;
            for (int b = 0; b < 3; b++) {
                // line fit to log values
                int sat = 240; // saturation threshold (avoid fitting to saturated colors)
                if (i1[b]>0 && i2[b]>0 // numbers must be positive
                    && i1[b] < sat && i2[b] < sat) {  // and below saturation threshold
                    double x = logtab[i1[b]];
                    double y = logtab[i2[b]];
                    s1[b] += 1;
                    sx[b] += x;
                    sy[b] += y;
                    sxx[b] += x*x;
                    sxy[b] += x*y;
                    syy[b] += y*y;
                }

                if (mode == 1) {
                  i1[b] = (int)((logtab[__max(10, i1[b])]-logtab[10])/(logtab[256]-logtab[10])*255.0);
                  i2[b] = (int)((logtab[__max(10, i2[b])]-logtab[10])/(logtab[256]-logtab[10])*255.0);
                }

                int *countpixel = &count.Pixel(i1[b], 255-i2[b], 0);
                countpixel[b]++;
                if (countpixel[b] > max[b])
                    max[b] = countpixel[b];

            }
        }
    }
/*    for (int b = 0; b < 3; b++) {
    // compute parameters of line
        double detx = (sxx[b] * s1[b] - sx[b] * sx[b]);
        double dety = (syy[b] * s1[b] - sy[b] * sy[b]);
        if ((fabs(detx) < 1e-6) && (fabs(dety) < 1e-6)) { // unstable
	        la[b] = 1.0;
	        lb[b] = 0.0;
        } else {
            la[b] = (s1[b] * sxy[b] -  sx[b] * sy[b]) / detx;
            lb[b] = (-sx[b] * sxy[b] + sxx[b] * sy[b]) / detx;
            if (la[b] > 1.0) {
                la[b] = (s1[b] * sxy[b] -  sx[b] * sy[b]) / dety;
                lb[b] = (-sy[b] * sxy[b] + syy[b] * sx[b]) / dety;
                swap[b] = 1;
            }
        }
        printf("line %d: a=%g, b=%g (%s)\n", b, la[b], lb[b], (swap[b] ? "swapped" : "not swapped"));
    }
*/

    for (int k = 0; k < 3; k++) {
	//max[k] /= 20; // increase contrast
	max[k] = npix/400; // normalize by constant amount instead of max
    }
    CByteImage result = normalize(&count, max);

    // skip line drawing for now
    if (0) {

    /********************/ 
    /* log fitting case */
    /********************/
    if (mode < 2) {
        double X[3][3];  
        logfit(&im1, &im2, X, mode);

        // plot lines
        for (int b = 0; b < 3; b++) {
            if (mode == 1) { // plot in log coordinates
                for (double x=0; x <= 255; x += .05) {
                    double y = X[0][b] * x + X[1][b];
                    int i1 = (int)((x-logtab[10])/(logtab[256]-logtab[10])*255.0);
                    int i2 = (int)((y-logtab[10])/(logtab[256]-logtab[10])*255.0);
                    if (i1 >= 0 && i1 < 256 && i2 >= 0 && i2 < 256) {
                        if (X[2][b] != 0) {
                            result.Pixel(i2, 255-i1, b) = 255;
                        } else {
                            result.Pixel(i1, 255-i2, b) = 255;
                        }
                    }
                }
            } else { // plot in regular coordinates
                for (int x=0; x <= 255; x += 4) {
                    int y = (int)(pow(x, X[0][b]) * exp(X[1][b]) + .5);
                    if (y >= 0 && y < 256) {
                        if (X[2][b] != 0) { //if swapped
                            result.Pixel(y, 255-x, b) = 255;
                        } else {
                            result.Pixel(x, 255-y, b) = 255;
                        }
                    }
                }
            }
        }
    } else {
        double X[mode][3];
        polyfit(&im1, &im2, mode, X);
        for (int b = 0; b < 3; b++) {
            for (int x = 0; x < 255; x += 4) {
                double y = 0;
                for (int d = 0; d < mode; d++) {
                    y += X[mode - d - 1][b] * pow((double)x, d);
                }
                if (y >= 0 && y < 256) {
                    result.Pixel(x, 255-((int)y), b) = 255;
                }
            } 
        }
    }
    }
    
    sh.nBands = 1;
    CByteImage red(sh);
    CByteImage green(sh);
    CByteImage blue(sh);

    BandSelect(result, red, 2, 0);
    BandSelect(result, green, 1, 0);
    BandSelect(result, blue, 0, 0);

    // set alpha channel to 255's
    for (int y = 0; y < 256; y++) {
        uchar *row = (&result.Pixel(0, y, 0)) + 3;
        for (int x = 0; x < 256; x++) {
            row[x*4] = 255;
        }
    }
    WriteImageVerb(red, "r.pgm", VERBOSE);
    WriteImageVerb(green, "g.pgm", VERBOSE);
    WriteImageVerb(blue, "b.pgm", VERBOSE);
    WriteImageVerb(result, "rgb.ppm", VERBOSE);
    // try correcting image 1 to image 2's color model
    // do whole image for now (probably should output cropped regions instead)

/*
    // if the image is greater than 1200x900, don't compute corrected and residual
    // skip for now
    if (1 &&w > 1200 && h > 900 || incr != 1) {
        printf("Not computing corrected and residual.\n");
        return 0;
    }

    // also compute residual
    CByteImage residual(sh1);
    CByteImage corrected(sh1);
    for (int y = 0; y < h; y++) {
        uchar *residualrow = &residual.Pixel(0, y, 0);
        uchar *corretedrow = &corrected.Pixel(0, y, 0);
        uchar *im1row = &im1.Pixel(0, y, 0);
        uchar *im2row = &im2.Pixel(0, y, 0);
        for (int x = 0; x < w; x++) {
            for (int b = 0; b < 3; b++) {
                double xa;
                double xb;
                if (X[b][3] != 0) {
                    xa = 1 / X[b][0];
                    xb = -X[b][1] / X[b][0];
                } else {
                    xa = X[b][0];
                    xb = X[b][1];
                }
	            int ix = im1row[b];
                int iy = (int)(pow(ix, xa) * exp(xb) + .5);
                iy = __max(0, __min(255, iy));
                correctedrow[b] = iy;

                int residscale = 3; // scale up residual error
                int resid = 128 + residscale * (iy - im2row[b]);
                residualrow[b] = __max(0, __min(255, resid));
            }
            residualrow[3] = 255; // full alpha
            correctedrow[3] = 255;
            residualrow += 4;
            correctedrow += 4;
            im1row += 4;
            im2row += 4;
        }
    }
    WriteImageVerb(im1, "im1-corrected.ppm", VERBOSE);
    WriteImageVerb(residual, "residual.ppm", VERBOSE);
*/
 } catch (CError &err) {
    fprintf(stderr, err.message);
    fprintf(stderr, "\n");
    return -1;
}

return 0;
}