/*
    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/render/texture.h>
#include <mitsuba/hw/gpuprogram.h>

MTS_NAMESPACE_BEGIN

/**
 * Mixture material, represents a linear combination of 
 * one or more BRDFs.
 */
class MixtureBSDF : public BSDF {
public:
	MixtureBSDF(const Properties &props) 
		: BSDF(props) {
		/* Parse the weight parameter */
		std::vector<std::string> weights = 
			tokenize(props.getString("weights", ""), " ,;");
		if (weights.size() == 0)
			Log(EError, "No weights were supplied!");
		m_weights.resize(weights.size());

		Float totalWeight = 0;
		char *end_ptr = NULL;
		for (size_t i=0; i<weights.size(); ++i) {
			Float weight = (Float) strtod(weights[i].c_str(), &end_ptr);
			if (*end_ptr != '\0')
				SLog(EError, "Could not parse the BRDF weights!");
			if (weight < 0)
				SLog(EError, "Invalid BRDF weight!");
			m_weights[i] = weight;
			totalWeight += weight;
		}

		if (totalWeight > 1)
			Log(EWarn, "Energy conservation is violated!");
		if (totalWeight == 0)
			Log(EError, "Combined weight must be > 0!");
	}

	MixtureBSDF(Stream *stream, InstanceManager *manager) 
	 : BSDF(stream, manager) {
		size_t bsdfCount = stream->readSize();
		m_weights.resize(bsdfCount);
		for (size_t i=0; i<bsdfCount; ++i) {
			m_weights[i] = stream->readFloat();
			BSDF *bsdf = static_cast<BSDF *>(manager->getInstance(stream));
			bsdf->incRef();
			m_bsdfs.push_back(bsdf);
		}
		configure();
	}

	virtual ~MixtureBSDF() {
		for (size_t i=0; i<m_bsdfs.size(); ++i)
			m_bsdfs[i]->decRef();
	}

	void serialize(Stream *stream, InstanceManager *manager) const {
		BSDF::serialize(stream, manager);
		
		stream->writeSize(m_bsdfs.size());
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			stream->writeFloat(m_weights[i]);
			manager->serialize(stream, m_bsdfs[i]);
		}
	}

	void configure() {
		m_usesRayDifferentials = false;
		size_t componentCount = 0;

		if (m_bsdfs.size() != m_weights.size())
			Log(EError, "BSDF count mismatch: " SIZE_T_FMT " bsdfs, but specified " SIZE_T_FMT " weights",
				m_bsdfs.size(), m_bsdfs.size());

		for (size_t i=0; i<m_bsdfs.size(); ++i)
			componentCount += m_bsdfs[i]->getComponentCount();

		m_pdf = DiscretePDF(m_bsdfs.size());
		m_components.reserve(componentCount);
		m_components.clear();
		m_indices.reserve(componentCount);
		m_indices.clear();
		m_offsets.reserve(m_bsdfs.size());
		m_offsets.clear();

		int offset = 0;
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			const BSDF *bsdf = m_bsdfs[i];
			m_offsets.push_back(offset);

			for (int j=0; j<bsdf->getComponentCount(); ++j) {
				int componentType = bsdf->getType(j);
				m_components.push_back(componentType);
				m_indices.push_back(std::make_pair((int) i, j));
			}

			offset += bsdf->getComponentCount();
			m_usesRayDifferentials |= bsdf->usesRayDifferentials();
			m_pdf[i] = m_weights[i];
		}
		m_pdf.build();
		BSDF::configure();
	}

	Spectrum eval(const BSDFQueryRecord &bRec, EMeasure measure) const {
		Spectrum result(0.0f);

		if (bRec.component == -1) {
			for (size_t i=0; i<m_bsdfs.size(); ++i)
				result += m_bsdfs[i]->eval(bRec, measure) * m_weights[i];
		} else {
			/* Pick out an individual component */
			int idx = m_indices[bRec.component].first;
			BSDFQueryRecord bRec2(bRec);
			bRec2.component = m_indices[bRec.component].second;
			return m_bsdfs[idx]->eval(bRec2, measure) * m_weights[idx];
		}

		return result;
	}

	Float pdf(const BSDFQueryRecord &bRec, EMeasure measure) const {
		Float result = 0.0f;

		if (bRec.component == -1) {
			for (size_t i=0; i<m_bsdfs.size(); ++i)
				result += m_bsdfs[i]->pdf(bRec, measure) * m_pdf[i];
		} else {
			/* Pick out an individual component */
			int idx = m_indices[bRec.component].first;
			BSDFQueryRecord bRec2(bRec);
			bRec2.component = m_indices[bRec.component].second;
			return m_bsdfs[idx]->pdf(bRec2, measure);
		}

		return result;
	}

	Spectrum sample(BSDFQueryRecord &bRec, Float &pdf, const Point2 &_sample) const {
		Point2 sample(_sample);
		if (bRec.component == -1) {
			/* Choose a component based on the normalized weights */
			size_t entry = m_pdf.sampleReuse(sample.x);

			Spectrum result = m_bsdfs[entry]->sample(bRec, pdf, sample);
			if (result.isZero()) // sampling failed
				return result;

			result *= m_weights[entry];
			pdf *= m_pdf[entry];

			EMeasure measure = BSDF::getMeasure(bRec.sampledType);
			for (size_t i=0; i<m_bsdfs.size(); ++i) {
				if (entry == i)
					continue;
				pdf += m_bsdfs[i]->pdf(bRec, measure) * m_pdf[i];
				result += m_bsdfs[i]->eval(bRec, measure) * m_weights[i];
			}

			bRec.sampledComponent += m_offsets[entry];
			return result;
		} else {
			/* Pick out an individual component */
			int requestedComponent = bRec.component;
			int bsdfIndex = m_indices[requestedComponent].first;
			bRec.component = m_indices[requestedComponent].second;
			Spectrum result = m_bsdfs[bsdfIndex]->sample(bRec, pdf, sample)
				* m_weights[bsdfIndex];
			bRec.component = bRec.sampledComponent = requestedComponent;
			return result;
		}
	}

	Spectrum sample(BSDFQueryRecord &bRec, const Point2 &_sample) const {
		Point2 sample(_sample);
		if (bRec.component == -1) {
			/* Choose a component based on the normalized weights */
			size_t entry = m_pdf.sampleReuse(sample.x);

			Float pdf;
			Spectrum result = m_bsdfs[entry]->sample(bRec, pdf, sample);
			if (result.isZero()) // sampling failed
				return result;

			result *= m_weights[entry];
			pdf *= m_pdf[entry];

			EMeasure measure = BSDF::getMeasure(bRec.sampledType);
			for (size_t i=0; i<m_bsdfs.size(); ++i) {
				if (entry == i)
					continue;
				pdf += m_bsdfs[i]->pdf(bRec, measure) * m_pdf[i];
				result += m_bsdfs[i]->eval(bRec, measure) * m_weights[i];
			}

			bRec.sampledComponent += m_offsets[entry];
			return result / pdf;
		} else {
			/* Pick out an individual component */
			int requestedComponent = bRec.component;
			int bsdfIndex = m_indices[requestedComponent].first;
			bRec.component = m_indices[requestedComponent].second;
			Spectrum result = m_bsdfs[bsdfIndex]->sample(bRec, sample)
				* m_weights[bsdfIndex];
			bRec.component = bRec.sampledComponent = requestedComponent;
			return result;
		}
	}

	void addChild(const std::string &name, ConfigurableObject *child) {
		if (child->getClass()->derivesFrom(MTS_CLASS(BSDF))) {
			BSDF *bsdf = static_cast<BSDF *>(child);
			m_bsdfs.push_back(bsdf);
			bsdf->incRef();
		} else {
			BSDF::addChild(name, child);
		}
	}

	std::string toString() const {
		std::ostringstream oss;
		oss << "MixtureBSDF[" << endl
			<< "  weights = {";
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			oss << " " << m_weights[i];
			if (i + 1 < m_bsdfs.size())
				oss << ",";
		}
		oss << " }," << endl
			<< "  bsdfs = {" << endl;
		for (size_t i=0; i<m_bsdfs.size(); ++i) 
			oss << "    " << indent(m_bsdfs[i]->toString(), 2) << "," << endl;
		oss << "  }" << endl
			<< "]";
		return oss.str();
	}

	Shader *createShader(Renderer *renderer) const;

	MTS_DECLARE_CLASS()
private:
	std::vector<Float> m_weights;
	std::vector<std::pair<int, int> > m_indices;
	std::vector<int> m_offsets;
	std::vector<BSDF *> m_bsdfs;
	DiscretePDF m_pdf;
};

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

class MixtureBSDFShader : public Shader {
public:
	MixtureBSDFShader(Renderer *renderer, const std::vector<BSDF *> &bsdfs, const std::vector<Float> &weights) 
		: Shader(renderer, EBSDFShader), m_bsdfs(bsdfs), m_weights(weights), m_complete(false) {
		m_bsdfShader.resize(bsdfs.size());
		for (size_t i=0; i<bsdfs.size(); ++i) {
			ref<Shader> shader = renderer->registerShaderForResource(bsdfs[i]);
			if (shader) {
				shader->incRef();
				/* At least one shader has a hardware implementation */
				m_complete = true;
			}
			m_bsdfShader[i] = shader;
		}
	}

	bool isComplete() const {
		return m_complete;
	}

	void cleanup(Renderer *renderer) {
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			renderer->unregisterShaderForResource(m_bsdfs[i]);
			if (m_bsdfShader[i])
				m_bsdfShader[i]->decRef();
		}
		m_bsdfShader.clear();
	}

	void putDependencies(std::vector<Shader *> &deps) {
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			if (m_bsdfShader[i])
				deps.push_back(m_bsdfShader[i]);
		}
	}

	void generateCode(std::ostringstream &oss,
			const std::string &evalName,
			const std::vector<std::string> &depNames) const {
		Assert(m_complete);
		int ctr = 0;
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			if (!m_bsdfShader[i])
				continue;
			oss << "uniform float " << evalName << "_weight_" << ctr++ << ";" << endl;
		}
		oss << endl;
		oss << "vec3 " << evalName << "(vec2 uv, vec3 wi, vec3 wo) {" << endl
			<< "    return ";
		ctr = 0;
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			if (!m_bsdfShader[i])
				continue;
			oss << endl << "      ";
			if (ctr != 0)
				oss << "+ ";
			else
				oss << "  ";
			oss << depNames[ctr] << "(uv, wi, wo) * "
							  << evalName << "_weight_" << ctr;
			ctr++;
		}
		oss << ";" << endl << "}" << endl << endl;
		oss << "vec3 " << evalName << "_diffuse(vec2 uv, vec3 wi, vec3 wo) {" << endl
			<< "    return ";
		ctr = 0;
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			if (!m_bsdfShader[i])
				continue;
			oss << endl << "      ";
			if (ctr != 0)
				oss << "+ ";
			else
				oss << "  ";
			oss << depNames[ctr] << "_diffuse(uv, wi, wo) * "
							  << evalName << "_weight_" << ctr;
			ctr++;
		}
		oss << ";" << endl << "}" << endl;
	}

	void resolve(const GPUProgram *program, const std::string &evalName, std::vector<int> &parameterIDs) const {
		int ctr = 0;
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			if (!m_bsdfShader[i])
				continue;
			parameterIDs.push_back(
				program->getParameterID(formatString("%s_weight_%i", evalName.c_str(), ctr++)));
		}
	}

	void bind(GPUProgram *program, const std::vector<int> &parameterIDs, int &textureUnitOffset) const {
		int ctr = 0;
		for (size_t i=0; i<m_bsdfs.size(); ++i) {
			if (!m_bsdfShader[i])
				continue;
			program->setParameter(parameterIDs[ctr++], m_weights[i]);
		}
	}


	MTS_DECLARE_CLASS()
private:
	std::vector<Shader *> m_bsdfShader;
	const std::vector<BSDF *> &m_bsdfs;
	const std::vector<Float> &m_weights;
	bool m_complete;
};

Shader *MixtureBSDF::createShader(Renderer *renderer) const { 
	return new MixtureBSDFShader(renderer, m_bsdfs, m_weights);
}

MTS_IMPLEMENT_CLASS(MixtureBSDFShader, false, Shader)
MTS_IMPLEMENT_CLASS_S(MixtureBSDF, false, BSDF)
MTS_EXPORT_PLUGIN(MixtureBSDF, "Mixture BRDF")
MTS_NAMESPACE_END