#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() {
    delete[] data;
}

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;
    }
}

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;
        Eigen::Vector3i rgb = ((1 - u) * (1 - v) * GetPixel(i, j).cast<double>() + u * (1 - v) * GetPixel(i + 1, j).cast<double>() + (1 - u) * v * GetPixel(i, j + 1).cast<double>() + u * v * GetPixel(i + 1, j + 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];
}