
// Copyright 2019 Shlomi Steinberg
// 

#pragma once

#include <complex>
#include <random>
#include <cmath>
#include "glm.hpp"
#include "fresnel.hpp"
#include "normal_distribution.hpp"
#include "common.hpp"
#include "surface.hpp"
#include "load2d.hpp"
#include "write2d.hpp"
#include "avx.hpp"

template <typename ctx>
class mc_k_correlation_sampler {
public:
	using vertex_t = typename ctx::vertex_t;
	using fp = typename ctx::fp;

private:
	// Config
	static constexpr fp fx_scale = 1.;
	static constexpr fp fy_scale = 1.;

	static constexpr fp sigma = 0.5;	// in um

    // Buddha
	// static constexpr fp C0 = 1.55;
	// static constexpr fp B0 = 7;
    // static constexpr fp C1 = 2.6;
	// static constexpr fp B1 = 18;

	static constexpr fp C0 = 2.2;
	static constexpr fp B0 = 10;
    static constexpr fp C1 = C0;
	static constexpr fp B1 = B0;
	// /Config

	static constexpr auto n = vertex_t{ 0,0,1 };

public:
    static auto Bf(float m) { return glm::mix(B0,B1,sqr(m)); }
    static auto Cf(float m) { return glm::mix(C0,C1,sqr(m)); }

	// static constexpr auto m2() noexcept {
	// 	if constexpr (C==3) {
	// 		return sqr(sigma)*(sqr(glm::pi<fp>())*(glm::log(sqr(2*B) + 1) + sqr(2*B)*glm::log(sqr(2*B) + 1) - sqr(2*B)))/(2*B*B*B*B);
	// 	}
	// 	else {
	// 		constexpr auto t = B*B*B*B*(sqr(C) - 4*C + 3);
	// 		return -sqr(sigma)*(2*sqr(B*glm::pi<fp>())*(2/t - (glm::pow(sqr(2*B) + 1, (1-C)/2)*(sqr(2*B)*C - sqr(2*B) + 2))/t)*(C - 1))/(glm::pow(sqr(2*B) + 1, (1-C)/2) - 1);
	// 	}
	// }

    static constexpr auto l_cor(float m) noexcept {
        const auto B = Bf(m), C = Cf(m);
        return (B*(C - 1)*sqr(sigma)) / (2*glm::sqrt(C*glm::pi<fp>())*sigma*glm::abs(glm::pow(100*sqr(B) + 1, fp(.5) - C/2) - 1));
    }

	static constexpr auto PSD(const fp f2, float m) noexcept {
        const auto B = Bf(m), C = Cf(m);
		return -fp(C-1) * sqr(B * sigma) / ((glm::pow<fp>(sqr(2*B)+1,(1-C)/2)-1) * glm::pow<fp>(1 + sqr(B)*f2, (C+1)/2));
	}
	static constexpr auto PSD(const fp fx, const fp fy, float m) noexcept { return PSD(fx*fx+fy*fy,m); }

    static auto PSD_avx8(const v8sf f2, float m) noexcept {
        const auto B = Bf(m), C = Cf(m);
        const auto Cs = _mm256_set1_ps(C);
        const auto Bs = _mm256_set1_ps(B);
        const auto B2s = _mm256_mul_ps(Bs, Bs);
        const auto sigmas = _mm256_set1_ps(sigma);
        const auto one_minus_C = _mm256_sub_ps(one8, Cs);
        const auto one_plus_C = _mm256_add_ps(one8, Cs);

        const auto b1 = glm::pow<fp>(sqr(2*B)+1,(1-C)/2)-1;
        const auto b = _mm256_set1_ps(b1);

        const auto c1 = _mm256_add_ps(one8, _mm256_mul_ps(f2,B2s));
        const auto c1pow = _mm256_mul_ps(one_plus_C, half8);
        const auto c = exp256_ps(_mm256_mul_ps(c1pow, log256_ps(c1)));
        const auto a = _mm256_mul_ps(one_minus_C, _mm256_mul_ps(B2s, _mm256_mul_ps(sigmas, sigmas)));

        return _mm256_div_ps(a, _mm256_mul_ps(b,c));
    }
    static auto PSD_avx8(const v8sf fx, const v8sf fy, float m) noexcept { return PSD_avx8(_mm256_add_ps(_mm256_mul_ps(fx,fx), _mm256_mul_ps(fy,fy)),m); }

	static constexpr auto C_S(const fp x, const fp y, float m) noexcept {
        const auto B = Bf(m), C = Cf(m);
		const auto l = glm::max(fp(1e-10), glm::sqrt(sqr(x*fx_scale) + sqr(y*fy_scale)));
		const auto c = 2 / std::tgamma(C/2) * 
                       glm::pow<fp>(glm::pi<fp>()*l/B,C/2) * 
                       std::cyl_bessel_k(C/2,glm::min(75.f, 2*glm::pi<fp>()*l/B));
		return glm::clamp<fp>(c, 0, 1);
	}
	static auto C_S_avx8(const v8sf x, const v8sf y, float m) noexcept {
        const auto B = Bf(m), C = Cf(m);
        const auto fx2s8 = _mm256_set1_ps(fx_scale*fx_scale);
        const auto fy2s8 = _mm256_set1_ps(fy_scale*fy_scale);

        const auto fs = _mm256_add_ps(_mm256_mul_ps(fx2s8,_mm256_mul_ps(x,x)), _mm256_mul_ps(fy2s8,_mm256_mul_ps(y,y)));
        const auto ls = _mm256_max_ps(_mm256_sqrt_ps(fs), _mm256_set1_ps(1e-10f));
        auto bessels = _mm256_mul_ps(_mm256_set1_ps(2*glm::pi<float>()/B),ls);
		bessels = _mm256_min_ps(_mm256_set1_ps(75.f),bessels);		// Clamp
        v8sf c;
        for (auto i=0;i<8;++i)
            c[i] = std::cyl_bessel_k(C/2,bessels[i]);

        const auto b1 = _mm256_mul_ps(_mm256_set1_ps(glm::pi<float>()/B),ls);
        const auto b = exp256_ps(_mm256_mul_ps(_mm256_set1_ps(C/2), log256_ps(b1)));
        const auto a = _mm256_set1_ps(2.f / std::tgamma(C/2.f));

        const auto ret = _mm256_mul_ps(a,_mm256_mul_ps(b,c));
        return _mm256_min_ps(one8,_mm256_max_ps(ret,_mm256_setzero_ps()));
	}

	static constexpr auto F(const vertex_t &u, const vertex_t &v, const std::complex<fp> &incident_eta, const std::complex<fp> &eta) noexcept {
		const auto e = eta/incident_eta;
		const auto proj_v = glm::vec<2,fp>{ v.x, v.y };
		const auto proj_u = glm::vec<2,fp>{ u.x, u.y };

        const auto l = glm::sqrt(glm::dot(proj_v,proj_v)*glm::dot(proj_u,proj_u));
		const auto cosPhis = l>1e-10f ? glm::clamp<fp>(glm::dot(-proj_v,proj_u)/l,-1,1) : fp(1);
		const auto cosThetai = glm::max(fp(0), glm::dot(n,u));
		const auto cosThetas = glm::max(fp(0), glm::dot(n,v));
		return Q_s(e, cosPhis, cosThetai, cosThetas) + Q_p(e, cosPhis, cosThetai, cosThetas);
	}
	static auto F_avx8(const vertex_t *u, const vertex_t *v, const std::complex<fp> &incident_eta, const std::complex<fp> *eta) noexcept {
        v8sf l2s, uv_dots;
        for (auto i=0;i<8;++i) {
            const auto proj_v = glm::vec<2,fp>{ v[i].x, v[i].y };
            const auto proj_u = glm::vec<2,fp>{ u[i].x, u[i].y };

            l2s[i] = glm::dot(proj_v,proj_v)*glm::dot(proj_u,proj_u);
            uv_dots[i] = glm::dot(-proj_v,proj_u);
        }
        const auto ls = _mm256_sqrt_ps(l2s);
        const auto uv_dots_norm = _mm256_div_ps(uv_dots,ls);
        v8sf qss, qps;
        for (auto i=0;i<8;++i) {
            const auto e = eta[i]/incident_eta;
            const auto cosPhis = ls[i]>1e-10f ? glm::clamp<fp>(uv_dots_norm[i],-1,1) : fp(1);
            const auto cosThetai = glm::max(fp(0), glm::dot(n,u[i]));
            const auto cosThetas = glm::max(fp(0), glm::dot(n,v[i]));
            qss[i] = Q_s(e, cosPhis, cosThetai, cosThetas);
            qps[i] = Q_p(e, cosPhis, cosThetai, cosThetas);
        }

        return _mm256_add_ps(qss,qps);
	}

public:
	// Lambda in um
	static constexpr auto BRDF(const vertex_t &u, const vertex_t &v, 
							   const std::complex<fp> &incident_eta, const std::complex<fp> &eta, 
							   const fp lambda, float m) noexcept {
		const auto proj_v = glm::vec<2,fp>{ v.x, v.y };
		const auto proj_u = glm::vec<2,fp>{ u.x, u.y };
        const auto fv = -(proj_u + proj_v) * glm::vec<2,fp>{ fx_scale,fy_scale } / fp(0.5);
		const auto f = glm::dot(fv,fv);
		return PSD(f,m) * F(u,v,incident_eta,eta);
	}
	// Lambda in um
	static auto BRDF_avx8(const vertex_t *u, const vertex_t *v, 
                          const std::complex<fp> &incident_eta, const std::complex<fp> *eta, 
                          const fp *lambda, float m) noexcept {
        v8sf f2s;
        for (auto i=0;i<8;++i) {
            const auto proj_v = glm::vec<2,fp>{ v[i].x, v[i].y };
            const auto proj_u = glm::vec<2,fp>{ u[i].x, u[i].y };
            const auto fv = -(proj_u + proj_v) * glm::vec<2,fp>{ fx_scale,fy_scale } / fp(0.5);
            f2s[i] = glm::dot(fv,fv);
        }
		return _mm256_mul_ps(PSD_avx8(f2s,m), F_avx8(u,v,incident_eta,eta));
	}

	// x in um, k in um^(-1)
	static auto Z_2(const vertex_t &u1, const vertex_t &v1, const vertex_t &u2, const vertex_t &v2, 
					const vertex_t &ki_1, const vertex_t &ks_1, const vertex_t &ki_2, const vertex_t &ks_2, 
					const glm::vec<2,fp> &x1, const glm::vec<2,fp> &x2,
					const std::complex<fp> &incident_eta, const std::complex<fp> &eta, const fp lambda, float m) noexcept {
		// A fast approximation
		const auto f1 = BRDF(u1,v1,incident_eta,eta,lambda,m);
		const auto f2 = BRDF(u2,v2,incident_eta,eta,lambda,m);

		const auto b = -glm::dot(ki_1+ks_1, vertex_t{ x1,0 }) + glm::dot(ki_2+ks_2,vertex_t{ x2,0 });
		const auto c = C_S(glm::abs(x2.x-x1.x),glm::abs(x2.y-x1.y),m);
		const auto a1 = sqr(sigma)/2 * (sqr(glm::dot(ki_1+ks_1,n)) + sqr(glm::dot(ki_2+ks_2,n)));
		const auto a2 = c * sqr(sigma) * glm::dot(ki_1+ks_1,n) * glm::dot(ki_2+ks_2,n);

		return std::polar(glm::sqrt(f1 * f2) * glm::exp(a2-a1), b);
	}
	// x in um, k in um^(-1)
	static auto Z_2_avx8(const vertex_t *u1, const vertex_t *v1, const vertex_t *u2, const vertex_t *v2, 
                         const vertex_t *ki_1, const vertex_t *ks_1, const vertex_t *ki_2, const vertex_t *ks_2, 
                         const glm::vec<2,fp> *x1, const glm::vec<2,fp> *x2,
                         const std::complex<fp> &incident_eta, const std::complex<fp> *eta, const fp *lambda, float m,
                         bool artistic) noexcept {
		// A fast approximation
		const auto f1 = BRDF_avx8(u1,v1,incident_eta,eta,lambda,m);
		const auto f2 = BRDF_avx8(u2,v2,incident_eta,eta,lambda,m);

        v8sf fxs,fys,bs,a1s,a2s;
        for (auto i=0;i<8;++i) {
		    bs[i] = -glm::dot(ki_1[i]+ks_1[i], vertex_t{ x1[i],0 }) + glm::dot(ki_2[i]+ks_2[i],vertex_t{ x2[i],0 });
            fxs[i] = glm::abs(x2[i].x-x1[i].x);
            fys[i] = glm::abs(x2[i].y-x1[i].y);
            a1s[i] = sqr(sigma)/2 * (sqr(glm::dot(ki_1[i]+ks_1[i],n)) + sqr(glm::dot(ki_2[i]+ks_2[i],n)));
            a2s[i] = sqr(sigma) * glm::dot(ki_1[i]+ks_1[i],n) * glm::dot(ki_2[i]+ks_2[i],n);
        }
		const auto cs = artistic ? one8 : C_S_avx8(fxs,fys,m);
        a2s = _mm256_mul_ps(cs, a2s);
        const auto f = _mm256_sqrt_ps(_mm256_mul_ps(f1,f2));
        const auto expa = exp256_ps(_mm256_sub_ps(a2s,a1s));
        const auto mag = _mm256_mul_ps(f,expa);

        v8sf sins, coss;
        sincos256_ps(bs, &sins, &coss);
        sins = _mm256_mul_ps(mag,sins);
        coss = _mm256_mul_ps(mag,coss);

        std::array<std::complex<float>,8> rets;
        for (auto i=0;i<8;++i)
            rets[i] = { coss[i], sins[i] };
        return rets;
	}
};
