
// Copyright 2019 Shlomi Steinberg
// 

#pragma once

#include <complex>
#include "glm.hpp"
#include "common.hpp"

template <typename T, typename RefractiveIndex>
auto fresnel_p(T cos_theta_incident, RefractiveIndex eta1, RefractiveIndex eta2) {
	const auto eta = eta1/eta2;
	const auto c = cos_theta_incident;
	const auto sin2 = T(1) - sqr(c);

	const auto t = std::sqrt(RefractiveIndex{ 1 } - sqr(eta)*sin2);
	return (t*eta - c) / (t*eta + c);
}
template <typename T, typename RefractiveIndex>
auto fresnel_s(T cos_theta_incident, RefractiveIndex eta1, RefractiveIndex eta2) {
	const auto eta = eta1/eta2;
	const auto c = cos_theta_incident;
	const auto sin2 = T(1) - sqr(c);

	const auto t = std::sqrt(RefractiveIndex{ 1 } - sqr(eta)*sin2);
	return (c*eta - t) / (c*eta + t);
}

template <typename RefractiveIndex, typename T>
auto q_ss(RefractiveIndex eta, T cosPhis, T cosThetai, T cosThetas) {
	eta = sqr(eta);
	const auto sinThetai2 = glm::min(T(1),T(1) - sqr(cosThetai)), sinThetas2 = glm::min(T(1),T(1) - sqr(cosThetas));
	return 	(eta - T(1)) * cosPhis /
				(
					(cosThetai + std::sqrt(eta - sinThetai2)) *
					(cosThetas + std::sqrt(eta - sinThetas2))
				);
}
template <typename RefractiveIndex, typename T>
auto q_sp(RefractiveIndex eta, T cosPhis, T cosThetai, T cosThetas) {
	eta = sqr(eta);
	const auto sinThetai2 = glm::min(T(1),T(1) - sqr(cosThetai)), sinThetas2 = glm::min(T(1),T(1) - sqr(cosThetas));
	const auto sinPhis    = std::sqrt(1 - glm::min(T(1),sqr(cosPhis)));
	return	(eta - T(1)) * std::sqrt(eta - sinThetas2) * sinPhis /
				(
					(cosThetai + std::sqrt(eta - sinThetai2)) *
					(eta * cosThetas + std::sqrt(eta - sinThetas2))
				);
}
template <typename RefractiveIndex, typename T>
auto q_ps(RefractiveIndex eta, T cosPhis, T cosThetai, T cosThetas) {
	eta = sqr(eta);
	const auto sinThetai2 = glm::min(T(1),T(1) - sqr(cosThetai)), sinThetas2 = glm::min(T(1),T(1) - sqr(cosThetas));
	const auto sinPhis    = std::sqrt(1 - glm::min(T(1),sqr(cosPhis)));
	return (eta - T(1)) * std::sqrt(eta - sinThetai2) * sinPhis /
				(
					(eta * cosThetai + std::sqrt(eta - sinThetai2)) * 
					(cosThetas + std::sqrt(eta - sinThetas2))
				);
}
template <typename RefractiveIndex, typename T>
auto q_pp(RefractiveIndex eta, T cosPhis, T cosThetai, T cosThetas) {
	eta = sqr(eta);
	const auto sinThetai2 = glm::min(T(1),T(1) - sqr(cosThetai)), sinThetas2 = glm::min(T(1),T(1) - sqr(cosThetas));
	const auto sinThetai = std::sqrt(1 - glm::min(T(1),sqr(cosThetai))), sinThetas = std::sqrt(1 - glm::min(T(1),sqr(cosThetas)));
	return ((eta - T(1)) * (std::sqrt(eta - sinThetai2) * std::sqrt(eta - sinThetas2)) * cosPhis - sinThetai*sinThetas * eta) /
				(
					(eta * cosThetai + std::sqrt(eta - sinThetai2)) * 
					(eta * cosThetas + std::sqrt(eta - sinThetas2))
				);
}
template <typename RefractiveIndex, typename T>
auto q_s(RefractiveIndex eta, T cosPhis, T cosThetai, T cosThetas) {
	return 	q_ss(eta, cosPhis, cosThetai, cosThetas) +
			q_ps(eta, cosPhis, cosThetai, cosThetas);
}
template <typename RefractiveIndex, typename T>
auto q_p(RefractiveIndex eta, T cosPhis, T cosThetai, T cosThetas) {
	return 	q_sp(eta, cosPhis, cosThetai, cosThetas) +
			q_pp(eta, cosPhis, cosThetai, cosThetas);
}
template <typename RefractiveIndex, typename T>
auto Q_s(RefractiveIndex eta, T cosPhis, T cosThetai, T cosThetas) {
	return 	std::norm(q_ss(eta, cosPhis, cosThetai, cosThetas)) +
			std::norm(q_ps(eta, cosPhis, cosThetai, cosThetas));
}
template <typename RefractiveIndex, typename T>
auto Q_p(RefractiveIndex eta, T cosPhis, T cosThetai, T cosThetas) {
	return 	std::norm(q_sp(eta, cosPhis, cosThetai, cosThetas)) +
			std::norm(q_pp(eta, cosPhis, cosThetai, cosThetas));
}
