
// Copyright 2019 Shlomi Steinberg
// 

#pragma once

#include <memory>
#include <cstddef>
#include <mutex>
#include <vector>
#include <array>
#include <algorithm>
#include <execution>
#include <numeric>
#include <iterator>
#include <atomic>
#include <type_traits>
#include <new>
#include <cmath>

#include "amplitude.hpp"
#include "spectrum.hpp"
#include "print_progress.hpp"
#include "create_array.hpp"
#include "write2d.hpp"

#include "orthogonal_frame.hpp"
#include "normal_distribution.hpp"

#include "fresnel.hpp"
#include "illuminant.hpp"

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wnarrowing"
#include "avx_mathfun.h"
#pragma GCC diagnostic pop

#include <fftw3.h>
#include <fftw3cxx.hh>


template <typename Rand, typename Ctx, typename Sampler, bool white_light, bool render_spatial=true>
class speckle_field_renderer {
private:
	static constexpr auto M = Ctx::M;
    static constexpr auto SS = Ctx::supersampling;
	static_assert(M>0);
	static_assert(SS>0);

	using fp = typename Ctx::fp;
	using spectrum_t = spectrum<amplitude_t<fp>, 3>;

	const std::array<std::array<std::complex<fp>, 8>, Ctx::materials> materials;
    const bool render_diffractions{ true };
    const bool artistic{};
    const fp artistic_glints;
    const fp artistic_magnitude;

	// Random sequences used for MC sampling
	struct random_data_t {
		std::vector<bool> mean_phases;
		std::vector<bool> spatial_phases;
		std::vector<bool> temporal_phases;
	};
	random_data_t phases;
    static constexpr auto temporal_phases_count = 1024;

	// Wavenumber
	static auto wavenumber(fp lambda) noexcept {
		return 2*glm::pi<fp>() / (fp(1e-3) * lambda);
	}

	template <std::size_t S>
	static auto draw_surface(const fp size_um, const bool write_out = false) {
		const fp f_max = fp(S) / (2*size_um);

		std::vector<std::complex<double>> X;
		{
			std::vector<std::complex<double>> Y;
			Y.resize(sqr(S+1));

			// Draw symmetrically-conjugate phases
			for (std::size_t y=0; y<=S; ++y)
				for (std::size_t x=0; x<=y; ++x)
					Y[y*(S+1)+x] = std::polar(fp(1), glm::pi<fp>() * std::uniform_real_distribution<fp>{ -1, 1 }(random_generator{}()));
			for (std::size_t y=0; y<=S; ++y)
				for (std::size_t x=y+1; x<=S; ++x)
					Y[y*(S+1)+x] = std::conj(Y[x*(S+1)+y]);

			for (std::size_t y=0; y<=S; ++y)
				for (std::size_t x=0; x<=S; ++x) {
					const auto fx = f_max * fp(int(x)-S/2) / fp(S+1) * Sampler::fx_scale;
					const auto fy = f_max * fp(int(y)-S/2) / fp(S+1) * Sampler::fy_scale;
					Y[y*(S+1)+x] *= static_cast<double>(glm::sqrt(Sampler::PSD(fx, fy)));
				}

			// Inverse FFT
			X.resize(Y.size());
			using complex = typename fftw3cxx::fftw<double>::complex;
			std::complex<double> *in = reinterpret_cast<std::complex<double>*>(Y.data());
			std::complex<double> *out = reinterpret_cast<std::complex<double>*>(X.data());
			int Ss[] = { S+1,S+1 };

			static std::mutex m;
			m.lock();
			auto p = fftw3cxx::plan<double>::plan_dft(2, Ss, in, out, FFTW_BACKWARD, FFTW_ESTIMATE);
			m.unlock();
			p.execute();
		}

		std::vector<float> surface;
		surface.resize(X.size());
		for (std::size_t idx=0; idx<surface.size(); ++idx)
			surface[idx] = static_cast<float>(glm::abs(std::real(X[idx])));

		if (write_out) {
			// Write out surface (for reasons)
			::surface<float> s = { std::make_unique<float[]>(sqr(S+1)), S+1, S+1, 1 };
			std::copy(surface.begin(), surface.end(), s.data.get());
			write_2d_hdr("../debug_surface_out.hdr", s);
		}

		return surface;
	}

public:
	// Coherence size in um
	speckle_field_renderer(bool render_diffractions,
                           const size_t resolution,
						   std::array<std::array<std::complex<fp>, 8>, Ctx::materials> ior,
                           bool artistic, fp artistic_glints, fp artistic_magnitude,
						   const size_t lambda=0) noexcept
	: materials(ior), 
      render_diffractions(render_diffractions),
      artistic(artistic), artistic_glints(artistic_glints), artistic_magnitude(artistic_magnitude)
	{
		if (render_diffractions) {
			static constexpr auto tolerance = 0;//static_cast<int>(0.04f * glm::sqrt(float(M)));
			std::cout << "Drawing integration phases (RNG tolerance: " << std::to_string(tolerance) << "/" << std::to_string(M) << ")..." << std::endl;

            std::bernoulli_distribution dist(0.5);

            // Draw random integration phases
            phases.mean_phases.resize(M);
            phases.spatial_phases.resize(M);
            phases.temporal_phases.resize(sqr(temporal_phases_count)*M);

			while (true) {
				std::int64_t a{};
				for (std::size_t m=0; m<M; ++m)
					a += (phases.mean_phases[m] = dist(Rand{}())) ? 1 : -1;
				if (glm::abs(a)==0 || !Ctx::draw_phases_with_toletance) break;
			}

			print_progress(.0f);
            static constexpr auto idxs = create_array<std::size_t,sqr(temporal_phases_count)>();
            std::atomic<int> prog=0;
            std::for_each(std::execution::par_unseq, idxs.begin(), idxs.end(), [&](const auto r) {
				while (true) {
					std::int64_t a{};
					for (std::size_t m=0; m<M; ++m)
						a += (phases.temporal_phases[r*M+m] = dist(Rand{}())) ? 1 : -1;
					if (glm::abs(a)<=tolerance || !Ctx::draw_phases_with_toletance) break;
				}
				print_progress(static_cast<float>(prog+++1)/static_cast<float>(sqr(temporal_phases_count)));
            });
            print_progress(1.f);
            std::cout << std::endl << "Drawing phases done." << std::endl;
        }
	}

    const auto get_material(std::size_t mat) const noexcept { return materials[mat]; }

	auto operator()(std::size_t pixel_idx, int x, int y,
					const Ctx::vertex_t &wp, const Ctx::vertex_t &n, const Ctx::vertex_t &v, const Ctx::vertex_t &l, 
					const Ctx::vertex_t &dwpdx, const Ctx::vertex_t &dwpdy, 
					const Ctx::vertex_t &dndx, const Ctx::vertex_t &dndy, 
					const Ctx::vertex_t &wvp, const Ctx::vertex_t &L,
    				const illuminant_base* illuminant,
					const float size_x, float size_y,
					const float coherence_area, float coherence_length,
                    const bool area_light, const float light_sr, const std::array<glm::vec3, 4> &light_vs,
                    const float mat) const noexcept {
		const auto coherence_radius = glm::sqrt(coherence_area/glm::pi<float>());
        const auto minus_half_over_sqr_coherence_length = -fp(0.5)/sqr(coherence_length);

        const auto ss_size = glm::vec<2,fp>{ size_x/fp(SS), size_y/fp(SS) };
        const auto ss_area = glm::max(fp(1e-20),ss_size.x*ss_size.y);

        const auto coherence_norm = static_cast<fp>((1 - 2*normal_distribution::cdf(-ss_size.x*.5)) * (1 - 2*normal_distribution::cdf(-ss_size.y*.5)));
        const auto g = artistic ? artistic_magnitude : glm::sqrt(2*glm::min(fp(1),coherence_area/ss_area));

        // Sampling strategy, coherence area or surface autocovariance
        const auto Cs_const = glm::pow(glm::sqrt(glm::pi<fp>()), Sampler::Cf(mat)) / std::tgamma(Sampler::Cf(mat)/2);
        const auto Cs_gamma = fp(2)*glm::pi<fp>()/Sampler::Bf(mat);
        const bool importance_sample_coherence = false;//fp(1)/sqr(coherence_radius) < Cs_gamma;

        const auto &material = get_material(static_cast<std::size_t>(glm::round(mat)));
		
		glm::vec<3,fp> J{};
		glm::vec<3,fp> Mean{};

        // Supersample
        for (std::size_t ssx=0; ssx<SS; ++ssx)
        for (std::size_t ssy=0; ssy<SS; ++ssy) {
            const auto temporal_phase_idx = ((x*SS + ssx)%temporal_phases_count) + (((y*SS + ssy)%temporal_phases_count) * temporal_phases_count);
            const auto ssxf = (fp(ssx)+fp(.5))/fp(SS)-fp(.5), ssyf = (fp(ssy)+fp(.5))/fp(SS)-fp(.5);
            const auto wpss = wp + ssxf*dwpdx + ssyf*dwpdy;
            const auto nss  = glm::normalize(n + ssxf*dndx + ssyf*dndy);

            glm::vec<3,fp> tss,bss;
            orthogonal_frame(nss, tss, bss);
            const auto tbn = glm::mat<3,3,fp>(tss,bss,nss);

            // Calculate mean field
            const auto vv = glm::normalize(-(wvp-wpss) * tbn);
            const auto ll = glm::normalize(-(L-wpss) * tbn);
            const auto h  = glm::normalize(vv+ll);
            const auto desired_area_samples = glm::max(1ul, static_cast<std::size_t>(14.f*glm::mix(0.5f,1.f,h.z)*sqr(glm::max(1.f, (light_sr/glm::pi<fp>())))));
            const auto light_area_samples = area_light ? desired_area_samples : 1;
            glm::vec<3,fp> q{};
            for (std::size_t ax=0;ax<light_area_samples;++ax)
            for (std::size_t ay=0;ay<light_area_samples;++ay) {
                static constexpr auto LS = 12;
                const auto point_on_light_x0 = glm::mix(light_vs[0],light_vs[2],(ay+.5f)/float(light_area_samples));
                const auto point_on_light_x1 = glm::mix(light_vs[1],light_vs[3],(ay+.5f)/float(light_area_samples));
                const auto point_on_light = glm::mix(point_on_light_x0,point_on_light_x1,(ax+.5f)/float(light_area_samples));
                const auto lll = area_light ? glm::normalize(-(point_on_light-wpss) * tbn) : ll;
                for (auto s=0; s<LS; ++s) {
                    const auto lambda = illuminant->draw((s+fp(0.5))/fp(LS)) * fp(1e-3);
                    const auto xyz = XYZ_1931(lambda * fp(1e+3));

                    const auto eta_idx = glm::mix(fp(0), fp(material.size()-1), glm::clamp<fp>(0,1,fp(lambda-.420)/(.700-.420)));
                    const auto eta = (fp(1)-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))];
                    
                    q += xyz * Sampler::BRDF(lll,vv, Ctx::incident_eta,eta, lambda, mat) / fp(LS*sqr(light_area_samples));
                }
            }
            Mean += q / fp(sqr(SS));

            if (render_diffractions) {
                static_assert(M%8==0);

                // Clamp grazing angles
                const auto clamp = sqr(sqr(glm::clamp<fp>((glm::dot(vv+ll, glm::vec<3,fp>{0,0,0.5}))/0.4,0,1)));

                typename Ctx::vertex_t l1[8], l2[8], v1[8], v2[8], ki_1[8], ks_1[8], ki_2[8], ks_2[8];
                glm::vec<2,fp> x1[8], x2[8];
                std::complex<fp> eta[8];
                fp lambda[8], k[8];
                fp normalizer[8];

                // MC iterations
		        glm::vec<3,fp> I{};
                for (std::size_t m0=0; m0<M; m0+=8) {
                    // Draw things
                    for (auto m=0;m<8;++m) {
                        lambda[m] = illuminant->draw(std::uniform_real_distribution<fp>{ 0,1 }(Rand{}())) * fp(1e-3);

                        const auto x1frac = glm::vec<2,fp>{ std::uniform_real_distribution<fp>{ -.5,.5 }(Rand{}()), std::uniform_real_distribution<fp>{ -.5,.5 }(Rand{}()) };
                        glm::vec<2,fp> x2frac;
                        if (artistic) {
                            const auto draw_cr = glm::vec<2,fp>{ std::normal_distribution<fp>{ 0,10 }(Rand{}()), 
                                                                 std::normal_distribution<fp>{ 0,10 }(Rand{}()) };
                            x2frac = x1frac + draw_cr / ss_size;

                            const auto d = glm::length(draw_cr);
                            normalizer[m] = glm::exp(-fp(0.5)*sqr(d/glm::mix<fp>(1e-3, 10, glm::log(artistic_glints+1)/glm::log(2)))) * clamp;
                        }
                        else if (importance_sample_coherence) {
                            // Importance sample w.r.t. the coherence area
                            const auto draw_cr = glm::vec<2,fp>{ std::normal_distribution<fp>{ 0,coherence_radius }(Rand{}()), 
                                                                 std::normal_distribution<fp>{ 0,coherence_radius }(Rand{}()) };
                            x2frac = x1frac + draw_cr / ss_size;
                            normalizer[m] = clamp / coherence_norm;
                        }
                        else {
                            // Importance sample w.r.t. the surface autocovariance function
                            // Roughly approximate the Bessel K function distribution and correct at computation
                            const auto rand_dir = std::uniform_real_distribution<fp>{ 0, 2*glm::pi<fp>() }(Rand{}());
                            const auto drawn_r  = std::exponential_distribution<fp>{ Cs_gamma }(Rand{}());
                            const auto rand_r   = Cs_const * drawn_r;
                            const auto draw_p   = glm::vec<2,fp>{ glm::cos(rand_dir), glm::sin(rand_dir) } * rand_r;
                            x2frac = x1frac + draw_p / ss_size;

                            // Take coherence into account
                            normalizer[m] = (clamp * glm::exp(-sqr(rand_r/coherence_radius))) / (Cs_gamma * glm::exp(-Cs_gamma*drawn_r));
                        }

                        x1[m] = x1frac * ss_size;
                        x2[m] = x2frac * ss_size;

                        // Choose light source
                        auto light = L;
                        if (area_light) {
                            const auto ax = std::uniform_real_distribution<fp>{ 0,1 }(Rand{}());
                            const auto ay = std::uniform_real_distribution<fp>{ 0,1 }(Rand{}());
                            const auto point_on_light_x0 = glm::mix(light_vs[0],light_vs[2],ay);
                            const auto point_on_light_x1 = glm::mix(light_vs[1],light_vs[3],ay);
                            light = glm::mix(point_on_light_x0,point_on_light_x1,ax);
                        }

                        const auto wp1 = wpss + (x1frac.x*dwpdx + x1frac.y*dwpdy)/fp(SS);
                        const auto wp2 = wpss + (x2frac.x*dwpdx + x2frac.y*dwpdy)/fp(SS);
                        v1[m] = -(wvp-wp1) * tbn;
                        l1[m] = -(light-wp1) * tbn;
                        v2[m] = -(wvp-wp2) * tbn;
                        l2[m] = -(light-wp2) * tbn;

                        k[m] = 2*glm::pi<fp>() / lambda[m];
                        const auto eta_idx = glm::mix(fp(0), fp(material.size()-1), glm::clamp<fp>(0,1,(lambda[m]-fp(.420))/fp(.700-.420)));
                        eta[m] = (fp(1)-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))];
                    }

                    // Compute samples
                    fp Z[8];
                    static_assert(std::is_same_v<fp,float>);
                    if constexpr (std::is_same_v<fp,float>) {
                        // AVX2
                        v8sf l1norms, l2norms, v1norms, v2norms;
                        for (auto m=0;m<8;++m) {
                            l1norms[m] = glm::dot(l1[m],l1[m]);
                            l2norms[m] = glm::dot(l2[m],l2[m]);
                            v1norms[m] = glm::dot(v1[m],v1[m]);
                            v2norms[m] = glm::dot(v2[m],v2[m]);
                        }
                        // Path lengths and temporal coherence
                        auto path_length = _mm256_sub_ps(_mm256_add_ps(l1norms,v1norms), _mm256_add_ps(l2norms,v2norms));
                        path_length = _mm256_mul_ps(_mm256_mul_ps(path_length,path_length), _mm256_set1_ps(minus_half_over_sqr_coherence_length));
                        const auto cl = exp256_ps(path_length);

                        // Normalize
                        const auto one = _mm256_set1_ps(1.f);
                        l1norms = _mm256_div_ps(one, _mm256_sqrt_ps(l1norms));
                        l2norms = _mm256_div_ps(one, _mm256_sqrt_ps(l2norms));
                        v1norms = _mm256_div_ps(one, _mm256_sqrt_ps(v1norms));
                        v2norms = _mm256_div_ps(one, _mm256_sqrt_ps(v2norms));

                        for (auto m=0;m<8;++m) {
                            l1[m] *= l1norms[m];
                            l2[m] *= l2norms[m];
                            v1[m] *= v1norms[m];
                            v2[m] *= v2norms[m];

                            ki_1[m] = l1[m]*k[m]; ki_2[m] = l2[m]*k[m];
                            ks_1[m] = v1[m]*k[m]; ks_2[m] = v2[m]*k[m];

                            // Account for temporal coherence
                            normalizer[m] *= cl[m];
                        }

                        const auto Z8 = Sampler::Z_2_avx8(l1,v1,l2,v2, 
                                                          ki_1,ks_1,ki_2,ks_2, x1,x2, 
                                                          Ctx::incident_eta,eta,lambda, mat,
                                                          artistic);
                        for (auto m=0;m<8;++m) Z[m] = std::real(Z8[m]);
                    }
                    // else {
                    //     // Non-vectorized
                    //     for (auto m=0;m<8;++m) 
                    //         Z[m] = std::real(Sampler::Z_2(glm::normalize(l1[m]),glm::normalize(v1[m]),glm::normalize(l2[m]),glm::normalize(v2[m]), 
                    //                          ki_1[m],ks_1[m],ki_2[m],ks_2[m], x1[m],x2[m], 
                    //                          Ctx::incident_eta,eta[m],lambda[m],mat));
                    // }

                    // Aggregate
                    for (auto m=0;m<8;++m) {
                        const auto xyz = XYZ_1931(lambda[m] * fp(1e+3));
                        const auto X = xyz * (normalizer[m] * Z[m]);
                        const auto Q = (phases.mean_phases[m0+m]?fp(-1):fp(1)) * q;

                        if constexpr (render_spatial)
                            I += (phases.spatial_phases[m0+m]?fp(-1):fp(1)) * X;
                        I += g * (phases.temporal_phases[temporal_phase_idx*M+m0+m]?fp(-1):fp(1)) * (X + Q);
                    }
                }

                // Fuctuating intensity
			    J += I / (glm::sqrt(fp(M)) * fp(sqr(SS)));
            }
		}
		
        // Write out colour values of the intensities
		const auto rgb_fluc = illuminant->XYZtoRGB(glm::vec3(J));
		const auto rgb_mean = illuminant->XYZtoRGB(glm::vec3(Mean));
		return std::make_pair(rgb_mean, rgb_fluc);
	}
};
