#include "../Log.h"
#include "Dither.h"

using namespace Generated;

void DitherGen(DitherData& data, int32_t recursion)
{
    data.Points[0] = {0.0f, 0.0f};
    data.Points[1] = {0.5f, 0.5f};
    data.Points[2] = {0.5f, 0.0f};
    data.Points[3] = {0.0f, 0.5f};
    data.PointCount = 4;
    for (int32_t i = 0; i < data.BrightnessBucketCount; ++i)
        data.BrightnessBuckets[i] = 0;

    for (int32_t recursionLevel = 0; recursionLevel < recursion; ++recursionLevel)
    {
        int32_t startCount = data.PointCount;
        float offset = bx::pow(0.5f, recursionLevel + 1);
        for (int32_t i = 0; i < 4; ++i)
        {
            for (int32_t j = 0; j < startCount; ++j)
            {
                data.Points[data.PointCount] = data.Points[j] + data.Points[i] * offset;
                data.PointCount++;
            }
        }
    }

    uint64_t dotsPerSide = bx::round(bx::pow(2, recursion));
    data.DitherTexDepth = dotsPerSide * dotsPerSide;
    data.DitherTexWH = 16 * dotsPerSide;
    uint64_t texPixelCount = data.DitherTexWH * data.DitherTexWH * data.DitherTexDepth;

    if (BX_COUNTOF(DitherData::DitherTex) < texPixelCount)
    {
        LOG_ERROR("Too many pixels: %llu", texPixelCount);
        return;
    }

    float invRes = 1.0f / data.DitherTexWH;
    for (int32_t z = 0; z < data.DitherTexDepth; ++z)
    {
        int32_t dotCount = z + 1;
        float dotArea = 0.5f / dotCount;
        float dotRadius = bx::sqrt(dotArea / bx::kPi);

        int zOffset = z * data.DitherTexWH * data.DitherTexWH;
        for (int32_t y = 0; y < data.DitherTexWH; ++y)
        {
            int32_t yOffset = y * data.DitherTexWH;
            for (int32_t x = 0; x < data.DitherTexWH; ++x)
            {
                Vec2 point{(x + 0.5f) * invRes, (y + 0.5f) * invRes};
                float dist = bx::kFloatInfinity;
                for (int32_t i = 0; i < dotCount; ++i)
                {
                    Vec2 vec = point - data.Points[i];
                    Vec2 wrappedVec{bx::mod(vec.x + 0.5f, 1.0f) - 0.5f, bx::mod(vec.y + 0.5f, 1.0f) - 0.5f};
                    float curDist = Magnitude(wrappedVec);
                    dist = bx::min(dist, curDist);
                }

                dist = dist / (dotRadius * 2.4f);
                float val = bx::clamp(1.0f - dist, 0.0f, 1.0f);
                data.DitherTex[x + yOffset + zOffset] = Vec4{val, val, val, 1.0f};
                // data.DitherTex[x + yOffset + zOffset] = Vec4{1.0, 0.0f, 0.0f, 1.0f};
                int32_t bucket = bx::clamp(
                    uint32_t(val * DitherData::BrightnessBucketCount), 0, DitherData::BrightnessBucketCount - 1);
                data.BrightnessBuckets[bucket] += 1;
            }
        }
    }

    // Brightness ramp
    int32_t sum = 0;
    for (int32_t i = 0; i < data.BrightnessBucketCount; ++i)
    {
        sum += data.BrightnessBuckets[data.BrightnessBucketCount - 1 - i];
        data.BrightnessRamp[i + 1] = sum / (float)texPixelCount;
    }

    // Upload textures
    if (isValid(data.PreviewTex))
    {
        bgfx::destroy(data.PreviewTex);
        data.PreviewTex = BGFX_INVALID_HANDLE;
    }
    if (isValid(data.FinalTex))
    {
        bgfx::destroy(data.FinalTex);
        data.FinalTex = BGFX_INVALID_HANDLE;
    }
    if (isValid(data.RampTex))
    {
        bgfx::destroy(data.RampTex);
        data.RampTex = BGFX_INVALID_HANDLE;
    }
    const bgfx::Memory* memPreview = bgfx::makeRef(data.DitherTex, texPixelCount * sizeof(Vec4));
    const bgfx::Memory* memFinal = bgfx::makeRef(data.DitherTex, texPixelCount * sizeof(Vec4));
    const bgfx::Memory* memRamp = bgfx::makeRef(data.BrightnessRamp, sizeof(data.BrightnessRamp));
    data.PreviewTex = bgfx::createTexture2D(data.DitherTexWH,
                                            data.DitherTexWH * data.DitherTexDepth,
                                            false,
                                            false,
                                            bgfx::TextureFormat::RGBA32F,
                                            0,
                                            memPreview);
    data.FinalTex = bgfx::createTexture3D(
        data.DitherTexWH, data.DitherTexWH, data.DitherTexDepth, false, bgfx::TextureFormat::RGBA32F, 0, memFinal);
    data.RampTex =
        bgfx::createTexture2D(BX_COUNTOF(data.BrightnessRamp), 1, false, 1, bgfx::TextureFormat::R32F, 0, memRamp);

    if (!isValid(data.DitherSampler))
    {
        data.DitherSampler = bgfx::createUniform("s_ditherSampler", bgfx::UniformType::Sampler);
    }
    if (!isValid(data.RampSampler))
    {
        data.RampSampler = bgfx::createUniform("s_rampSampler", bgfx::UniformType::Sampler);
    }
}