#include <Bitmap/Bitmap.hpp>
#include <Log/Log.hpp>
#include <optional>
#include <random>

Bitmap::Bitmap(glm::uvec2 size, const uint8_t *rgba) : size(size)
{
    rgbaPixels.resize(4 * size.x * size.y);

    if (rgba != nullptr)
    {
        std::memcpy(rgbaPixels.data(), rgba, 4 * size.x * size.y);
    }
}

std::optional<Bitmap> Bitmap::fromFile(std::string path)
{
    int width, height, comp;
    unsigned char *pixels = stbi_load(path.c_str(), &width, &height, &comp, STBI_rgb_alpha);
    
    if (pixels == nullptr)
    {
        return std::nullopt;
    }

    if (width == 0 || height == 0)
    {
        crash(formatString("width == 0 || height == 0 for '$+'!", path));
    }

    Bitmap result = Bitmap(glm::uvec2(width, height), pixels);
    stbi_image_free(pixels);
    return result;
}

Bitmap Bitmap::fromFileOrDefault(std::string path)
{
    std::optional<Bitmap> optBitmap = fromFile(path);
    if (optBitmap.has_value())
    {
        return *optBitmap;
    }

    Bitmap bitmap(glm::uvec2(8, 8));
    glm::uvec2 pos;

    for (pos.y = 0; pos.y < 8; pos.y++)
    {
        for (pos.x = 0; pos.x < 8; pos.x++)
        {
            if ((pos.x >> 2) ^ (pos.y >> 2))
            {
                bitmap.setColor(pos, 0xFF, 0x00, 0xFF, 0xFF);
            }
            else
            {
                bitmap.setColor(pos, 0xFF, 0xFF, 0xFF, 0xFF);
            }
        }
    }

    Log::warn(formatString("Bitmap::fromFileOrDefault: Failed to load '$+' - the default bitmap will be used.", path));
    return bitmap;
}

glm::vec4 Bitmap::decodeXYZW8(uint32_t rgba8)
{
    return glm::vec4(
        (rgba8 >> 24) & 0xFF,
        (rgba8 >> 16) & 0xFF,
        (rgba8 >> 8) & 0xFF,
        rgba8 & 0xFF
    ) / 255.0f;
}

uint32_t Bitmap::encodeXYZW8(glm::vec4 v)
{
    uint32_t x = std::min(255, std::max(0, static_cast<int32_t>(v.x * 255.0f)));
    uint32_t y = std::min(255, std::max(0, static_cast<int32_t>(v.y * 255.0f)));
    uint32_t z = std::min(255, std::max(0, static_cast<int32_t>(v.z * 255.0f)));
    uint32_t w = std::min(255, std::max(0, static_cast<int32_t>(v.w * 255.0f)));

    return (x << 24) | (y << 16) | (z << 8) | w;
}

glm::vec3 Bitmap::sRGBToLinear(glm::vec3 srgb)
{
    return glm::mix(srgb / 12.92f, glm::pow((srgb + 0.055f) / 1.055f, glm::vec3(2.4f)), glm::step(glm::vec3(0.04045f), srgb));
}

glm::vec3 Bitmap::linearToSRGB(glm::vec3 linearRgb)
{
    glm::vec3 sRgb;
    for (int i = 0; i < 3; i++)
    {
        float c = linearRgb[i];
        if (c <= 0.0031308f)
        {
            sRgb[i] = 12.92f * c;
        }
        else
        {
            sRgb[i] = 1.055f * std::pow(c, 1.0f / 2.4f) - 0.055f;
        }
    }

    return sRgb;
}

glm::vec3 Bitmap::rgb2hsv(glm::vec3 c)
{
    glm::vec4 K = glm::vec4(0.0, -1.0 / 3.0, 2.0 / 3.0, -1.0);
    glm::vec4 p = mix(glm::vec4(c.b, c.g, K.w, K.z), glm::vec4(c.g, c.b, K.x, K.y), glm::step(c.b, c.g));
    glm::vec4 q = mix(glm::vec4(p.x, p.y, p.w, c.r), glm::vec4(c.r, p.y, p.z, p.x), glm::step(p.x, c.r));

    float d = q.x - std::min(q.w, q.y);
    float e = 1.0e-10;
    return glm::vec3(abs(q.z + (q.w - q.y) / (6.0 * d + e)), d / (q.x + e), q.x);
}

glm::vec3 Bitmap::hsv2rgb(glm::vec3 c)
{
    glm::vec4 K = glm::vec4(1.0, 2.0 / 3.0, 1.0 / 3.0, 3.0);
    glm::vec3 p = glm::abs(glm::fract(glm::vec3(c.x + K.x, c.x + K.y, c.x + K.z)) * 6.0f - glm::vec3(K.w, K.w, K.w));
    return c.z * mix(glm::vec3(K.x), glm::clamp(p - glm::vec3(K.x), 0.0f, 1.0f), c.y);
}

glm::uvec2 Bitmap::getSize()
{
    return size;
}

uint32_t Bitmap::getRGBA8(glm::uvec2 pos)
{
    size_t offset = 4 * (pos.y * size.x + pos.x);
    uint32_t result = 0;

    result |= ((uint32_t) rgbaPixels[offset + 0]) << 24;
    result |= ((uint32_t) rgbaPixels[offset + 1]) << 16;
    result |= ((uint32_t) rgbaPixels[offset + 2]) << 8;
    result |= ((uint32_t) rgbaPixels[offset + 3]);

    return result;
}

RGBABytes Bitmap::getColor(glm::uvec2 pos)
{
    RGBABytes color;
    std::memcpy(&color, &rgbaPixels[4 * (pos.y * size.x + pos.x)], 4);
    return color;
}

void Bitmap::setColor(glm::uvec2 pos, RGBABytes color)
{
    uint8_t *put = &rgbaPixels[4 * (pos.y * size.x + pos.x)];
    put[0] = color.red;
    put[1] = color.green;
    put[2] = color.blue;
    put[3] = color.alpha;
}

void Bitmap::setColor(glm::uvec2 pos, uint8_t r, uint8_t g, uint8_t b, uint8_t a)
{
    uint8_t *put = &rgbaPixels[4 * (pos.y * size.x + pos.x)];
    put[0] = r;
    put[1] = g;
    put[2] = b;
    put[3] = a;
}

glm::vec4 Bitmap::getLinearColor(glm::uvec2 pos)
{
    glm::vec4 xyzw = decodeXYZW8(getRGBA8(pos));
    return glm::vec4(sRGBToLinear(glm::vec3(xyzw)), xyzw.a);
}

void Bitmap::setSRGBColor(glm::uvec2 pos, glm::vec4 srgba)
{
    uint8_t x = std::min(255, std::max(0, static_cast<int32_t>(srgba.x * 255.0f)));
    uint8_t y = std::min(255, std::max(0, static_cast<int32_t>(srgba.y * 255.0f)));
    uint8_t z = std::min(255, std::max(0, static_cast<int32_t>(srgba.z * 255.0f)));
    uint8_t w = std::min(255, std::max(0, static_cast<int32_t>(srgba.w * 255.0f)));

    setColor(pos, x, y, z, w);
}

void Bitmap::setLinearColor(glm::uvec2 pos, glm::vec4 rgba)
{
    glm::vec4 srgba = glm::vec4(linearToSRGB(glm::vec3(rgba)), rgba.a);
    setSRGBColor(pos, srgba);
}

void Bitmap::writeToPNG(std::string path)
{
    stbi_write_png(path.c_str(), size.x, size.y, 4, rgbaPixels.data(), 4 * size.x);
}

void Bitmap::applyNoiseHSV(Bitmap &biasMap)
{
    static std::random_device rd;
    static std::mt19937 eng(rd());
    static std::uniform_real_distribution<float> distr(-1.0f, 1.0f);

    glm::uvec2 pos;
    for (pos.y = 0; pos.y < size.y; pos.y++)
    {
        for (pos.x = 0; pos.x < size.x; pos.x++)
        {
            glm::vec4 rgba = getLinearColor(pos);
            glm::vec3 hsv = rgb2hsv(glm::vec3(rgba));

            glm::vec3 noise = glm::vec3(distr(eng), distr(eng), distr(eng));
            glm::vec3 bias = glm::vec3(decodeXYZW8(biasMap.getRGBA8(pos % biasMap.getSize())));

            hsv += bias * noise;

            rgba = glm::vec4(hsv2rgb(hsv), rgba.a);
            setLinearColor(pos, rgba);
        }
    }
}

uint8_t* Bitmap::data()
{
    return rgbaPixels.data();
}