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

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 "libmlhe.h"

//Class constructor
mlhe::mlhe(unsigned char *Iinput, int w, int h): w(w), h(h)
{
    //initialize output image
    I=new unsigned char[w*h];
    memcpy(I, Iinput, w*h);
    
    //auxiliary variable
    marked = new bool[w*h];
}

//Class destructor
mlhe::~mlhe()
{
    delete[] I;
    delete[] marked;
}

//Return processed image
unsigned char *mlhe::get_result()
{
    unsigned char *Iout=new unsigned char[w*h];
    memcpy(Iout, I, w*h);
    return Iout;
}

//First call to recursive algorithm (Algorithm 2)
void mlhe::apply_algorithm(int L, int A, float rrmax, float rrmin, int option,
                           int PAE_N, float PAE_smin, float PAE_smax,
                           float CLAHE_clip)
{
    Lmax = L;
    Amin = A;
    rmax = rrmax;
    rmin = rrmin;
    optionHE = option;
    N = PAE_N;
    smin = PAE_smin;
    smax = PAE_smax;
    clip = CLAHE_clip;
    
    
    //initialize auxiliary variables
    int *selected = new int[w*h]; //list of selected pixels
    //initially all pixels are selected
    for (int n=0; n < w*h; n++) selected[n]=n;
    int nselected = w*h;
    
    //list of 'active' pixels for the  extraction of connected components
    for (int n=0; n < w*h; n++) marked[n]=false;
    
    
    //recursive_histo on selected pixels
    recursive_histo(selected, nselected, 0, 255);
    
    delete[] selected;
}


//Algorithm 7: contrast limited histogram equalization
void mlhe::clahe(int *selected, int nselected, int min, int max)
{
    float hist[256];
    float histAc[256];

    //check range of output values
    vector <int> values;
    for (int n=0; n < nselected; n++) values.push_back(I[selected[n]]);
    sort(values.begin(), values.end());
    float rangeI = (float) (values[nselected-1] - values[0]);
    if (rangeI == 0) return; //return if all the pixels have the same value

    //compute cumulative histogram of values
    memset(hist, 0, 256*sizeof(float));
    for (int n=0; n < nselected; n++) hist[I[selected[n]]]+=1.0f;
    //normalize histogram
    for (int n=0; n < 256; n++) hist[n] /= (float) nselected;
    
    //CLAHE step
    //clip values above threshold
    float pexcess = 0.0f;
    for (int n=min; n <= max; n++)
        if (hist[n] > clip) {
            pexcess += (hist[n]-clip);
            hist[n] = clip;
        }
    //distribute excess values among all levels
    float peach = pexcess / (max-min+1);
    for (int n=min; n <= max; n++) hist[n] += peach;
    
    //compute cumulative histogram of values
    histAc[0]= (float) hist[0];
    for (int n=1; n < 256; n++) histAc[n] = (float) hist[n] + histAc[n-1];  
    
    //check range of output values
    int v, vmin=255, vmax=0;
    values.clear();
    for (int n=0; n < nselected; n++) {
        //apply tone mapping (i.e. cumulative distribution function)
        v = min + (int) ((histAc[I[selected[n]]] * (float) (max-min)) + 0.5f);
        values.push_back(v);
        if (v < vmin) vmin = v;
        if (v > vmax) vmax = v;
    }
    if (vmin == vmax) return; //this can happen due to the rounding of 
                              //floats to integers
    
    //get output values
    for (int n=0; n < nselected; n++) {
        I[selected[n]] = values[n];
    }
    
    

}


//auxiliary data structure for storing linear segments used in PAE
struct pointxy {
    float x, y;
    float slopetonext;
};

//Algorithm 8: piecewise affine histogram equalization (PAE)
void mlhe::piecewise_affine_histogram_equalization(int *selected, int nselected,
                                                   int min, int max)
{
    int hist[256];
    float histAc[256];
        
    //check range of input values
    vector <int> values;
    for (int n=0; n < nselected; n++) values.push_back(I[selected[n]]);
    sort(values.begin(), values.end());
    float rangeI = (float) (values[nselected-1] - values[0]);
    if (rangeI == 0) return; //return if all the pixels have the same value
    
    //compute histogram of values
    memset(hist, 0, 256*sizeof(int));
    for (int n=0; n < nselected; n++) hist[I[selected[n]]]++;
    //compute cumulative histogram of values
    histAc[0]= (float) hist[0];
    for (int n=1; n < 256; n++) histAc[n] = (float) hist[n] + histAc[n-1];
    for (int n=0; n < 256; n++) histAc[n]/= (float) nselected;
    
    //get (x_i, y_i) points for the piecewise affine approximation of 
    //the cumulative histogram
    //initial point: (min, min)
    //separation between y_i points = 1/N, N=number of intervals (N >= 1)
    vector < pointxy > points;
    struct pointxy p;
    p.x = min;
    p.y = min;
    points.push_back(p);
    float next_y = 1.0f/(float) N;
    int i=1;
    int n=0;
    float slope;
    while (i < N) {
        for (;(n < 256) && (histAc[n] < next_y); n++);
        struct pointxy p;
        p.x = (float) n;
        p.y = (float) min + next_y * (max-min);
        if (p.x != points[i-1].x) { //not infinite slope
            slope = (p.y - points[i-1].y)/(p.x - points[i-1].x);
        } else slope = -1.0; //infinite slope
        if ((slope > smax) || (slope == -1)) 
            p.y = points[i-1].y + smax * (p.x - points[i-1].x);
        if (slope < smin)  
            p.y = points[i-1].y + smin * (p.x - points[i-1].x);
        points.push_back(p);
        i++;
        next_y += 1.0f/(float) N;
    }
    //add last interval with x_i=max
    p.x = (float) max;
    p.y = (float) max;
    if (p.x != points[i-1].x) { //not infinite slope
        slope = (p.y - points[i-1].y)/(p.x - points[i-1].x);
    } else slope = -1.0; //infinite slope
    if ((slope > smax) || (slope == -1)) 
        p.y = points[i-1].y + smax * (p.x - points[i-1].x);
    if (slope < smin)  
        p.y = points[i-1].y + smin * (p.x - points[i-1].x);
    points.push_back(p);
    //check that final value is = max
    //reject if > max, since it implies saturation
    //reject if < max since it implies that the dynamic range is not fully used
    //(it may happen when too few values are in the cc)
    if (p.y > max) {
        //rescale all y values so that final value=max, but keeping 
        //the initial value = min
        for (int k = 0; k < points.size(); k++) 
            points[k].y = (float) min + (points[k].y - (float) min) * \
                          ((float) max - (float) min)/(p.y - (float) min);
    }
    if (p.y < max) return;
        
    
    //assign corresponding interval to each value 
    //(assign index to initial point of interval)
    int interval[256];
    for (int i=0; i < points.size()-1; i++) {
        int i_initial = (int) (points[i].x + 0.5f);
        int i_final = (int) (points[i+1].x + 0.5f);
        for (int k=i_initial; k <= i_final; k++) interval[k] = i;
    }
    
    //add slope information
    for (int i=0; i < points.size(); i++) {
        if (i < points.size()-1) {
            if (points[i+1].x != points[i].x) { //not infinite slope
                points[i].slopetonext = \
                    (points[i+1].y - points[i].y)/(points[i+1].x - points[i].x);
            } else points[i].slopetonext = -1.0; //infinite slope
        } else points[i].slopetonext=0; //last point
    }
    
    //Apply transform
    float rangeF;
    int v, vmin=255, vmax=0;
    values.clear();
    for (int n=0; n < nselected; n++) {
        int i = interval[I[selected[n]]]; //get interval
        v = (int) ( points[i].y + points[i].slopetonext * \
            ( (float) I[selected[n]] - points[i].x ) + 0.5f);
        values.push_back(v);
        if (v < vmin) vmin = v;
        if (v > vmax) vmax = v;
    }
    if (vmin == vmax) return; //this can happen due to the rounding of 
                              //floats to integers
    
    //get output values
    for (int n=0; n < nselected; n++) {
        I[selected[n]] = values[n];
    }

}



//Algorithm 5: classical histogram equalization
void mlhe::histogram_equalization(int *selected, 
                                  int nselected, int min, int max)
{
    int hist[256];
    float histAc[256];

    //check range of input values
    vector <int> values;
    for (int n=0; n < nselected; n++) values.push_back(I[selected[n]]);
    sort(values.begin(), values.end());
    float rangeI = (float) (values[nselected-1] - values[0]);
    if (rangeI == 0) return; //return if all the pixels have the same value

    //compute histogram of values
    memset(hist, 0, 256*sizeof(int));
    for (int n=0; n < nselected; n++) hist[I[selected[n]]]++;
    //compute cumulative histogram of values
    histAc[0]= (float) hist[0];
    for (int n=1; n < 256; n++) histAc[n] = (float) hist[n] + histAc[n-1];
    for (int n=0; n < 256; n++) histAc[n]/= (float) nselected;
        
    //check range of output values
    float rangeF;
    int v, vmin=255, vmax=0;
    values.clear();
    for (int n=0; n < nselected; n++) {
        //apply tone mapping (i.e. cumulative distribution function)
        v = min + (int) ((histAc[I[selected[n]]] * (float) (max-min)) + 0.5f);
        values.push_back(v);
        if (v < vmin) vmin = v;
        if (v > vmax) vmax = v;
    }
    
    float ratioranges;
    if (vmin != vmax) { //vmin = vmax can happen due to the rounding of 
                        //floats to integers
    
        rangeF = (float) (vmax-vmin);

        ratioranges = (rangeI != 0)?(rangeF/rangeI):(-1);
        
        //check ratio between original and processed ranges of values
        if ((ratioranges != -1) && (ratioranges <= rmax) && \
            (ratioranges >= rmin)) {
            //get output values
            for (int n=0; n < nselected; n++) {
                I[selected[n]] = values[n];
            }
        }
    
    }
    
}


//Algorithm 4 (lines 8 to 15): fill connected component
void mlhe::fill_cc(vector<int> &cc, int p0)
{
    //recursively grow connected component, use 4-connectivity
    int p = p0;
    marked[p] = false; //remove from list so it is not used again
    bool grow = true;
    queue < int > listtogrow;
    int k=1;
    while (grow) {
        cc.push_back(p); 
        int x = p%w;
        int y = p/w;
        if ((x > 0) && (marked[x-1+y*w])) { //add left neighbor to list 
            listtogrow.push(x-1+y*w); 
            marked[x-1+y*w] = false;
        }
        if ((x+1 < w) && (marked[x+1+y*w])) { //add right neighbor to list 
            listtogrow.push(x+1+y*w); 
            marked[x+1+y*w] = false; 
        }
        if ((y > 0) && (marked[x+(y-1)*w])) { //add upper neighbor to list 
            listtogrow.push(x+(y-1)*w); 
            marked[x+(y-1)*w] = false; 
        }
        if ((y+1 < h) && (marked[x+(y+1)*w])) { //add lower neighbor to list 
            listtogrow.push(x+(y+1)*w); 
            marked[x+(y+1)*w] = false; 
        }
        grow = !listtogrow.empty();
        if (grow) {
            p = listtogrow.front(); //get first element of the list
            listtogrow.pop(); //remove from list
        }
    } 
}

//Algorithm 3: 
// - get bi-level sets
// - extract connected components (Algorithm 4)
// - apply Algorithm 2 to connected components (use recursivity)
void mlhe::process_bilevelset(int *selected, int nselected, 
                              int min, int max)
{
    //mark pixels belonging to the level set
    for (int n=0; n < nselected; n++)
        if ((I[selected[n]] >= min) && (I[selected[n]] <= max)) 
            marked[selected[n]]=true;
   
    //get connected components of level set (Algorithm 4: lines 1 to 7)
    vector< vector<int> > ccs; //a list of lists of pixels
    for (int n=0; n < nselected; n++) 
        if (marked[selected[n]]) {
            //new connected component: a list of pixels 
            vector<int> cc;
            fill_cc(cc, selected[n]);
            ccs.push_back(cc);
        }
    
   //apply recursive histo equalization to each connected component
   for (int i=0; i < ccs.size(); i++) {
        //list of selected pixels for further histogram equalization
        int nselected2 = ccs[i].size();
        if (nselected2 >= Amin) { //do not process small connected components
            int *selected2 = new int[nselected2];
            for (int n = 0; n < nselected2; n++) selected2[n] = ccs[i][n];
            //recursive_histo on selected pixels
            recursive_histo(selected2, nselected2, min, max);
            delete[] selected2;
        }
    }

}

//Algorithm 2: 
// - apply histogram equalization in range [min, max]
// - process bi-level set [min, (min+max)/2]
// - process bi-level set [(min+max)/2 + 1, max]
void mlhe::recursive_histo(int *selected, int nselected,
                           int min, int max)
{

    int level = (int) (log2(256.0 / (double) (max-min+1)) + 0.5);

    //histogram equalization
    switch (optionHE) {
        case 1: //classic method (Algorithm 5)
            histogram_equalization(selected, nselected, min, max);
            break;
        
        case 2: //PAE method (Algorithm 8)
            piecewise_affine_histogram_equalization(selected, nselected, min, max);
            break;
        
        case 3: //CLAHE-based method (Algorithm 7)
            clahe(selected, nselected, min, max);
            break;
            
        default: //this should never happen
            fprintf(stderr, "Invalid Equalization Option\n");
            exit(EXIT_FAILURE);
    }
    
    if (level+1 > Lmax) return;
    
    if (max - min <= 2) return; //not enough values in range for dyadic split
    
    //center of interval of values
    int med=(min+max)/2;
    
    //Process connected components [min, med] and [med+1, max]
    process_bilevelset(selected, nselected, min, med);
    process_bilevelset(selected, nselected, med+1, max);
     
}