#include <mitsuba/render/scene.h>
#include <mitsuba/render/texture.h>
#include <mitsuba/core/warp.h>
#include <mitsuba/hw/basicshader.h>

MTS_NAMESPACE_BEGIN



static float pdot(Vector a, Vector b) {
	return std::max(0.f, dot(a, b));
}

static Vector wt(Vector wp) {
	return normalize(Vector(-wp.x, -wp.y, 0.f));
}

static float G1(Vector wp, Vector w) {
	return std::min(1.f,
		std::max(0.f, Frame::cosTheta(w)) * std::max(0.f, Frame::cosTheta(wp))
		/ (pdot(w, wp) + pdot(w, wt(wp)) * Frame::sinTheta(wp))
	);
}

static float lambda_p(Vector wp, Vector wi) {
	float i_dot_p = pdot(wp, wi);
	return i_dot_p / (i_dot_p + pdot(wt(wp), wi) * Frame::sinTheta(wp));
}

static BSDFSamplingRecord getPerturbedRecord(const BSDFSamplingRecord& bRec, Frame frame, Intersection& perturbedIts) {
	perturbedIts.geoFrame = frame;
	perturbedIts.shFrame = frame;

	BSDFSamplingRecord query_wm(perturbedIts,
		frame.toLocal(bRec.its.toWorld(bRec.wi)),
		frame.toLocal(bRec.its.toWorld(bRec.wo)), bRec.mode);
	query_wm.sampler = bRec.sampler;
	query_wm.typeMask = bRec.typeMask;
	query_wm.component = bRec.component;
	return query_wm;
}

class NormalMapMicrofacet : public BSDF {
public:
	NormalMapMicrofacet(const Properties &props) : BSDF(props), rng(new Random) {
		maxScatteringOrder = props.getInteger("maxScatteringOrder", 2);
		saturation = props.getFloat("saturation", 0);
	}

	NormalMapMicrofacet(Stream *stream, InstanceManager *manager)
			: BSDF(stream, manager), rng(new Random) {
		m_normals = static_cast<Texture *>(manager->getInstance(stream));
		m_nested = static_cast<BSDF *>(manager->getInstance(stream));
		maxScatteringOrder = stream->readInt();
		saturation = stream->readFloat();

		configure();
	}

	void configure() {
		if (!m_nested)
			Log(EError, "A child BSDF instance is required");
		if (!m_normals)
			Log(EError, "A normal map texture must be specified");

		m_components.clear();
		for (int i=0; i<m_nested->getComponentCount(); ++i)
			m_components.push_back(m_nested->getType(i) | ESpatiallyVarying | EAnisotropic);

		m_usesRayDifferentials = true;
		BSDF::configure();
	}

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

		manager->serialize(stream, m_nested.get());
		manager->serialize(stream, m_normals.get());
		stream->writeInt(maxScatteringOrder);
		stream->writeFloat(saturation);
	}

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

	Spectrum getSpecularReflectance(const Intersection &its) const {
		return m_nested->getSpecularReflectance(its);
	}

	void addChild(const std::string &name, ConfigurableObject *child) {
		if (child->getClass()->derivesFrom(MTS_CLASS(BSDF))) {
			if (m_nested != NULL)
				Log(EError, "Only a single nested BSDF can be added!");
			m_nested = static_cast<BSDF *>(child);
		} else if (child->getClass()->derivesFrom(MTS_CLASS(Texture))) {
			if (m_normals != NULL)
				Log(EError, "Only a single normal texture can be specified!");
			const Properties &props = child->getProperties();
			if (props.getPluginName() == "bitmap" && !props.hasProperty("gamma"))
				Log(EError, "When using a bitmap texture as a normal map, please explicitly specify "
						"the 'gamma' parameter of the bitmap plugin. In most cases the following is the correct choice: "
						"<float name=\"gamma\" value=\"1.0\"/>");
			m_normals = static_cast<Texture *>(child);
		} else {
			BSDF::addChild(name, child);
		}

	}

	Spectrum eval(const BSDFSamplingRecord &bRec_, EMeasure measure) const {
		BSDFSamplingRecord bRec(bRec_);
		float weight = 2.;
		if (generateRandomNumber() < .5) {
			bRec.reverse();
			weight *= Frame::cosTheta(bRec.wi) / Frame::cosTheta(bRec.wo);
		}
		if (saturation > generateRandomNumber()) {
			return weight * evalSameMaterial(bRec, measure);
		} else {
			return weight * evalMirror(bRec, measure);
		}
	}

	Float pdf(const BSDFSamplingRecord &bRec, EMeasure measure) const {
		return (1.f - saturation) * pdfMirror(bRec, measure)
			+ saturation * pdfSameMaterial(bRec, measure);
	}

	Spectrum sample(BSDFSamplingRecord &bRec, const Point2 &sample) const {
		bRec.eta = 1.0f;
		bRec.sampledComponent = 0;
		bRec.sampledType = ESmooth;
		if (saturation > generateRandomNumber()) {
			return sampleSameMaterial(bRec, sample);
		} else {
			return sampleMirror(bRec, sample);
		}
	}

	Spectrum sample(BSDFSamplingRecord &bRec, Float &pdf_, const Point2 &sample_) const {
		Spectrum weight = sample(bRec, sample_);
		pdf_ = pdf(bRec, ESolidAngle);
		return weight;
	}

	Float getRoughness(const Intersection &its, int component) const {
		return m_nested->getRoughness(its, component);
	}

	std::string toString() const {
		std::ostringstream oss;
		oss << "NormalMapMicrofacet[" << endl
			<< "  id = \"" << getID() << "\"," << endl
			<< "  normals = " << indent(m_normals->toString()) << "," << endl
			<< "  nested = " << indent(m_nested->toString()) << ","  << endl
			<< "  maxScatteringOrder = " << maxScatteringOrder << ","  << endl
			<< "  saturation = " << saturation << endl
			<< "]";
		return oss.str();
	}

	Shader *createShader(Renderer *renderer) const;

	MTS_DECLARE_CLASS()
protected:
	ref<BSDF> m_nested;
	ref<Texture> m_normals;
	int maxScatteringOrder;
	float saturation;
	ref<Random> rng;

private:
	Spectrum evalSameMaterial(const BSDFSamplingRecord &bRec, EMeasure measure) const {
		if (measure != ESolidAngle
			|| Frame::cosTheta(bRec.wi) <= 0
			|| Frame::cosTheta(bRec.wo) <= 0)
			return Spectrum(0.0f);
		Vector wp;
		m_normals->eval(bRec.its, false).toLinearRGB(wp.x, wp.y, wp.z);
		wp = normalize(2.f * wp - Vector(1.f));

		if (Frame::cosTheta(wp) <= 0 || (std::abs(wp.x) < 1e-6 && std::abs(wp.y) < 1e-6)) {
			return m_nested->eval(bRec, measure);
		}

		Vector wr = -bRec.its.toWorld(bRec.wi);
		Spectrum energy(1.f);
		Spectrum value(0.f);
		float pdfWi;
		Frame frame_wp(bRec.its.toWorld(wp));
		Intersection perturbed_wp(bRec.its);
		perturbed_wp.geoFrame = frame_wp;
		perturbed_wp.shFrame = frame_wp;
		Frame frame_wt(bRec.its.toWorld(wt(wp)));
		Intersection perturbed_wt(bRec.its);
		perturbed_wt.geoFrame = frame_wt;
		perturbed_wt.shFrame = frame_wt;
		Vector currentWo = perturbed_wp.toLocal(bRec.its.toWorld(bRec.wo));
		Vector nextWo = perturbed_wt.toLocal(bRec.its.toWorld(bRec.wo));
		Intersection* currentIts = &perturbed_wp;
		Intersection* nextIts = &perturbed_wt;
		if (generateRandomNumber() > lambda_p(wp, bRec.wi)) {
			std::swap(currentIts, nextIts);
			std::swap(currentWo, nextWo);
		}
		// height \in [0, 1]
		float max_height = 1.f;
		float min_height = 1.f - G1(wp, bRec.wi);
		float facet_width = 1.f / Frame::tanTheta(wp);
		int i = 0;
		float shadowing = G1(wp, bRec.wo);
		float G1_;
		do {
			Point2 sample(generateRandomNumber(), generateRandomNumber());
			currentIts->wi = -currentIts->toLocal(wr);
			BSDFSamplingRecord samplingQuery(*currentIts, NULL);
			BSDFSamplingRecord evalQuery(*currentIts, currentIts->wi, currentWo);
			float weight;
			if (i > 0) {
				float pdfWo = m_nested->pdf(evalQuery, measure);
				pdfWo /= Frame::cosTheta(evalQuery.wo);
				weight = pdfWi / (pdfWi + pdfWo);
			} else {
				weight = .5f;
			}

			if (max_height > 1.f - shadowing) {
				float masking_shadowing = std::max(0.f,
					(max_height - std::max(min_height, 1.f - shadowing)));
				Spectrum contribution = m_nested->eval(evalQuery, measure)
					* energy * (masking_shadowing / (max_height - min_height) * weight);
				if (std::isfinite(contribution[0])) {
					value += contribution;
				}
			}

			Spectrum throughput = m_nested->sample(samplingQuery, sample);
			if (throughput.isZero()) {
				break;
			} else if (i == 0) {
				BSDFSamplingRecord queryEvalWeight(*currentIts, samplingQuery.wo);
				pdfWi = m_nested->pdf(queryEvalWeight, measure);
				pdfWi /= Frame::cosTheta(queryEvalWeight.wo);
			}

			energy *= throughput;
			wr = currentIts->toWorld(samplingQuery.wo);
			Vector wr_local = bRec.its.toLocal(wr);
			Vector position_min = Vector(0, 0, min_height);
			Vector position_max = Vector(0, 0, max_height);
			Vector other_normal = wp;
			if (currentIts == &perturbed_wp) {
				other_normal = wt(wp);
				position_min += min_height * facet_width * other_normal;
				position_max += max_height * facet_width * other_normal;
			}
			float d_min = -dot(position_min, other_normal) / dot(wr_local, other_normal);
			float d_max = -dot(position_max, other_normal) / dot(wr_local, other_normal);
			if (d_min < 0 || d_max < 0) {
				break;
			} else {
				min_height += d_min * wr_local.z;
				max_height += d_max * wr_local.z;
			}
			std::swap(currentIts, nextIts);
			std::swap(currentWo, nextWo);
			i++;
			G1_ = (max_height - std::max(min_height, 1.f)) / (max_height - min_height);
			max_height = std::min(1.f, max_height);
		} while (i < maxScatteringOrder && generateRandomNumber() > G1_);
		return value;
	}

	Spectrum evalMirror(const BSDFSamplingRecord &bRec, EMeasure measure) const {
		if (measure != ESolidAngle
			|| Frame::cosTheta(bRec.wi) <= 0
			|| Frame::cosTheta(bRec.wo) <= 0)
			return Spectrum(0.0f);
		Vector wp;
		m_normals->eval(bRec.its, false).toLinearRGB(wp.x, wp.y, wp.z);
		wp = normalize(2.f * wp - Vector(1.f));

		if (Frame::cosTheta(wp) <= 0 || (std::abs(wp.x) < 1e-6 && std::abs(wp.y) < 1e-6)) {
			return m_nested->eval(bRec, measure);
		}

		Frame frame_wp(bRec.its.toWorld(wp));
		Intersection perturbed_wp(bRec.its);
		perturbed_wp.geoFrame = frame_wp;
		perturbed_wp.shFrame = frame_wp;

		Vector wt_ = bRec.its.toWorld(wt(wp));
		Frame frame_wt(wt_);
		Intersection perturbed_wt(bRec.its);
		perturbed_wt.geoFrame = frame_wt;
		perturbed_wt.shFrame = frame_wt;

		Vector wo_wp = perturbed_wp.toLocal(bRec.its.toWorld(bRec.wo));
		Vector wo = bRec.its.toWorld(bRec.wo);
		Vector woReflected = normalize(wo - 2.0f * dot(wo, wt_) * wt_);
		float shadowing = G1(wp, bRec.wo);
		float reflected_min_height = 1.f - shadowing;
		float reflected_max_height = 1.f;
		{
			Vector wr_local = bRec.its.toLocal(-woReflected);
			Vector position_min = Vector(0, 0, reflected_min_height);
			Vector position_max = Vector(0, 0, reflected_max_height);
			Vector other_normal = wp;
			float d_min = -dot(position_min, other_normal) / dot(wr_local, other_normal);
			float d_max = -dot(position_max, other_normal) / dot(wr_local, other_normal);
			reflected_min_height += d_min * wr_local.z;
			reflected_max_height += d_max * wr_local.z;
		}
		woReflected = perturbed_wp.toLocal(woReflected);

		// height \in [0, 1]
		float max_height = 1.f;
		float min_height = 1.f - G1(wp, bRec.wi);
		float facet_width = 1.f / Frame::tanTheta(wp);
		int i = 0;
		float G1_;
		Spectrum energy(1.f);
		Spectrum value(0.f);
		float pdfWi;
		Vector wr = -bRec.its.toWorld(bRec.wi);
		// do one reflection if we start at wt
		if (generateRandomNumber() > lambda_p(wp, bRec.wi)) {
			wr = normalize(wr - 2.0f * dot(wr, wt_) * wt_);
			Vector wr_local = bRec.its.toLocal(wr);
			Vector position_min = Vector(0, 0, min_height);
			Vector position_max = Vector(0, 0, max_height);
			Vector other_normal = wp;
			float d_min = -dot(position_min, other_normal) / dot(wr_local, other_normal);
			float d_max = -dot(position_max, other_normal) / dot(wr_local, other_normal);
			min_height += d_min * wr_local.z;
			max_height += d_max * wr_local.z;
			i++;
		}

		bool singleScattering = true;
		do {
			Point2 sample(generateRandomNumber(), generateRandomNumber());
			perturbed_wp.wi = -perturbed_wp.toLocal(wr);
			BSDFSamplingRecord query_wp(perturbed_wp, NULL);

			// eval for wp
			BSDFSamplingRecord queryEval(perturbed_wp, perturbed_wp.wi, wo_wp);
			float weight;
			if (!singleScattering) {
				float pdfWo = m_nested->pdf(queryEval, measure);
				pdfWo /= Frame::cosTheta(queryEval.wo);
				weight = pdfWi / (pdfWi + pdfWo);
			} else {
				weight = .5f;
			}

			if (max_height > 1.f - shadowing) {
				float masking_shadowing = std::max(0.f,
					(max_height - std::max(min_height, 1.f - shadowing)));
				Spectrum contribution = m_nested->eval(queryEval, measure)
					* energy * (masking_shadowing / (max_height - min_height) * weight);
				if (std::isfinite(contribution[0])) {
					value += contribution;
				}
			}

			i++;
			if (i >= maxScatteringOrder) {
				break;
			}

			// eval reflection of wo
			BSDFSamplingRecord queryEvalReflection(perturbed_wp, perturbed_wp.wi, woReflected);
			if (!singleScattering) {
				float pdfWo = m_nested->pdf(queryEvalReflection, measure);
				pdfWo /= Frame::cosTheta(queryEvalReflection.wo);
				weight = pdfWi / (pdfWi + pdfWo);
			} else {
				weight = .5f;
			}
			if (min_height < reflected_max_height  && reflected_min_height < max_height) {
				// the intervals intersect
				float masking_shadowing = std::max(0.f,
					(std::min(max_height, reflected_max_height) - std::max(min_height, reflected_min_height)));
				Spectrum contribution = m_nested->eval(queryEvalReflection, measure)
					* energy * (masking_shadowing / (max_height - min_height) * weight);
				if (std::isfinite(contribution[0])) {
					value += contribution;
				}
			}

			Spectrum throughput = m_nested->sample(query_wp, sample);
			if (throughput.isZero()) {
				break;
			} else if (singleScattering) {
				BSDFSamplingRecord queryEvalWeight(perturbed_wp, query_wp.wo);
				pdfWi = m_nested->pdf(queryEvalWeight, measure);
				pdfWi /= Frame::cosTheta(queryEvalWeight.wo);
				singleScattering = false;
			}
			energy *= throughput;
			wr = perturbed_wp.toWorld(query_wp.wo);
			Vector wr_local = bRec.its.toLocal(wr);
			Vector position_min = Vector(0, 0, min_height);
			Vector position_max = Vector(0, 0, max_height);
			Vector other_normal = wt(wp);
			position_min += min_height * facet_width * other_normal;
			position_max += max_height * facet_width * other_normal;
			float d_min = -dot(position_min, other_normal) / dot(wr_local, other_normal);
			float d_max = -dot(position_max, other_normal) / dot(wr_local, other_normal);
			if (d_min < 0 || d_max < 0) {
				break;
			} else {
				min_height += d_min * wr_local.z;
				max_height += d_max * wr_local.z;
			}
			G1_ = (max_height - std::max(min_height, 1.f)) / (max_height - min_height);
			max_height = std::min(1.f, max_height);
			if (i >= maxScatteringOrder || generateRandomNumber() < G1_) break;

			// reflect on wt
			wr = normalize(wr + 2.0f * dot(-wr, wt_) * wt_);
			wr_local = bRec.its.toLocal(wr);
			position_min = Vector(0, 0, min_height);
			position_max = Vector(0, 0, max_height);
			other_normal = wp;
			d_min = -dot(position_min, other_normal) / dot(wr_local, other_normal);
			d_max = -dot(position_max, other_normal) / dot(wr_local, other_normal);
			if (d_min < 0 || d_max < 0) {
				break;
			} else {
				min_height += d_min * wr_local.z;
				max_height += d_max * wr_local.z;
			}
			G1_ = (max_height - std::max(min_height, 1.f)) / (max_height - min_height);
			max_height = std::min(1.f, max_height);
			i++;
		} while (i < maxScatteringOrder && generateRandomNumber() > G1_);
		return value;
	}

	Float pdfSameMaterial(const BSDFSamplingRecord &bRec, EMeasure measure) const {
		if (measure != ESolidAngle
			|| Frame::cosTheta(bRec.wi) <= 0
			|| Frame::cosTheta(bRec.wo) <= 0)
			return 0.0f;
		Vector wp;
		m_normals->eval(bRec.its, false).toLinearRGB(wp.x, wp.y, wp.z);
		wp = normalize(2.f * wp - Vector(1.f));

		if (Frame::cosTheta(wp) <= 0 || (std::abs(wp.x) < 1e-6 && std::abs(wp.y) < 1e-6)) {
			return m_nested->pdf(bRec, measure);
		}

		Float probability_wp = lambda_p(wp, bRec.wi);
		float pdf = 0.f;
		if (probability_wp > 0.f) {
			Intersection perturbed_wp(bRec.its);
			Frame frame_wp(bRec.its.toWorld(wp));
			BSDFSamplingRecord query_wp = getPerturbedRecord(bRec, frame_wp, perturbed_wp);
			query_wp.sampledComponent = bRec.sampledComponent;
			query_wp.sampledType = bRec.sampledType;
			query_wp.eta = bRec.eta;
			pdf += probability_wp * m_nested->pdf(query_wp, measure);
		}
		if (probability_wp < 1.f) {
			Intersection perturbed_wt(bRec.its);
			Frame frame_wt(bRec.its.toWorld(wt(wp)));
			BSDFSamplingRecord query_wt = getPerturbedRecord(bRec, frame_wt, perturbed_wt);
			query_wt.sampledComponent = bRec.sampledComponent;
			query_wt.sampledType = bRec.sampledType;
			query_wt.eta = bRec.eta;
			pdf += (1.f - probability_wp) * m_nested->pdf(query_wt, measure);
		}
		pdf *= G1(wp, bRec.wo);
		// single scattering separable + diffuse
		return pdf + .1f * warp::squareToCosineHemispherePdf(bRec.wo);
	}

	Float pdfMirror(const BSDFSamplingRecord &bRec, EMeasure measure) const {
		if (measure != ESolidAngle
			|| Frame::cosTheta(bRec.wi) <= 0
			|| Frame::cosTheta(bRec.wo) <= 0)
			return 0.0f;
		Vector wp;
		m_normals->eval(bRec.its, false).toLinearRGB(wp.x, wp.y, wp.z);
		wp = normalize(2.f * wp - Vector(1.f));

		if (Frame::cosTheta(wp) <= 0 || (std::abs(wp.x) < 1e-6 && std::abs(wp.y) < 1e-6)) {
			return m_nested->pdf(bRec, measure);
		}

		Float probability_wp = lambda_p(wp, bRec.wi);
		Frame frameWp(bRec.its.toWorld(wp));
		Intersection perturbed_wp(bRec.its);
		perturbed_wp.geoFrame = frameWp;
		perturbed_wp.shFrame = frameWp;
		Vector wi = bRec.its.toWorld(bRec.wi);
		Vector wo = bRec.its.toWorld(bRec.wo);
		Vector wt_ = bRec.its.toWorld(wt(wp));
		float pdf = 0.f;
		if (probability_wp > 0.f) {
			BSDFSamplingRecord queryWp(perturbed_wp, perturbed_wp.toLocal(wi),
				perturbed_wp.toLocal(wo));
			pdf += probability_wp * m_nested->pdf(queryWp, measure);

			if (dot(wo, wt_) > 1e-6) {
				Vector woReflected = normalize(wo - 2.0f * dot(wo, wt_) * wt_);
				BSDFSamplingRecord queryWpt(perturbed_wp, perturbed_wp.toLocal(wi),
					perturbed_wp.toLocal(woReflected));
				pdf += probability_wp * m_nested->pdf(queryWpt, measure)
					* (1.f - G1(wp, bRec.its.toLocal(woReflected)));
			}
		}

		if (probability_wp < 1.f && dot(wi, wt_) > 1e-6) {
			Vector wiReflected = normalize(wi - 2.0f * dot(wi, wt_) * wt_);
			BSDFSamplingRecord queryWtp(perturbed_wp, perturbed_wp.toLocal(wiReflected),
				perturbed_wp.toLocal(wo));
			pdf += (1.f - probability_wp) * m_nested->pdf(queryWtp, measure);
		}
		pdf *= G1(wp, bRec.wo);

		// double scattering separable + diffuse
		return pdf + .1f * warp::squareToCosineHemispherePdf(bRec.wo);
	}

	Spectrum sampleSameMaterial(BSDFSamplingRecord &bRec, const Point2 &sample) const {
		if (Frame::cosTheta(bRec.wi) <= 0)
			return Spectrum(0.0f);

		Vector wp;
		m_normals->eval(bRec.its, false).toLinearRGB(wp.x, wp.y, wp.z);
		wp = normalize(2.f * wp - Vector(1.f));

		if (Frame::cosTheta(wp) <= 0 || (std::abs(wp.x) < 1e-6 && std::abs(wp.y) < 1e-6)) {
			return m_nested->sample(bRec, sample);
		}

		Frame frame_wp(bRec.its.toWorld(wp));
		Intersection perturbed_wp(bRec.its);
		perturbed_wp.geoFrame = frame_wp;
		perturbed_wp.shFrame = frame_wp;

		Frame frame_wt(bRec.its.toWorld(wt(wp)));
		Intersection perturbed_wt(bRec.its);
		perturbed_wt.geoFrame = frame_wt;
		perturbed_wt.shFrame = frame_wt;

		Intersection* currentIts = &perturbed_wp;
		Intersection* nextIts = &perturbed_wt;
		if (generateRandomNumber() > lambda_p(wp, bRec.wi)) {
			std::swap(currentIts, nextIts);
		}
		// height \in [0, 1]
		float height = 1.f - generateRandomNumber() * G1(wp, bRec.wi);
		float facet_width = 1.f / Frame::tanTheta(wp);
		int i = 0;
		Vector wr = -bRec.its.toWorld(bRec.wi);
		Spectrum energy(1.f);
		do {
			if (i >= maxScatteringOrder) {
				return Spectrum(0.f);
			}
			currentIts->wi = -currentIts->toLocal(wr);
			BSDFSamplingRecord samplingQuery(*currentIts, NULL);
			Point2 nextSample(generateRandomNumber(), generateRandomNumber());
			energy *= m_nested->sample(samplingQuery, nextSample);
			if (energy.isZero()) {
				return energy;
			}
			wr = currentIts->toWorld(samplingQuery.wo);
			Vector wr_local = bRec.its.toLocal(wr);
			Vector position = Vector(0, 0, height);
			Vector other_normal = wp;
			if (currentIts == &perturbed_wp) {
				other_normal = wt(wp);
				position += height * facet_width * other_normal;
			}
			float d = -dot(position, other_normal) / dot(wr_local, other_normal);
			if (d < 0) {
				break;
			} else {
				height += d * wr_local.z;
			}
			std::swap(currentIts, nextIts);
			i++;
		} while (height <= 1.f);
		bRec.wo = bRec.its.toLocal(wr);
		if (Frame::cosTheta(bRec.wo) <= 0.) {
			return Spectrum(0.f);
		}
		return energy;
	}

	Spectrum sampleMirror(BSDFSamplingRecord &bRec, const Point2 &sample) const {
		if (Frame::cosTheta(bRec.wi) <= 0)
			return Spectrum(0.0f);

		Vector wp;
		m_normals->eval(bRec.its, false).toLinearRGB(wp.x, wp.y, wp.z);
		wp = normalize(2.f * wp - Vector(1.f));

		if (Frame::cosTheta(wp) <= 0 || (std::abs(wp.x) < 1e-6 && std::abs(wp.y) < 1e-6)) {
			return m_nested->sample(bRec, sample);
		}

		Frame frame_wp(bRec.its.toWorld(wp));
		Intersection perturbed_wp(bRec.its);
		perturbed_wp.geoFrame = frame_wp;
		perturbed_wp.shFrame = frame_wp;

		Vector wt_ = bRec.its.toWorld(wt(wp));
		Frame frame_wt(wt_);
		Intersection perturbed_wt(bRec.its);
		perturbed_wt.geoFrame = frame_wt;
		perturbed_wt.shFrame = frame_wt;

		// height \in [0, 1]
		float height = 1.f - generateRandomNumber() * G1(wp, bRec.wi);
		float facet_width = 1.f / Frame::tanTheta(wp);
		int i = 0;
		Vector wr = -bRec.its.toWorld(bRec.wi);
		Spectrum energy(1.f);
		// do one reflection if we start at wt
		if (generateRandomNumber() > lambda_p(wp, bRec.wi)) {
			wr = normalize(wr - 2.0f * dot(wr, wt_) * wt_);
			Vector wr_local = bRec.its.toLocal(wr);
			Vector position = Vector(0, 0, height);
			Vector other_normal = wp;
			float d = -dot(position, other_normal) / dot(wr_local, other_normal);
			height += d * wr_local.z;
			i++;
		}
		do {
			// sample brdf on wp
			if (i >= maxScatteringOrder) {
				return Spectrum(0.f);
			}
			Point2 sample(generateRandomNumber(), generateRandomNumber());
			perturbed_wp.wi = -perturbed_wp.toLocal(wr);
			BSDFSamplingRecord query_wp(perturbed_wp, NULL);
			energy *= m_nested->sample(query_wp, sample);
			if (energy.isZero()) {
				return energy;
			}
			wr = perturbed_wp.toWorld(query_wp.wo);
			Vector wr_local = bRec.its.toLocal(wr);
			Vector position = Vector(0, 0, height);
			Vector other_normal = wt(wp);
			position += height * facet_width * other_normal;
			float d = -dot(position, other_normal) / dot(wr_local, other_normal);
			if (d < 0) {
				break;
			} else {
				height += d * wr_local.z;
			}
			i++;

			// reflect on wt
			if (height > 1.f) break;
			if (i >= maxScatteringOrder) {
				return Spectrum(0.f);
			}
			wr = normalize(wr - 2.0f * dot(wr, wt_) * wt_);
			wr_local = bRec.its.toLocal(wr);
			position = Vector(0, 0, height);
			other_normal = wp;
			d = -dot(position, other_normal) / dot(wr_local, other_normal);
			if (d < 0) {
				break;
			} else {
				height += d * wr_local.z;
			}
			i++;
		} while (height <= 1.f);
		bRec.wo = bRec.its.toLocal(wr);
		if (Frame::cosTheta(bRec.wo) <= 0.) {
			return Spectrum(0.f);
		}
		return energy;
	}

	float generateRandomNumber() const
	{
		return const_cast<Random*>(rng.get())->nextFloat();
	}
};

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

/**
 * This is a quite approximate version of the normal map model -- it likely
 * won't match the reference exactly, but it should be good enough for
 * preview purposes
 */
class NormalMapShader : public Shader {
public:
	NormalMapShader(Renderer *renderer, const BSDF *nested, const Texture *normals)
		: Shader(renderer, EBSDFShader), m_nested(nested), m_normals(normals) {
		m_nestedShader = renderer->registerShaderForResource(m_nested.get());
		m_normalShader = renderer->registerShaderForResource(m_normals.get());
	}

	bool isComplete() const {
		return m_nestedShader.get() != NULL;
	}

	void cleanup(Renderer *renderer) {
		renderer->unregisterShaderForResource(m_nested.get());
		renderer->unregisterShaderForResource(m_normals.get());
	}

	void putDependencies(std::vector<Shader *> &deps) {
		deps.push_back(m_nestedShader.get());
		deps.push_back(m_normalShader.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
			<< "    vec3 n = normalize(2.0*" << depNames[1] << "(uv) - vec3(1.0));" << endl
			<< "    vec3 s = normalize(vec3(1.0-n.x*n.x, -n.x*n.y, -n.x*n.z)); " << endl
			<< "    vec3 t = cross(s, n);" << endl
			<< "    wi = vec3(dot(wi, s), dot(wi, t), dot(wi, n));" << endl
			<< "    wo = vec3(dot(wo, s), dot(wo, t), dot(wo, n));" << endl
			<< "    return " << depNames[0] << "(uv, wi, wo);" << endl
			<< "}" << endl
			<< endl
		    << "vec3 " << evalName << "_diffuse(vec2 uv, vec3 wi, vec3 wo) {" << endl
			<< "    vec3 n = normalize(2.0*" << depNames[1] << "(uv) - vec3(1.0));" << endl
			<< "    vec3 s = normalize(vec3(1.0-n.x*n.x, -n.x*n.y, -n.x*n.z)); " << endl
			<< "    vec3 t = cross(s, n);" << endl
			<< "    wi = vec3(dot(wi, s), dot(wi, t), dot(wi, n));" << endl
			<< "    wo = vec3(dot(wo, s), dot(wo, t), dot(wo, n));" << endl
			<< "    return " << depNames[0] << "_diffuse(uv, wi, wo);" << endl
			<< "}" << endl
			<< endl;
	}

	MTS_DECLARE_CLASS()
private:
	ref<const BSDF> m_nested;
	ref<const Texture> m_normals;
	ref<Shader> m_nestedShader;
	ref<Shader> m_normalShader;
};

Shader *NormalMapMicrofacet::createShader(Renderer *renderer) const {
	return new NormalMapShader(renderer, m_nested.get(), m_normals.get());
}

MTS_IMPLEMENT_CLASS(NormalMapShader, false, Shader)
MTS_IMPLEMENT_CLASS_S(NormalMapMicrofacet, false, BSDF)
MTS_EXPORT_PLUGIN(NormalMapMicrofacet, "Normal map modifier");
MTS_NAMESPACE_END
