/*----------------------------------------------------------------------------

Copyright (c) 2018 J.L. Lisani (joseluis.lisani@uib.es)

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.

----------------------------------------------------------------------------*/

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


#include "library/io_png/io_png.h"
#include "library/io_RGB.h"
#include "library/libmlhe.h"
#include "library/parser.h"


//Read input data
unsigned char *getImageData(const char *name,
                  unsigned char **R,
                  unsigned char **G,
                  unsigned char **B,
                  unsigned char **I,
                  int &w, int &h)
{
    size_t ww, hh;
    unsigned char *img = io_png_read_u8_rgb(name, &ww, &hh);
    w=(int) ww;
    h=(int) hh;
    input2RGB(img, R, G, B, w*h);
    RGBtoI(*R, *G, *B, I, w*h);
    return img;
}

//Save processed result
void saveImageData(const char *name,
                   unsigned char *R,
                   unsigned char *G,
                   unsigned char *B,
                   unsigned char *imgIn,
                   int w, int h)
{
    unsigned char *img = imgIn;
    if (!imgIn) img=new unsigned char[3*w*h];
    RGB2output(R, G, B, img, w*h);
    io_png_write_u8(name, img, w, h, 3); //3 channels
    if (!imgIn) free(img);
}



//Algorithm 6: 
//use original and processed intensities to compute
//the colors of the ouput image
void colorprocess(unsigned char *R, unsigned char *G, unsigned char *B, 
                  unsigned char *I, unsigned char *Iout, int w, int h)
{
    float scale;
    int size=w*h;
    unsigned char *Rout=R;
    unsigned char *Gout=G;
    unsigned char *Bout=B;
    
    for (int n=0; n < size; n++) {
        if (I[n] != 0) {
            scale=(float) Iout[n]/ (float) I[n];
            float rr= (float) R[n] *scale;
            float gg= (float) G[n] *scale;
            float bb= (float) B[n] *scale;
            float outmax=(rr > gg)?((rr > bb)?(rr):(bb)):((gg > bb)?(gg):(bb));
            //assume 8-bits images
            if (outmax > 255) { //out of range: rescale but keeping 
                                //the R/G/B ratio
                rr*=255.0f/outmax;
                gg*=255.0f/outmax;
                bb*=255.0f/outmax;
                
            }
            Rout[n]=(int) (rr+0.5f);
            Gout[n]=(int) (gg+0.5f);
            Bout[n]=(int) (bb+0.5f);
        } else {
            Rout[n]=0;
            Gout[n]=0;
            Bout[n]=0;
        }
    }
    
}

//Main function (Algorithm 1):
// - read names of input and output images and parameters
// - read input image
// - apply MLHE to intensity component
// - modify color of input image using processed and original intensities
// - save result
int main(int argc, char **argv)
{
    std::vector<ParStruct *> parameters;
    ParStruct pInput = {(char *) "input", NULL, (char *) "input image"};
    parameters.push_back(&pInput);
    ParStruct pOutput = {(char *) "output", NULL, (char *) "output image"};
    parameters.push_back(&pOutput);
    
    std::vector <OptStruct *> options;
    OptStruct oO  = {(char *) "o:", 0, (char *) "1", NULL, (char *) "Histogram Processing Option \n \
                     1: Classical Histogram Equalization\n \
                     2: Piecewise Affine Histogram Equalization\n \
                     3: CLAHE"}; 
    options.push_back(&oO);
    OptStruct oA  = {(char *) "a:", 0, (char *) "20", NULL, (char *) "Minimum area of processed connected components"}; 
    options.push_back(&oA);
    OptStruct oL  = {(char *) "l:", 0, (char *) "7", NULL, (char *)  "Maximum recursion depth (maximum=7)"}; 
    options.push_back(&oL);
    OptStruct oR  = {(char *) "r:", 0, (char *) "3.0", NULL, (char *)  "Maximum allowed ratio between ranges of values of connected components before and after processing"}; 
    options.push_back(&oR);
    OptStruct oM  = {(char *) "m:", 0, (char *) "0.8", NULL, (char *) "Minimum allowed ratio between ranges of values of connected components before and after processing"}; 
    options.push_back(&oM);
    OptStruct oN  = {(char *) "n:", 0, (char *) "5", NULL, (char *) "Number of segments in Piecewise Affine Histogram Equalization"}; 
    options.push_back(&oN);
    OptStruct oS  = {(char *) "s:", 0, (char *) "3.0", NULL, (char *) "Maximum slope in Piecewise Affine Histogram Equalization"}; 
    options.push_back(&oS);
    OptStruct oT  = {(char *) "t:", 0, (char *) "1.0", NULL, (char *) "Minimum slope in Piecewise Affine Histogram Equalization"}; 
    options.push_back(&oT);
    OptStruct oC  = {(char *) "c:", 0, (char *) "0.01", NULL, (char *) "Clip factor for CLAHE"}; 
    options.push_back(&oC);
    
    if (!parsecmdline((char *) "mlhe", (char *) "Morpho Local Histogram Equalization", \
                  argc, argv, options, parameters))
        return EXIT_FAILURE;
    
    const char *namein = pInput.value;
    const char *nameout = pOutput.value;
    int HE_option = atoi(oO.value);
    if ((HE_option != 1) && (HE_option != 2) && (HE_option != 3)) {
        fprintf(stderr, "Invalid Histogram Processing option\n");
        return EXIT_FAILURE;
    }

    //Parameters of the method
    int Amin = atoi(oA.value);
    int Lmax = atoi(oL.value);
    float rmax = atof(oR.value);
    float rmin = atof(oM.value);
    int PAE_N = atoi(oN.value);
    float PAE_smax = atof(oS.value);
    float PAE_smin = atof(oT.value);
    float CLAHE_clip = atof(oC.value);
    
    //define variables
    unsigned char *R, *G, *B, *I, *Iout, *img;
    int w, h;

    //read input image
    img = getImageData(namein, &R, &G, &B, &I, w, h);
    
    //apply algorithm
    mlhe a(I, w, h);
    a.apply_algorithm(Lmax, Amin, rmax, rmin, HE_option, PAE_N, \
                      PAE_smin, PAE_smax, CLAHE_clip);
    Iout = a.get_result();

    //apply changes to color channels
    colorprocess(R, G, B, I, Iout, w, h);

    //save result
    saveImageData(nameout, R, G, B, img, w, h);
    
    //clean up
    delete[] R;
    delete[] G;
    delete[] B;
    delete[] I;
    free(img);
    delete[] Iout;
    
    return EXIT_SUCCESS;
    
}