/*
    This file is part of Mitsuba, a physically based rendering system.

    Copyright (c) 2007-2011 by Wenzel Jakob and others.

    Mitsuba is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License Version 3
    as published by the Free Software Foundation.

    Mitsuba is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#include <mitsuba/render/bsdf.h>
#include <mitsuba/hw/basicshader.h>
#include <boost/algorithm/string.hpp>

MTS_NAMESPACE_BEGIN

/*!\plugin{ward}{Anisotropic Ward BRDF}
 * \order{11}
 * \parameters{
 *     \parameter{variant}{\String}{
 *         Determines the variant of the Ward model to use:
 *         \begin{enumerate}[(i)]
 *             \item \code{ward}: The original model by Ward \cite{Ward1992Measuring}
 *             --- suffers from energy loss at grazing angles.
 *             \item \code{ward-duer}: Corrected Ward model with lower energy loss 
 *             at grazing angles \cite{Dur2006Improved}.
 *             Does not always conserve energy. 
 *             \item \code{balanced}: Improved version of the \code{ward-duer}
 *             model with energy balance at all angles \cite{Geisler2010New}.
 *             \vspace{-4mm}
 *         \end{enumerate}
 *     }
 *     \parameter{alphaU, alphaV}{\Float\Or\Texture}{
 *         Specifies the anisotropic roughness values along the tangent and 
 *         bitangent directions. 
 *         \default{0.1}. 
 *     }
 *     \parameter{specular\showbreak Reflectance}{\Spectrum\Or\Texture}{
 *         Specifies the weight of the specular reflectance component.\default{0.2}}
 *     \parameter{diffuse\showbreak Reflectance}{\Spectrum\Or\Texture}{
 *         Specifies the weight of the diffuse reflectance component\default{0.5}}
 * }
 * \renderings{
 *     \rendering{$\alpha_u=0.1,\ \alpha_v=0.3$}{bsdf_ward_01_03}
 *     \rendering{$\alpha_u=0.3,\ \alpha_v=0.1$}{bsdf_ward_03_01}
 * }

 * This plugin implements the anisotropic Ward reflectance model and
 * several extensions. They are described in the papers
 * \begin{enumerate}[(i)]
 *    \item ``Measuring and Modeling Anisotropic Reflection'' 
 *      by Greg Ward \cite{Ward1992Measuring}
 *    \item ``Notes on the Ward BRDF'' by Bruce Walter \cite{Walter2005Notes}
 *    \item ``An Improved Normalization for the Ward Reflectance Model''
 *      by Arne D\"ur \cite{Dur2006Improved}
 *    \item ``A New Ward BRDF Model with Bounded Albedo'' by
 *      Geisler-Moroder et al. \cite{Geisler2010New}
 * \end{enumerate}
 *
 * Like the Phong BRDF, the Ward model does not take the Fresnel reflectance
 * of the material into account. In an experimental study by Ngan et al. 
 * \cite{Ngan2005Experimental}, the Ward model performed noticeably worse than
 * models based on microfacets.
 *
 * For this reason, it is usually preferable to switch to a microfacet model 
 * that incorporates knowledge about the material's index of refraction. In Mitsuba, 
 * two such alternatives to \pluginref{ward} are given by the plugins 
 * \pluginref{roughconductor} and \pluginref{roughplastic} (depending on the
 * material type).
 *
 * When using this plugin, note that the diffuse and specular reflectance
 * components should add up to a value less than or equal to one (for each 
 * color channel). Otherwise, they will automatically be scaled appropriately 
 * to ensure energy conservation.
 */
class Ward : public BSDF {
public:
	/// Supported model types
	enum EModelVariant {
		/// The original Ward model
		EWard = 0,
		/// Ward model with correction by Arne Duer
		EWardDuer = 1,
		/// Energy-balanced Ward model
		EBalanced = 2
	};

	Ward(const Properties &props) 
		: BSDF(props) {
		m_diffuseReflectance = new ConstantSpectrumTexture(
			props.getSpectrum("diffuseReflectance", Spectrum(0.5f)));
		m_specularReflectance = new ConstantSpectrumTexture(
			props.getSpectrum("specularReflectance", Spectrum(0.2f)));
		
		std::string type = 
			boost::to_lower_copy(props.getString("variant", "balanced"));
		if (type == "ward")
			m_modelVariant = EWard;
		else if (type == "ward-duer")
			m_modelVariant = EWardDuer;
		else if (type == "balanced")
			m_modelVariant = EBalanced;
		else
			Log(EError, "Specified an invalid model type \"%s\", must be "
				"\"ward\", \"ward-duer\", or \"balanced\"!", type.c_str());

		Float alpha = props.getFloat("alpha", 0.1f),
			  alphaU = props.getFloat("alphaU", alpha),
			  alphaV = props.getFloat("alphaV", alpha);

		m_alphaU = new ConstantFloatTexture(alphaU);
		if (alphaU == alphaV)
			m_alphaV = m_alphaU;
		else
			m_alphaV = new ConstantFloatTexture(alphaV);
		m_specularSamplingWeight = 0.0f;
	}

	Ward(Stream *stream, InstanceManager *manager) 
	 : BSDF(stream, manager) {
		m_modelVariant = (EModelVariant) stream->readUInt();
		m_diffuseReflectance = static_cast<Texture *>(manager->getInstance(stream));
		m_specularReflectance = static_cast<Texture *>(manager->getInstance(stream));
		m_alphaU = static_cast<Texture *>(manager->getInstance(stream));
		m_alphaV = static_cast<Texture *>(manager->getInstance(stream));

		configure();
	}

	void configure() {
		unsigned int extraFlags = 0;
		if (m_alphaU != m_alphaV)
			extraFlags |= EAnisotropic;

		m_components.clear();
		m_components.push_back(EGlossyReflection | EFrontSide | extraFlags
			| ((!m_specularReflectance->isConstant() || !m_alphaU->isConstant()
			  || !m_alphaV->isConstant()) ? ESpatiallyVarying : 0));
		m_components.push_back(EDiffuseReflection | EFrontSide | extraFlags
			| (m_diffuseReflectance->isConstant() ? 0 : ESpatiallyVarying));

		/* Verify the input parameters and fix them if necessary */
		std::pair<Texture *, Texture *> result = ensureEnergyConservation(
			m_specularReflectance, m_diffuseReflectance,
			"specularReflectance", "diffuseReflectance", 1.0f);
		m_specularReflectance = result.first;
		m_diffuseReflectance = result.second;

		/* Compute weights that steer samples towards
		   the specular or diffuse components */
		Float dAvg = m_diffuseReflectance->getAverage().getLuminance(),
			  sAvg = m_specularReflectance->getAverage().getLuminance();
		m_specularSamplingWeight = sAvg / (dAvg + sAvg);

		m_usesRayDifferentials = 
			m_diffuseReflectance->usesRayDifferentials() ||
			m_specularReflectance->usesRayDifferentials() ||
			m_alphaU->usesRayDifferentials() ||
			m_alphaV->usesRayDifferentials();

		BSDF::configure();
	}

	Spectrum getDiffuseReflectance(const Intersection &its) const {
		return m_diffuseReflectance->getValue(its);
	}


	Spectrum eval(const BSDFQueryRecord &bRec, EMeasure measure) const {
		if (Frame::cosTheta(bRec.wi) <= 0 ||
			Frame::cosTheta(bRec.wo) <= 0 || measure != ESolidAngle)
			return Spectrum(0.0f);

		bool hasSpecular = (bRec.typeMask & EGlossyReflection)
				&& (bRec.component == -1 || bRec.component == 0);
		bool hasDiffuse  = (bRec.typeMask & EDiffuseReflection)
				&& (bRec.component == -1 || bRec.component == 1);

		Spectrum result(0.0f);
		if (hasSpecular) {
			Vector H = bRec.wi+bRec.wo;
			Float alphaU = m_alphaU->getValue(bRec.its).average();
			Float alphaV = m_alphaV->getValue(bRec.its).average();
			
			Float factor1 = 0.0f;
			switch (m_modelVariant) {
				case EWard:
					factor1 = 1.0f / (4.0f * M_PI * alphaU * alphaV * 
						std::sqrt(Frame::cosTheta(bRec.wi)*Frame::cosTheta(bRec.wo)));
					break;
				case EWardDuer:
					factor1 = 1.0f / (4.0f * M_PI * alphaU * alphaV * 
						Frame::cosTheta(bRec.wi)*Frame::cosTheta(bRec.wo));
					break;
				case EBalanced:
					factor1 = dot(H,H) / (M_PI * alphaU * alphaV 
						* std::pow(Frame::cosTheta(H),4));
					break;
				default:
					Log(EError, "Unknown model type!");
			}

			Float factor2 = H.x / alphaU, factor3 = H.y / alphaV;
			Float exponent = -(factor2*factor2+factor3*factor3)/(H.z*H.z);
			Float specRef = factor1 * std::exp(exponent);
			/* Important to prevent numeric issues when evaluating the
			   sampling density of the Ward model in places where it takes
			   on miniscule values (Veach-MLT does this for instance) */
			if (specRef > 1e-10f)
				result += m_specularReflectance->getValue(bRec.its) * specRef;
		}

		if (hasDiffuse) 
			result += m_diffuseReflectance->getValue(bRec.its) * INV_PI;

		return result * Frame::cosTheta(bRec.wo);
	}

	Float pdf(const BSDFQueryRecord &bRec, EMeasure measure) const {
		if (Frame::cosTheta(bRec.wi) <= 0 ||
			Frame::cosTheta(bRec.wo) <= 0 || measure != ESolidAngle)
			return 0.0f;

		bool hasSpecular = (bRec.typeMask & EGlossyReflection)
				&& (bRec.component == -1 || bRec.component == 0);
		bool hasDiffuse  = (bRec.typeMask & EDiffuseReflection)
				&& (bRec.component == -1 || bRec.component == 1);

		Float diffuseProb = 0.0f, specProb = 0.0f;

		if (hasSpecular) {
			Float alphaU = m_alphaU->getValue(bRec.its).average();
			Float alphaV = m_alphaV->getValue(bRec.its).average();
			Vector H = normalize(bRec.wi+bRec.wo);
			Float factor1 = 1.0f / (4.0f * M_PI * alphaU * alphaV * 
				dot(H, bRec.wi) * std::pow(Frame::cosTheta(H), 3));
			Float factor2 = H.x / alphaU, factor3 = H.y / alphaV;

			Float exponent = -(factor2*factor2+factor3*factor3)/(H.z*H.z);
			specProb = factor1 * std::exp(exponent);
		}

		if (hasDiffuse) 
			diffuseProb = Frame::cosTheta(bRec.wo) * INV_PI;

		if (hasDiffuse && hasSpecular)
			return m_specularSamplingWeight * specProb + 
				   (1-m_specularSamplingWeight) * diffuseProb;
		else if (hasDiffuse)
			return diffuseProb;
		else if (hasSpecular)
			return specProb;
		else
			return 0.0f;
	}

	inline Spectrum sample(BSDFQueryRecord &bRec, Float &_pdf, const Point2 &_sample) const {
		Point2 sample(_sample);

		bool hasSpecular = (bRec.typeMask & EGlossyReflection)
				&& (bRec.component == -1 || bRec.component == 0);
		bool hasDiffuse  = (bRec.typeMask & EDiffuseReflection)
				&& (bRec.component == -1 || bRec.component == 1);

		if (!hasSpecular && !hasDiffuse)
			return Spectrum(0.0f);

		bool choseSpecular = hasSpecular;

		if (hasDiffuse && hasSpecular) {
			if (sample.x <= m_specularSamplingWeight) {
				sample.x /= m_specularSamplingWeight;
			} else {
				sample.x = (sample.x - m_specularSamplingWeight)
					/ (1-m_specularSamplingWeight);
				choseSpecular = false;
			}
		}

		if (choseSpecular) {
			Float alphaU = m_alphaU->getValue(bRec.its).average();
			Float alphaV = m_alphaV->getValue(bRec.its).average();

			Float phiH = std::atan(alphaV/alphaU 
				* std::tan(2.0f * M_PI * sample.y));
			if (sample.y > 0.5f)
				phiH += M_PI;
			Float cosPhiH = std::cos(phiH);
			Float sinPhiH = std::sqrt(std::max((Float) 0.0f, 
				1.0f-cosPhiH*cosPhiH));

			Float thetaH = std::atan(std::sqrt(std::max((Float) 0.0f, 
				-std::log(sample.x) / (
					(cosPhiH*cosPhiH) / (alphaU*alphaU) +
					(sinPhiH*sinPhiH) / (alphaV*alphaV)
			))));
			Vector H = sphericalDirection(thetaH, phiH);
			bRec.wo = H * (2.0f * dot(bRec.wi, H)) - bRec.wi;

			bRec.sampledComponent = 1;
			bRec.sampledType = EGlossyReflection;

			if (Frame::cosTheta(bRec.wo) <= 0.0f)
				return Spectrum(0.0f);
		} else {
			bRec.wo = squareToHemispherePSA(sample);
			bRec.sampledComponent = 0;
			bRec.sampledType = EDiffuseReflection;
		}

		_pdf = pdf(bRec, ESolidAngle);

		if (_pdf == 0) 
			return Spectrum(0.0f);
		else
			return eval(bRec, ESolidAngle);
	}

	Spectrum sample(BSDFQueryRecord &bRec, const Point2 &sample) const {
		Float pdf=0;
		Spectrum result = Ward::sample(bRec, pdf, sample);

		if (result.isZero())
			return Spectrum(0.0f);
		else
			return result / pdf;
	}

	void addChild(const std::string &name, ConfigurableObject *child) {
		if (child->getClass()->derivesFrom(MTS_CLASS(Texture))) {
			if (name == "alphaU")
				m_alphaU = static_cast<Texture *>(child);
			else if (name == "alphaV")
				m_alphaV = static_cast<Texture *>(child);
			else if (name == "diffuseReflectance")
				m_diffuseReflectance = static_cast<Texture *>(child);
			else if (name == "specularReflectance")
				m_specularReflectance = static_cast<Texture *>(child);
			else
				BSDF::addChild(name, child);
		} else {
			BSDF::addChild(name, child);
		}
	}

	void serialize(Stream *stream, InstanceManager *manager) const {
		BSDF::serialize(stream, manager);

		stream->writeUInt(m_modelVariant);
		manager->serialize(stream, m_diffuseReflectance.get());
		manager->serialize(stream, m_specularReflectance.get());
		manager->serialize(stream, m_alphaU.get());
		manager->serialize(stream, m_alphaV.get());
	}

	Shader *createShader(Renderer *renderer) const; 

	std::string toString() const {
		std::ostringstream oss;
		oss << "Ward[" << endl
			<< "  name = \"" << getName() << "\"," << endl
			<< "  variant = ";

		switch (m_modelVariant) {
			case EWard: oss << "ward," << endl; break;
			case EWardDuer: oss << "wardDuer," << endl; break;
			case EBalanced: oss << "balanced," << endl; break;
			default: Log(EError, "Unknown model type!");
		}

		oss << "  diffuseReflectance = " << indent(m_diffuseReflectance->toString()) << "," << endl
			<< "  specularReflectance = " << indent(m_specularReflectance->toString()) << "," << endl
			<< "  specularSamplingWeight = " << m_specularSamplingWeight << "," << endl
			<< "  alphaU = " << indent(m_alphaU->toString()) << "," << endl
			<< "  alphaV = " << indent(m_alphaV->toString()) << endl
			<< "]";
		return oss.str();
	}

	MTS_DECLARE_CLASS()
private:
	EModelVariant m_modelVariant;
	ref<Texture> m_diffuseReflectance;
	ref<Texture> m_specularReflectance;
	ref<Texture> m_alphaU;
	ref<Texture> m_alphaV;
	Float m_specularSamplingWeight;
};

// ================ Hardware shader implementation ================ 

/**
 * GLSL port of the Ward shader. This version only implements the variant
 * with energy balance. When the roughness is lower than 
 * \alpha < 0.2, the shader clamps it to 0.2 so that it will still perform
 * reasonably well in a VPL-based preview.
 */
class WardShader : public Shader {
public:
	WardShader(Renderer *renderer, 
			const Texture *diffuseColor,
			const Texture *specularColor,
			const Texture *alphaU,
			const Texture *alphaV) : Shader(renderer, EBSDFShader), 
			m_diffuseReflectance(diffuseColor),
			m_specularReflectance(specularColor),
			m_alphaU(alphaU), m_alphaV(alphaV) {
		m_diffuseReflectanceShader = renderer->registerShaderForResource(m_diffuseReflectance.get());
		m_specularReflectanceShader = renderer->registerShaderForResource(m_specularReflectance.get());
		m_alphaUShader = renderer->registerShaderForResource(m_alphaU.get());
		m_alphaVShader = renderer->registerShaderForResource(m_alphaV.get());
	}

	bool isComplete() const {
		return m_diffuseReflectanceShader.get() != NULL &&
			   m_specularReflectanceShader.get() != NULL &&
			   m_alphaU.get() != NULL &&
			   m_alphaV.get() != NULL;
	}

	void putDependencies(std::vector<Shader *> &deps) {
		deps.push_back(m_diffuseReflectanceShader.get());
		deps.push_back(m_specularReflectanceShader.get());
		deps.push_back(m_alphaUShader.get());
		deps.push_back(m_alphaVShader.get());
	}

	void cleanup(Renderer *renderer) {
		renderer->unregisterShaderForResource(m_diffuseReflectance.get());
		renderer->unregisterShaderForResource(m_specularReflectance.get());
		renderer->unregisterShaderForResource(m_alphaU.get());
		renderer->unregisterShaderForResource(m_alphaV.get());
	}

	void generateCode(std::ostringstream &oss,
			const std::string &evalName,
			const std::vector<std::string> &depNames) const {
		oss << "vec3 " << evalName << "(vec2 uv, vec3 wi, vec3 wo) {" << endl
			<< "    if (wi.z <= 0.0 || wo.z <= 0.0)" << endl
			<< "    	return vec3(0.0);" << endl
			<< "    vec3 H = wi + wo;" << endl
			<< "    float cosSqr = H.z * H.z;" << endl
			<< "    float alphaU = max(0.3, " << depNames[2] << "(uv)[0]);" << endl
			<< "    float alphaV = max(0.3, " << depNames[3] << "(uv)[0]);" << endl
			<< "    float factor1 = dot(H, H)/(3.1415*alphaU*alphaV*cosSqr*cosSqr);"  << endl
			<< "    float factor2 = H.x / alphaU, factor3 = H.y / alphaV;" << endl
			<< "    float exponent = -(factor2*factor2 + factor3*factor3)/(H.z*H.z);" << endl
			<< "    float specRef = factor1 * exp(exponent);" << endl 
			<< "    return (" << depNames[0] << "(uv) * 0.31831" << endl
			<< "           + " << depNames[1] << "(uv) * specRef) * cosTheta(wo);" << endl
			<< "}" << endl
			<< "vec3 " << evalName << "_diffuse(vec2 uv, vec3 wi, vec3 wo) {" << endl
			<< "    if (wi.z <= 0.0 || wo.z <= 0.0)" << endl
			<< "    	return vec3(0.0);" << endl
			<< "    return " << depNames[0] << "(uv) * (0.31831 * cosTheta(wo));" << endl
			<< "}" << endl;
	}

	MTS_DECLARE_CLASS()
private:
	ref<const Texture> m_diffuseReflectance;
	ref<const Texture> m_specularReflectance;
	ref<const Texture> m_alphaU;
	ref<const Texture> m_alphaV;
	ref<Shader> m_diffuseReflectanceShader;
	ref<Shader> m_specularReflectanceShader;
	ref<Shader> m_alphaUShader;
	ref<Shader> m_alphaVShader;
};

Shader *Ward::createShader(Renderer *renderer) const { 
	return new WardShader(renderer, m_diffuseReflectance.get(),
		m_specularReflectance.get(), m_alphaU.get(), m_alphaV.get());
}

MTS_IMPLEMENT_CLASS(WardShader, false, Shader)
MTS_IMPLEMENT_CLASS_S(Ward, false, BSDF);
MTS_EXPORT_PLUGIN(Ward, "Anisotropic Ward BRDF");
MTS_NAMESPACE_END