#include <Image.hpp>

Image::Image(int width_, int height_, int depth_) : width(width_), height(height_), depth(depth_) {
    data = new unsigned char[width * height * depth];
}

Image::Image(Image const& other) : width(other.width), height(other.height), depth(other.depth) {
    data = new unsigned char[width * height * depth];
    std::memcpy(data, other.data, width * height * depth);
}

Image::Image(std::string const& filename) {
    // Load image data using stb_image.h
    unsigned char* image_data = stbi_load(filename.c_str(), &width, &height, &depth, 0);

    // Allocate memory using new[] and copy the data from image.data. This avoids allocating memory with malloc and freeing it with delete.
    if (image_data != nullptr) {
        data = new unsigned char[width * height * depth];
        std::memcpy(data, image_data, width * height * depth);
        stbi_image_free(image_data);
    } else {    // If the image loading process has failed
        width = 1;
        height = 1;
        depth = 1;
        data = new unsigned char[width * height * depth];
        data[0] = 0;
    }
}

Image::~Image() {
    delete[] data;
}

Image & Image::operator=(Image const& other) {
    // Delete current data
    delete[] data;

    // Copy image attributes
    width = other.width;
    height = other.height;
    depth = other.depth;

    // Allocate new data and copy the data from other.data.
    data = new unsigned char[width * height * depth];
    std::memcpy(data, other.data, width * height * depth);
    return *this;
}

int Image::Save(std::string const& filename) const {
    if(filename.size() > 4 && filename.substr(filename.size() - 4) == ".png")
        return stbi_write_png(filename.c_str(), width, height, depth, (const void *)data, 0);
    if(filename.size() > 4 && filename.substr(filename.size() - 4) == ".jpg")
        return stbi_write_jpg(filename.c_str(), width, height, depth, (const void *)data, 100);
    if(filename.size() > 4 && filename.substr(filename.size() - 4) == ".bmp")
        return stbi_write_bmp(filename.c_str(), width, height, depth, (const void *)data);
    if(filename.size() > 4 && filename.substr(filename.size() - 4) == ".tga")
        return stbi_write_tga(filename.c_str(), width, height, depth, (const void *)data);
    return -1;// If the file extension is not supported
}

int Image::GetWidth() const {
    return width;
}

int Image::GetHeight() const {
    return height;
}

int Image::GetDepth() const {
    return depth;
}

Image & Image::Fill(unsigned char value) {
    std::memset(data, value, width * height * depth);
    return *this;
}

Image & Image::Fill(Eigen::Vector3i const& rgb) {
    for(int i = 0 ; i < height ; i++)
        for(int j = 0 ; j < width ; j++)
            SetPixel(i, j, rgb);
    return *this;
}

unsigned char Image::GetPixelValue(int i, int j, int c) const {
    return data[i*width*depth + j*depth + c];
}

void Image::SetPixelValue(int i, int j, int c, unsigned char value) {
    data[i*width*depth + j*depth + c] = value;
}

unsigned char & Image::PixelValue(int i, int j, int c) {
    return data[i*width*depth + j*depth + c];
}

Eigen::Vector3i Image::GetPixel(int i, int j) const {
    int index = i*width*depth + j*depth;
    return Eigen::Vector3i((int)data[index], (int)data[index + 1], (int)data[index + 2]);
}

Eigen::Vector3i Image::GetPixel(Eigen::Vector2i const& pos) const {
    return GetPixel(pos[0], pos[1]);
}

Eigen::Vector3i Image::GetPixelInterp(Eigen::Vector2d const& pos, InterpMethod const& interp_method) const {
    if(interp_method == Image::InterpMethod::BILINEAR) {
        int i = (int)std::floor(pos[0]);
        int j = (int)std::floor(pos[1]);
        double u = pos[0] - i;
        double v = pos[1] - j;
        int i_plus_1 = (i + 1) % height;
        int j_plus_1 = (j + 1) % width;
        Eigen::Vector3i rgb = ((1 - u) * (1 - v) * GetPixel(i, j).cast<double>() + u * (1 - v) * GetPixel(i_plus_1, j).cast<double>() + (1 - u) * v * GetPixel(i, j_plus_1).cast<double>() + u * v * GetPixel(i_plus_1, j_plus_1).cast<double>()).cast<int>();
        return rgb;
    }
    else    // By default, use nearest neighbor interpolation
        return GetPixel((int)std::round(pos[0]), (int)std::round(pos[1]));
}

Image & Image::SetPixel(int i, int j, Eigen::Vector3i const& rgb) {
    SetPixelValue(i, j, 0, (unsigned char)rgb[0]);
    SetPixelValue(i, j, 1, (unsigned char)rgb[1]);
    SetPixelValue(i, j, 2, (unsigned char)rgb[2]);
    return *this;
}

Image & Image::SetPixel(Eigen::Vector2i const& pos, Eigen::Vector3i const& rgb) {
    SetPixel(pos[0], pos[1], rgb);
    return *this;
}

unsigned char Image::operator[](int i) const {
    return data[i];
}

Image Image::Downsampled(int factor) const {
    factor = std::max(factor, 1);
    Image downsampled(width / factor, height / factor, depth);
    for(int i = 0 ; i < downsampled.GetHeight() ; i++)
        for(int j = 0 ; j < downsampled.GetWidth() ; j++)
            downsampled.SetPixel(i, j, GetPixel(i * factor, j * factor));
    return downsampled;
}

Image Image::Grayscale() const {
    Image grayscale(width, height, 1);
    for(int i = 0 ; i < height ; i++)
        for(int j = 0 ; j < width ; j++) {
            Eigen::Vector3i rgb = GetPixel(i, j);
            grayscale.SetPixelValue(i, j, 0, ((int)rgb[0] + (int)rgb[1] + (int)rgb[2]) / 3);
        }
    return grayscale;
}

Image Image::LumaREC709() const {
    Image grayscale(width, height, 1);
    for(int i = 0 ; i < height ; i++)
        for(int j = 0 ; j < width ; j++) {
            Eigen::Vector3i rgb = GetPixel(i, j);
            grayscale.SetPixelValue(i, j, 0, ((double)rgb[0]*0.2126 + (double)rgb[1]*0.7152 + (double)rgb[2]*0.0722));
        }
    return grayscale;
}

Image & Image::HistogramNormalize(int downsampling_factor) {
    // First, compute the histogram
    std::vector<uint64_t> hist = ComputeHistogram(downsampling_factor);
    // Then, compute the cumulative histogram
    std::vector<uint64_t> cumul_hist = CumulativeHistogram(hist);
    // Finally, normalize the image
    uint64_t max_cum_hist_val = cumul_hist[cumul_hist.size()-1];
    for(int i = 0 ; i < height ; i++)
        for(int j = 0 ; j < width ; j++) {
            Eigen::Vector3i rgb = GetPixel(i, j);
            SetPixel(i, j, Eigen::Vector3i(cumul_hist[rgb[0]] * 255 / max_cum_hist_val, cumul_hist[rgb[1]] * 255 / max_cum_hist_val, cumul_hist[rgb[2]] * 255 / max_cum_hist_val));
        }
    return *this;
}

Image Image::HistogramNormalized(int downsampling_factor) const {
    Image normalized(*this);
    normalized.HistogramNormalize(downsampling_factor);
    return normalized;
}

Image & Image::Normalize() {
    // First, compute the min and max values using all channels
    unsigned char min_val = 255, max_val = 0;
    for(int i = 0 ; i < width*height*depth ; i++) {
        if(data[i] < min_val)
            min_val = data[i];
        else if(data[i] > max_val)
            max_val = data[i];
    }

    // Then, normalize the image
    for(int i = 0 ; i < width*height*depth ; i++)
        data[i] = ((int)data[i] - min_val) * 255 / (max_val - min_val);
    
    return *this;
}

Image Image::Normalized() const {
    Image normalized(*this);
    normalized.Normalize();
    return normalized;
}

Image Image::Colorized(unsigned char const* colormap) const {
    if (depth != 1)
        return *this;
    Image colorized(width, height, 3);
    for(int i = 0 ; i < height ; i++)
        for(int j = 0 ; j < width ; j++)
            for(int c = 0 ; c < 3 ; c++)
                colorized.SetPixelValue(i, j, c, colormap[3*GetPixelValue(i, j) + c]);
    return colorized;
}

std::vector<uint64_t> Image::ComputeHistogram(int downsampling_factor) const {
    std::vector<uint64_t> hist(256, 0);
    for(int i = 0 ; i < width*height*depth ; i += downsampling_factor)
        hist[data[i]]++;
    return hist;
}

std::vector<uint64_t> Image::CumulativeHistogram(std::vector<uint64_t> const& histogram) {
    std::vector<uint64_t> cumul_hist(histogram.size(), 0);
    cumul_hist[0] = histogram[0];
    for(unsigned int i = 1 ; i < histogram.size() ; i++)
        cumul_hist[i] = cumul_hist[i - 1] + histogram[i];
    return cumul_hist;
}