
// Copyright 2019 Shlomi Steinberg
// 

#include <functional>
#include <tuple>
#include <vector>

#include <uraster.hpp>

#include "speckle_field_renderer.hpp"
#include "tuple_hash.hpp"

#include "light.hpp"
#include "model.hpp"
#include "envmap.hpp"

#include "glm.hpp"


namespace renderer {

static constexpr auto envmap_scale = .9f;

struct vert_vsout_t {
	Eigen::Vector4d p{};
	Eigen::Vector3d wp{};
	Eigen::Vector3d n{};

	const auto& position() const { return p; }
	vert_vsout_t& operator+=(const vert_vsout_t& tp) {
		wp+=tp.wp;
		p+=tp.p;
		n+=tp.n;
		return *this;
	}
	vert_vsout_t& operator-=(const vert_vsout_t& tp) {
		wp-=tp.wp;
		p-=tp.p;
		n-=tp.n;
		return *this;
	}
	vert_vsout_t& operator*=(const double& f) {
		wp*=f;
		p*=f;
        n*=f;
		return *this;
	}
};

struct pixel_t {
    glm::vec4 c0{ .0f };
    glm::vec4 c1{ .0f };
    glm::vec4 c2{ .0f };
	float& depth() { return c0.w; }
};

auto vertex_program(const vert_t& vin,const Eigen::Matrix4d& mvp) {
    vert_vsout_t vout;
    vout.wp={ vin.x,vin.y,vin.z };
    vout.p=mvp*Eigen::Vector4d{ vin.x,vin.y,vin.z,1.0f };
    vout.n={ vin.nx,vin.ny,vin.nz };
    return vout;
}

inline auto tovec(const Eigen::Vector3d &v) noexcept {
	return glm::vec3{ static_cast<float>(v[0]),static_cast<float>(v[1]),static_cast<float>(v[2]) };
}

inline auto material_at_point(const glm::vec3 &v) noexcept {
    return 1.f - glm::clamp((-.35f*v.x+.25f*v.z+(v.y-0.16f))/0.003f,.0f,1.f);
}

pixel_t fragment_program_light(const vert_vsout_t& fsin,
	                           const vert_vsout_t& dVdx,const vert_vsout_t& dVdy,
                         	   std::size_t pixel_idx, int x,int y) {
    pixel_t p{};
	p.c0 = p.c1 = p.c2 = { 1,1,1,0 };
	return p;
}

template <typename SpeckleRenderer>
pixel_t fragment_program(const vert_vsout_t& fsin,
                         const vert_vsout_t& dVdx,const vert_vsout_t& dVdy,
                         std::size_t pixel_idx, int x,int y,
                         const SpeckleRenderer &speckle_renderer,
                         const glm::vec3 &wvp,
                         float scale,
						 float t,
                         const std::vector<light> &lights,
                         const envmap *env,
                         std::size_t materials) {
    const auto wp = tovec(fsin.wp);
    const auto n = tovec(fsin.n.normalized());
    const auto v = glm::normalize(-(wvp-wp));
    const auto mat = materials>1 ? material_at_point(wp) : 0.f;

    // Deal with envmap first (hacky)
    glm::vec4 envmap_final_rgb{};
    const auto reflv = glm::normalize(glm::reflect(v,n)*glm::vec3{1,-1,1});
    for (std::size_t m=0; !!env && m<materials; ++m) {
        const auto w = glm::max(.0f, 1.f - glm::abs(mat - float(m)));
        if (w==.0f) continue;
        const auto& material = speckle_renderer.get_material(m);
        const auto envmap_colour = env[m](reflv) * envmap_scale; 

        static constexpr auto LS = 12;
        for (auto s=0; s<LS; ++s) {
            const auto lambda = glm::mix(.42f,.70f,(s+.5f)/float(LS));
            const auto rgb = glm::mat3{
                     2.0413690, -0.5649464, -0.3446944,
                    -0.9692660,  1.8760108,  0.0415560,
                     0.0134474, -0.1183897,  1.0154096,
                } * XYZ_1931(lambda*1e+3f);
            const auto eta_idx = glm::mix(0.f, (float)(material.size()-1), (lambda-.42f)/(.70f-.42f));
            const auto eta = (1.f-glm::fract(eta_idx)) * material[glm::clamp(0ul,material.size()-1,static_cast<std::size_t>(eta_idx))] +
                             glm::fract(eta_idx) * material[glm::clamp(0ul,material.size()-1,static_cast<std::size_t>(eta_idx+1))];
            
            const auto F = std::norm(fresnel_p(glm::dot(v,n),std::complex<float>(1.),eta)) +
                           std::norm(fresnel_s(glm::dot(v,n),std::complex<float>(1.),eta));
            envmap_final_rgb += w * envmap_colour * glm::vec4{ rgb, 0 } * static_cast<float>(F) / float(LS);
        }
    }

    pixel_t p{};
    for (const auto &light : lights) {
        const auto &L = light.position(t)/scale;
        const auto sr = light.solid_angle(t,wp*scale);
        const auto coherence_length = light.coherence_length();
        const auto coherence_radius = light.coherence_radius(sr);
        const auto coherence_area   = glm::pi<float>()*sqr(coherence_radius);
        const auto l = glm::normalize(-(L-wp));
        const auto dot_ln = glm::dot(l,n);

        // Generate fragments
        const auto Is = speckle_renderer(pixel_idx,x,y,
                                         wp, n, v, l,
                                         tovec(dVdx.wp),tovec(dVdy.wp),
                                         tovec(dVdx.n),tovec(dVdy.n),
                                         wvp, L,
										 light.get_illuminant(),
                                         dVdx.wp.norm()*scale, dVdy.wp.norm()*scale,
                                         coherence_area, coherence_length,
                                         light.is_area(), sr, light.is_area() ? light.vertices(t,1.f/scale) : std::array<glm::vec3, 4>{},
                                         mat);
        const auto mean = glm::max(glm::vec3{ 0 }, static_cast<glm::vec3>(Is.first)  * light.get_intensity(t, 1.f/scale));
        const auto fluc = static_cast<glm::vec3>(Is.second) * light.get_intensity(t, 1.f/scale);

        p.c2 += glm::vec4{ mean, .0f };
        p.c1 += glm::vec4{ fluc, .0f };
    }
    // const auto v_dot_n = glm::max(glm::dot(v,n),.0f);
    // p.c1 *= v_dot_n;
    // p.c2 *= v_dot_n;
    p.c2 += envmap_final_rgb;
	p.c0 = p.c1+p.c2;

	return p;
}

using render_output_t = uraster::Framebuffer<pixel_t>;

template <typename SpeckleRenderer>
auto do_render(const std::size_t res_x, const std::size_t res_y, const SpeckleRenderer &speckle_renderer, 
               const model_t &model, const glm::mat4 &modelmatrix, const glm::mat4 &viewmatrix,
               const std::vector<light> &lights, const envmap *env, std::size_t materials,
			   float t,
			   float scale) {
    auto vc = std::make_unique<vert_vsout_t[]>(model.verts.size());
	const auto* vbb=model.verts.data();
	const auto* vbe=model.verts.data()+model.verts.size();
	const auto* ibb=model.inds.data();
	const auto* ibe=model.inds.data()+model.inds.size();

	vert_vsout_t* vcb=vc.get();
	vert_vsout_t* vce=vc.get()+model.verts.size();

    Eigen::Matrix4d modelmat, viewmat;
    for (std::size_t x=0;x<4;++x)
    for (std::size_t y=0;y<4;++y) {
        modelmat(x,y)=modelmatrix[x][y];
        viewmat(x,y)=viewmatrix[x][y];
    }

	auto tp = std::make_unique<render_output_t>(res_x,res_y);
    const auto camera_matrix=viewmat*modelmat;

    const Eigen::Vector4d wvp4 = camera_matrix.inverse() * Eigen::Vector4d{0,0,0,1};
    const auto wvp = tovec({ wvp4[0]/wvp4[3], wvp4[1]/wvp4[3], wvp4[2]/wvp4[3] });

    // Render envmap
    if (env) {
        for (std::size_t x=0;x<res_x;++x)
        for (std::size_t y=0;y<res_y;++y) {
            const auto refl = env[0](glm::vec2{ x,y } / glm::vec2{ res_x,res_y }) * .75f;
            (*tp)(x,y).c0 = (*tp)(x,y).c2 = glm::vec4{ refl.x, refl.y, refl.z, .0 };
        }
    }

	// Render lights
	// for (const auto &l : lights) {
	// 	if (!l.should_render()) continue;

	// 	vert_vsout_t vos[4];
	// 	const auto vs = l.vertices(t,1/scale);
	// 	vert_t vv[4];
	// 	for (int i=0;i<4;++i) {
	// 		vv[i].x = vs[i].x;
	// 		vv[i].y = vs[i].y;
	// 		vv[i].z = vs[i].z;
	// 	}
	// 	std::size_t ind[6] = { 0,1,2,2,1,3 };

	// 	uraster::draw(*tp,
	// 		vv,vv+4,
	// 		ind,ind+6,
	// 		vos,vos+4,
	// 		std::bind(vertex_program, std::placeholders::_1, Eigen::Matrix4d::Identity()),
	// 		fragment_program_light);
	// }

	// Render
    uraster::draw(*tp,
        vbb,vbe,
        ibb,ibe,
        vcb,vce,
        std::bind(vertex_program, std::placeholders::_1, camera_matrix),
        std::bind(fragment_program<SpeckleRenderer>, 
                  std::placeholders::_1,
                  std::placeholders::_2,
                  std::placeholders::_3,
                  std::placeholders::_4,
                  std::placeholders::_5,
                  std::placeholders::_6,
                  std::ref(speckle_renderer),
                  std::ref(wvp),
                  scale,
				  t,
                  std::ref(lights),
                  env, 
                  materials)
    );

	return tp;
}

}
