/*
 * Copyright 2009-2015 IPOL Image Processing On Line http://www.ipol.im/
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/**
 * @mainpage Implementation of Nonlocal Pansharpening Image Fusion (version 2)
 *
 * README.txt:
 * @verbinclude README.txt
 */

/**
 * @file   pansharpening_ipol.cpp
 * @brief  Main executable file
 *
 * @author Joan Duran <joan.duran@uib.es>
 */


#include "libpansharpening.h"
#include "libauxiliar.h"
#include "io_png.h"

// Usage: pansharpening_ipol truth.png pan.png spectral.png IHS.png
//        pansharpened.png sampling

int main(int argc, char **argv)
{
    if(argc < 7)
	{
        printf("usage: pansharpening_ipol truth.png pan.png spectral.png "
               "IHS.png pansharpened.png sampling\n\n");
        printf("truth.png    :: true high-resolution RGB image.\n");
        printf("pan.png      :: panchromatic grayscale image.\n");
        printf("spectral.png :: low-resolution RGB image.\n");
		printf("IHS.png      :: initialized image by IHS.\n");
        printf("pansharpened :: pansharpened fused image.\n");
		printf("sampling     :: sampling factor (2 or 4).\n");
        printf("\n");
        printf("The following parameters are fixed in the main demo "
               "function:\n");
        printf("h              :: filtering parameter of regularizing term.\n");
        printf("lmb            :: trade-off parameter of panchromatic\n"
               "                  constraint.\n");
        printf("mu             :: trade-off parameter of spectral\n"
               "                  constraint.\n");
        printf("std            :: standard deviation of the Gaussian kernel\n"
               "                  used to simulate the low-resolution\n"
               "                  multispectral image.\n");
        printf("num_iterations :: number of iterations of the explicit\n"
               "                  scheme.\n");
        printf("reswind        :: half-size of the support zone used for\n"
               "                  searching similarities among pixels in the\n"
               "                  image.\n");
        printf("compwind       :: half-size of the comparison window used for\n"
               "                  computing distances on the panchromatic.\n");
        printf("dt             :: time-step parameter of the explicit\n"
               "                  scheme.\n");
        printf("alpha          :: mixing coefficients that give the\n"
               "                  panchromatic image as a linear combination\n"
               "                  of the spectral components.\n");
	    
        return EXIT_FAILURE;
	}

    
    // Read truth image
    size_t nx, ny, nc;
    float *d_v = NULL;
    
    d_v = io_png_read_f32(argv[1], &nx, &ny, &nc);
    
    if(!d_v)
	{
	    fprintf(stderr, "Error - %s not found  or not a correct png image.\n",
                argv[1]);
        return EXIT_FAILURE;
    }
    
    if(nc > 3)  // Use only RGB image
        nc = 3;
    
    
    // Variables of high-resolution image
    int high_w = (int) nx;
    int high_h = (int) ny;
    int num_channels = (int) nc;
	int high_wh = high_w * high_h;
    int high_whc = num_channels * high_wh;
    
    float **truth = new float*[num_channels];
    for(int c = 0; c < num_channels; c++)
        truth[c] = &d_v[c*high_wh];
    
    
    // Parameters
    int sampling_factor = atoi(argv[6]);
    
    if(sampling_factor != 2 && sampling_factor != 4)
    {
        fprintf(stderr, "Error - sampling must be 2 or 4.\n");
        return EXIT_FAILURE;
    }
    
    if((high_w % sampling_factor != 0) || (high_h % sampling_factor != 0))
    {
        // Check validity of sampling factor
        fprintf(stderr, "Error - invalid sampling factor. Sampling must be a "
                "divisor of both width and height of truth image\n");
        return EXIT_FAILURE;
    }
    
    float std, h;

    if(sampling_factor == 2)
    {
        std = 1.2f;
        h = 1.25f;
        
    } else
    {
        std = 2.2f;
        h = 6.0f;
    }
    
    float lmb = 100.0f;
    float mu = 100.0f;
    int num_iterations = 100;
    float tol = 0.001f;
    int reswind = 3;
    int compwind = 1;
    float dt = 0.01f;
    
    float *alpha = new float[num_channels];
    for (int c = 0; c < num_channels; c++)
        alpha[c] = (float) 1/num_channels;
    
    
    // Simulate panchromatic image
    float *pan = new float[high_wh];
    
    if(panchromatic(truth, pan, alpha, num_channels, high_w, high_h) != 1)
        return EXIT_FAILURE;

    
    // Simulate low-resolution RGB image
    int low_w, low_h;
    
    float **spectral;
    spectral = low_resolution(sampling_factor, std, truth, num_channels,
                              high_w, high_h, low_w, low_h);
    
    int low_wh = low_w * low_h;
    int low_whc = num_channels * low_wh;
    
    
	// Initialization for IHS transform. We use a cubic spline interpolation
    float **initial_splines = new float*[num_channels];
    for(int c = 0; c < num_channels; c++)
        initial_splines[c] = new float[high_wh];
    
    if(init_splines(sampling_factor, spectral, initial_splines, num_channels,
                    high_w, high_h, low_w, low_h) != 1)
        return EXIT_FAILURE;
    
    
	// Pansharpening by IHS transform
    // The result provided by IHS is used as initialization of the proposed
    // nonlocal pansharpening algorithm
    float **initial_ihs = new float*[num_channels];
    for(int c = 0; c < num_channels; c++)
        initial_ihs[c] = new float[high_wh];
    
    if(ihs(initial_splines, initial_ihs, pan, num_channels, high_w,
           high_h) != 1)
        return EXIT_FAILURE;
    
    
    // Quantization of IHS
    float **quantized_ihs = new float*[num_channels];
    
    for(int c = 0; c < num_channels; c++)
    {
        quantized_ihs[c] = new float[high_wh];
    
        for(int i = 0; i < high_wh; i++)
            quantized_ihs[c][i] = rintf(MAX(0, MIN(initial_ihs[c][i], 255)));
    }
    
    
    // Extended image from low-resolution to high-resolution by replication
    float **replicated = new float*[num_channels];
    for(int c = 0; c < num_channels; c++)
        replicated[c] = new float[high_wh];
    
    if(extension(sampling_factor, spectral, replicated, num_channels, high_w,
                 high_h, low_w) != 1)
        return EXIT_FAILURE;

    
    // Pansharpening by nonlocal algorithm
    
    float **pansharpened = new float*[num_channels];
	for(int c = 0; c < num_channels; c++)
    pansharpened[c] = new float[high_wh];

    if(nonlocal_pansharpening(sampling_factor, std, h, lmb, mu, num_iterations,
                              tol, reswind, compwind, dt, alpha, pan,
                              quantized_ihs, replicated, pansharpened,
                              num_channels, high_w, high_h) != 1)
        return EXIT_FAILURE;
    
    
    // Save images as PNG
    
    // Saving panchromatic image
    if(io_png_write_f32(argv[2], pan, (size_t)high_w, (size_t)high_h, 1) != 0)
    {
        fprintf(stderr, "Error - Failed to save png image %s.\n", argv[2]);
        return EXIT_FAILURE;
    }
    
    // Saving low-spectral resolution image
    float *low_png = new float[low_whc];
    int k = 0;
    for(int c = 0; c < num_channels; c++)
        for(int i = 0; i < low_wh; i++)
        {
            low_png[k] = spectral[c][i];
            k++;
        }
    
    if(io_png_write_f32(argv[3], low_png, (size_t) low_w,
                        (size_t) low_h, (size_t) num_channels) != 0)
    {
        fprintf(stderr, "Error - Failed to save png image %s.\n", argv[3]);
        return EXIT_FAILURE;
    }
    
    // Saving IHS image
    float *high_png = new float[high_whc];
    k = 0;
	for(int c = 0; c < num_channels; c++)
        for(int i = 0; i < high_wh; i++)
        {
            high_png[k] = initial_ihs[c][i];
            k++;
        }
    
    if(io_png_write_f32(argv[4], high_png, (size_t) high_w, (size_t) high_h,
                        (size_t) num_channels) != 0)
    {
        fprintf(stderr, "Error - Failed to save png image %s.\n", argv[4]);
        return EXIT_FAILURE;
    }
    
    // Saving pansharpened image    
    k = 0;
	for(int c = 0; c < num_channels; c++)
        for(int i = 0; i < high_wh; i++)
        {
            high_png[k] = pansharpened[c][i];
            k++;
        }
    
	if(io_png_write_f32(argv[5], high_png, (size_t) high_w, (size_t) high_h,
                        (size_t) num_channels) != 0)
    {
        fprintf(stderr, "Error - Failed to save png image %s.\n", argv[4]);
        return EXIT_FAILURE;
    }
    

	// Free memory
    delete[] alpha;
    delete[] truth;
    free(d_v);
    
    for(int c = 0; c < num_channels; c++)
    {
        delete[] spectral[c];
        delete[] initial_splines[c];
        delete[] initial_ihs[c];
        delete[] quantized_ihs[c];
        delete[] replicated[c];
        delete[] pansharpened[c];
    }
    
    delete[] spectral;
    delete[] initial_splines;
    delete[] initial_ihs;
    delete[] quantized_ihs;
    delete[] replicated;
    delete[] pansharpened;
    delete[] pan;    
    delete[] low_png;
    delete[] high_png;
    

	return EXIT_SUCCESS;
}