#include <iostream>
#include <fstream>
#include <algorithm>
#include <stdlib.h>
#include <time.h>
#include <math.h>

#include "haader.h"

using namespace std;
using namespace haader;


int ResponseFunction::lookup(double value) const {
    // + 0.5 for rounding to nearest integer
    int bin = (double)m_lookup_table.size() * (value - m_min_value) * m_recip_range + 0.5;

    if (bin < 0) {
        return 0;
    } else if ((unsigned int)bin >= m_lookup_table.size()) {
        return 255;
    } else {
        return m_lookup_table[bin];
    }
}

double InverseResponseFunction::lookup(int index) const {
    return m_lookup_table[index];
}

int InverseResponseFunction::inverse_lookup(double value) const {
    if (value <= m_lookup_table[0]) {
        return 0;
    }
    if (value >= m_lookup_table[255]) {
        return 255;
    }

    for (unsigned int i = 1; i < 256; i++) {
        if (value > m_lookup_table[i - 1] && value <= m_lookup_table[i]) {
            double diff_a = value - m_lookup_table[i - 1];
            double diff_b = m_lookup_table[i] - value;
            if (diff_a < diff_b) {
                return i - 1;
            } else {
                return i;
            }
        }
    }

    // should not happen
    return 0;
}

ResponseFunction InverseResponseFunction::to_response_function(int bins) {
    if (!is_monotonic()) {
        cerr << "InverseResponseFunction::to_response_function: not monotonic." << endl;
        return ResponseFunction();
    }

    double min_value = m_lookup_table[0];
    double max_value = m_lookup_table[m_lookup_table.size() - 1];

    ResponseFunction rf;
    rf.m_min_value = min_value;
    rf.m_max_value = max_value;
    rf.m_recip_range = 1.0 / (max_value - min_value);
    rf.m_lookup_table.resize(bins);


    //TODO reduce quadratic complexity to linear complexity
    for (int i = 0; i < bins; i++) {
        double y = min_value + ((max_value - min_value) * i) / (double)(bins - 1.0);
        rf.m_lookup_table[i] = inverse_lookup(y);
    }

    return rf;
}

bool InverseResponseFunction::repair() {

    // Fill missing values
    fill_zeros();

    if (is_monotonic()) {
        return true;
    }

    for (unsigned int i = 0; i < 128; i++) {
        cout << "low pass filter " << i + 1 << endl;
        low_pass_filter();
        if (is_monotonic()) {
            // repair success

            // offset values so that m_lookup_table[128] is 0.0 again
            double offset = -m_lookup_table[128];
            for (unsigned int i = 0; i < 256; i++) {
                m_lookup_table[i] += offset;
            }

            return true;
        }
    }

    // repair failed
    return false;
}

void InverseResponseFunction::fill_zeros() {
    //TODO test correctness

    // fill at front
    if (m_lookup_table[0] == 0) {
        for (unsigned int i = 0; i < 128; i++) {
            if (m_lookup_table[i] != 0.0) {
                for (unsigned int k = 0; k <= i; k++) {
                    m_lookup_table[k] = m_lookup_table[i];
                }
                break;
            }
        }
    }

    // fill at back
    if (m_lookup_table[255] == 0) {
        for (unsigned int i = 255; i > 128; i--) {
            if (m_lookup_table[i] != 0.0) {
                for (unsigned int k = 255; k >= i; k--) {
                    m_lookup_table[k] = m_lookup_table[i];
                }
                break;
            }
        }
    }

    {
        bool inside_zeros = false;
        int start_index = 0;
        for (unsigned int i = 1; i < 256; i++) {
            if (m_lookup_table[i] == 0.0 && !inside_zeros && i != 128) {
                start_index = i;
                inside_zeros = true;
            } else if ((m_lookup_table[i] != 0.0 || i == 128) && inside_zeros) {
                inside_zeros = false;
                double a_val = m_lookup_table[start_index - 1];
                double b_val = m_lookup_table[i];
                for (unsigned int k = start_index; k < i; k++) {
                    double t = (double)(k - start_index + 1) / (double)(i - start_index + 1);
                    m_lookup_table[k] = (1.0 - t) * a_val + t * b_val;
                }
            }
        }
    }
}

void InverseResponseFunction::low_pass_filter() {
    vector<double> newvals;
    newvals.resize(256);

    newvals[0] = m_lookup_table[0];
    newvals[255] = m_lookup_table[255];

    for (unsigned int i = 1; i < 255; i++) {
        newvals[i] = 0.25 * m_lookup_table[i - 1] +  0.5 * m_lookup_table[i] + 0.25 * m_lookup_table[i + 1];
    }

    for (unsigned int i = 0; i < 256; i++) {
        m_lookup_table[i] = newvals[i];
    }
}

bool InverseResponseFunction::is_monotonic() {
    for (unsigned int i = 1; i < 256; i++) {
        if (m_lookup_table[i - 1] > m_lookup_table[i]) {
            return false;
        }
    }
    return true;
}


HdrImage::HdrImage():
    m_width(0),
    m_height(0),
    m_components(0)
{
}

Image HdrImage::get_log_image() {
    if (m_image_data.size() == 0) {
        return Image();
    }

    Image output;
    output.m_image_data.resize(m_image_data.size());
    output.m_width = m_width;
    output.m_height = m_height;
    output.m_components = m_components;

    double min_value = m_image_data[0];
    double max_value = m_image_data[0];

    unsigned int size = m_image_data.size();
    for (unsigned int i = 1; i < size; i++) {
        double val = m_image_data[i];
        if (val < min_value) {
            min_value = val;
        }
        if (val > max_value) {
            max_value = val;
        }
    }

    double scale = 256.0 / (max_value - min_value);

    for (unsigned int i = 0; i < size; i++) {
        output.m_image_data[i] = (int)((m_image_data[i] - min_value) * scale);
    }

    return output;
}

Image HdrImage::expose(double exposure_time, const ResponseFunction &rf, double compression) {
    if (m_image_data.size() == 0) {
        return Image();
    }

    Image output;
    output.m_image_data.resize(m_image_data.size());
    output.m_width = m_width;
    output.m_height = m_height;
    output.m_components = m_components;

    double log_exposure = log(exposure_time);

    unsigned int size = m_image_data.size();
    for (unsigned int i = 0; i < size; i++) {
        output.m_image_data[i] = rf.lookup((m_image_data[i] + log_exposure) * compression);
    }

    return output;
}


HdrImageStack::HdrImageStack() {
}

bool HdrImageStack::read_from_files(char * const* file_paths, int number_of_paths) {
    m_images.clear();

    if (number_of_paths < 0) {
        cerr << "HdrImageStack::read_from_files: number_of_paths is negative" << endl;
        return false;
    }

    for (int i = 0; i < number_of_paths; i++) {
        m_images.push_back(Image());
        if (m_images[i].read_from_file(file_paths[i])) {
            cout << "Read \"" << file_paths[i] << "\"" << endl;
            if (!m_images[0].has_equal_dimensions(m_images[i])) {
                cerr << "HdrImageStack::read_from_files: Dimensions do not match (\""
                     << file_paths[i]
                     << "\")"
                     << endl;
                return false;
            }
            if (m_images[i].get_exposure_time() == 0.0) {
                cerr << "HdrImageStack::read_from_files: Could not get exposure time (\""
                     << file_paths[i]
                     << "\")"
                     << endl;
                return false;
            }
        } else {
            return false;
        }
    }

    sort_by_exposure();

    return true;
}

bool HdrImageStack::get_average_image_slow(Image &output) {
    if (m_images.size() == 0) {
        return false;
    }

    output.m_image_data.resize(m_images[0].m_image_data.size());
    output.m_width = m_images[0].m_width;
    output.m_height = m_images[0].m_height;
    output.m_components = m_images[0].m_components;

    unsigned int bytes_size = m_images[0].m_image_data.size();
    unsigned int images_size = m_images.size();
    for (unsigned int i = 0; i < bytes_size; i++) {
        int v = 0;
        for (unsigned int k = 0; k < images_size; k++) {
            v += m_images[k].m_image_data[i];
        }
        output.m_image_data[i] = v / (int)images_size;
    }

    return true;
}

bool HdrImageStack::get_average_image(Image &output) {
    if (m_images.size() == 0) {
        return false;
    }

    output.m_image_data.resize(m_images[0].m_image_data.size());
    output.m_width = m_images[0].m_width;
    output.m_height = m_images[0].m_height;
    output.m_components = m_images[0].m_components;

    unsigned int bytes_size = m_images[0].m_image_data.size();
    unsigned int images_size = m_images.size();

    const unsigned int buf_size = 128;
    int buf[buf_size];

    for (unsigned int i = 0; i + buf_size <= bytes_size; i += buf_size) {
        for (unsigned int m = 0; m < buf_size; m++) {
            buf[m] = m_images[0].m_image_data[i + m];
        }
        for (unsigned int k = 1; k < images_size; k++) {
            for (unsigned int m = 0; m < buf_size; m++) {
                buf[m] += m_images[k].m_image_data[i + m];
            }
        }
        for (unsigned int m = 0; m < buf_size; m++) {
            output.m_image_data[i + m] = buf[m] / images_size;
        }
    }

    // TODO test this with an image that has prime numbered dimensions
    unsigned int rest = bytes_size % buf_size;
    for (unsigned int i = bytes_size - rest; i < bytes_size; i++) {
        int v = 0;
        for (unsigned int k = 0; k < images_size; k++) {
            v += m_images[k].m_image_data[i];
        }
        output.m_image_data[i] = v / (int)images_size;
    }

    return true;
}

bool HdrImageStack::get_inverse_response_function(InverseResponseFunction &output, unsigned int num_samples) {
    srand(time(0));

    unsigned int images_size = m_images.size();
    unsigned int c = m_images[0].m_components;
    unsigned int width = m_images[0].m_width;
    unsigned int height = m_images[0].m_height;

    vector<vector<double> > resp;
    resp.resize(256);

    for (unsigned int i = 0; i < num_samples; i++) {
        vector<int> values;
        vector<double> exposure_times;
        unsigned int x = rand() % width;
        unsigned int y = rand() % height;
        unsigned int index = (y * width + x) * c;
        for (unsigned int k = 0; k < images_size; k++) {
            int val = m_images[k].m_image_data[index];
            if (val == 255) {
                break;
            }

            values.push_back(val);
            exposure_times.push_back(log(m_images[k].get_exposure_time()));
        }

        if (values.size() > 0 && values[0] < 128 && values[values.size() - 1] > 128) {
            double offset = 0.0;
            for (unsigned int k = 0; k < values.size(); k++) {
                if (values[k] == 128) {
                    offset = -exposure_times[k];
                    break;
                } else if (values[k] > 128) {
                    double delta_val = values[k] - values[k -1];
                    double t = (double)(128 - values[k - 1]) / delta_val;
                    offset = -((1.0 - t) * exposure_times[k - 1] + t * exposure_times[k]);
                    break;
                }
            }

            for (unsigned int k = 0; k < values.size(); k++) {
                resp[values[k]].push_back(exposure_times[k] + offset);
            }
        }
    }

    output.m_lookup_table.resize(256);

    for (unsigned int i = 0; i < 256; i++) {
        double avg = 0.0;
        for (unsigned int k = 0; k < resp[i].size(); k++) {
            avg += resp[i][k];
        }
        if (resp[i].size() > 0) {
            output.m_lookup_table[i] = avg / (double)resp[i].size();
        } else {
            output.m_lookup_table[i] = 0.0;
        }
    }

    for (unsigned int i = 0; i < 256; i++) {
        double avg = 0.0;
        for (unsigned int k = 0; k < resp[i].size(); k++) {
            avg += resp[i][k];
        }
        if (resp[i].size() > 0) {
            output.m_lookup_table[i] = avg / (double)resp[i].size();
        } else {
            output.m_lookup_table[i] = 0.0;
        }
    }

    if (output.repair()) {
        cout << "Inverse response function" << endl;
        cout << endl;
        for (unsigned int i = 0; i < 256; i++) {
            cout << i << "," << output.m_lookup_table[i] << endl;
        }
        cout << endl;

        return true;
    } else {
        return false;
    }
}

HdrImage HdrImageStack::get_hdr_image(const InverseResponseFunction &invRespFunc) {
    if (m_images.size() == 0) {
        return HdrImage();
    }

    HdrImage output;

    output.m_image_data.resize(m_images[0].m_image_data.size());
    output.m_width = m_images[0].m_width;
    output.m_height = m_images[0].m_height;
    output.m_components = m_images[0].m_components;

    unsigned int bytes_size = m_images[0].m_image_data.size();
    unsigned int images_size = m_images.size();

    // precompute log of exposure times
    vector<double> log_exposures;
    log_exposures.resize(images_size);
    for (unsigned int k = 0; k < images_size; k++) {
        log_exposures[k] = log(m_images[k].get_exposure_time());
    }

    for (unsigned int i = 0; i < bytes_size; i++) {

        double scale_sum = 0.0;
        double output_val = 0.0;

        for (unsigned int k = 0; k < images_size; k++) {
            int val = m_images[k].m_image_data[i];
            double scale = 1.0 - fabs((double)(128.0 - val) * (1.0 / 128.0));
            output_val += scale * (invRespFunc.lookup(val) - log_exposures[k]);
            scale_sum += scale;
        }

        output.m_image_data[i] = output_val / scale_sum;
    }

    return output;
}

static bool compare_image_by_exposure(const Image &a, const Image &b) {
    return a.get_exposure_time() < b.get_exposure_time();
}

void HdrImageStack::sort_by_exposure() {
    std::sort(m_images.begin(), m_images.end(), compare_image_by_exposure);
}