
// Copyright 2019 Shlomi Steinberg
// 

#pragma once

#include <memory>

#include <gli/load.hpp>
#include <gli/texture2d.hpp>
#include <gli/sampler2d.hpp>

#include <opencv2/opencv.hpp>

#include <filesystem>
#include <stdexcept>

#include "glm.hpp"
#include "load2d.hpp"

class envmap {
    using sampler_t = gli::sampler2d<float>;
    struct maps_t {
        sampler_t tex;
        sampler_t blurred_tex;
    };

private:
    std::unique_ptr<maps_t> maps;
    float rot{};
    float dy{}, size{}, aspect{};

private:
    static auto blur(const surface<> &surf, float blur_angle) {
        const auto extend = glm::vec2{ surf.width, surf.height };
        const auto row    = surf.width*4;

        // Gaussian blur
        auto gaussian_kernel_size = glm::max(glm::uvec2{ 1,1 }, 
                                             static_cast<glm::uvec2>(extend * blur_angle / glm::vec2{ 1, 1 } / glm::pi<float>()));
        if (gaussian_kernel_size.x%2 == 0) --gaussian_kernel_size.x;
        if (gaussian_kernel_size.y%2 == 0) --gaussian_kernel_size.y;
        std::cout << "Blur kernel: " << std::to_string(gaussian_kernel_size.x) << " x " << std::to_string(gaussian_kernel_size.y) << std::endl;
		cv::Mat image(surf.width*3, surf.height, CV_8UC4, cv::Scalar(0, 0, 0, 0));
        for (std::size_t y=0;y<surf.height;++y) 
        for (std::size_t x=0;x<surf.width;++x)
        for (std::size_t c=0;c<4;++c) {
            image.at<cv::Vec4b>(x+0*surf.width,y)[c] = surf.data.get()[4*x+y*row+c];
            image.at<cv::Vec4b>(x+1*surf.width,y)[c] = surf.data.get()[4*x+y*row+c];
            image.at<cv::Vec4b>(x+2*surf.width,y)[c] = surf.data.get()[4*x+y*row+c];
        }
        cv::Mat image_blurred;
        cv::GaussianBlur(image, image_blurred, cv::Size(gaussian_kernel_size.x, gaussian_kernel_size.y), 0);

        // Create blurred
        gli::texture2d b{ gli::format::FORMAT_RGBA8_SRGB_PACK32, { surf.width, surf.height }, 1 };
        for (std::size_t y=0;y<surf.height;++y)
        for (std::size_t x=0;x<surf.width;++x)
        for (std::size_t c=0;c<4;++c)
            static_cast<std::uint8_t*>(b.data())[4*x+y*row+c] = image_blurred.at<cv::Vec4b>(x+surf.width,y)[c];

        return b;
    }
    static auto load_envmaps(const std::filesystem::path &path, float blur_angle, float blur_base) {
        const auto surf = load2d(path);
        if (!surf.data || surf.components!=4)
            throw new std::runtime_error("Failed to load envmap");

        auto t = blur(surf, blur_base);
        auto b = blur(surf, blur_angle);
        
        return std::make_unique<maps_t>(maps_t{
                sampler_t(t,
                          gli::wrap::WRAP_REPEAT,
                          gli::filter::FILTER_LINEAR,gli::filter::FILTER_LINEAR),
                sampler_t(b,
                          gli::wrap::WRAP_REPEAT,
                          gli::filter::FILTER_LINEAR,gli::filter::FILTER_LINEAR),
        });
    }

public:
    envmap(const std::filesystem::path &path, 
           float blur_angle, float blur_base, 
           const float dy, const float size, const float aspect,
           float rotation)
        : maps(load_envmaps(path, blur_angle, blur_base)), rot(rotation), dy(dy), size(size), aspect(aspect) {}
    envmap(envmap &&) noexcept = default;
    envmap(const envmap&) noexcept = default;
    envmap &operator=(envmap &&) noexcept = default;
    envmap &operator=(const envmap&) noexcept = default;
    ~envmap() noexcept = default;

    // Directional lookup
    auto operator()(const glm::vec3 &x) const noexcept {
        const auto u = glm::atan(x.x,x.z) + rot;
        const auto v = glm::acos(x.y);
        const auto uv = .5f * glm::vec2{ u,v } / glm::pi<float>() + glm::vec2{ .5,.5 };

        return maps->blurred_tex.texture_lod(uv, 0);
    }
    // Linear slice lookup
    auto operator()(const glm::vec2 &uv) const noexcept {
        const auto p = size * glm::vec2{ .25f * aspect, .5f };
        return maps->tex.texture_lod(uv*p + .5f*p + glm::vec2{ .5f*rot/glm::pi<float>(), dy }, 0);
    }
};
