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

    Copyright (c) 2007-2014 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 "microfacet.h"
#include "rtrans.h"
#include "ior.h"

MTS_NAMESPACE_BEGIN

/*!\plugin{roughcoating}{Rough dielectric coating}
 * \order{11}
 * \icon{bsdf_roughcoating}
 * \parameters{
 *     \parameter{distribution}{\String}{
 *          Specifies the type of microfacet normal distribution
 *          used to model the surface roughness.
 *          \vspace{-1mm}
 *       \begin{enumerate}[(i)]
 *           \item \code{beckmann}: Physically-based distribution derived from
 *               Gaussian random surfaces. This is the default.\vspace{-1.5mm}
 *           \item \code{ggx}: The GGX \cite{Walter07Microfacet} distribution (also known as
 *               Trowbridge-Reitz \cite{Trowbridge19975Average} distribution)
 *               was designed to better approximate the long tails observed in measurements
 *               of ground surfaces, which are not modeled by the Beckmann distribution.
 *           \vspace{-1.5mm}
 *           \item \code{phong}: Classical Phong distribution.
 *              In most cases, the \code{ggx} and \code{beckmann} distributions
 *              should be preferred, since they provide better importance sampling
 *              and accurate shadowing/masking computations.
 *              \vspace{-4mm}
 *       \end{enumerate}
 *     }
 *     \parameter{alpha}{\Float\Or\Texture}{
 *         Specifies the roughness of the unresolved surface micro-geometry.
 *         When the Beckmann distribution is used, this parameter is equal to the
 *         \emph{root mean square} (RMS) slope of the microfacets.
 *         \default{0.1}.
 *     }
 *     \parameter{sampleVisible}{\Boolean}{
 *         Enables an improved importance sampling technique. Refer to
 *         pages \pageref{plg:roughconductor} and \pageref{sec:visiblenormal-sampling}
 *         for details. \default{\code{true}}
 *     }
 *     \parameter{intIOR}{\Float\Or\String}{Interior index of refraction specified
 *      numerically or using a known material name. \default{\texttt{bk7} / 1.5046}}
 *     \parameter{extIOR}{\Float\Or\String}{Exterior index of refraction specified
 *      numerically or using a known material name. \default{\texttt{air} / 1.000277}}
 *     \parameter{thickness}{\Float}{Denotes the thickness of the layer (to
 *      model absorption --- should be specified in inverse units of \code{sigmaA})\default{1}}
 *     \parameter{sigmaA}{\Spectrum\Or\Texture}{The absorption coefficient of the
 *      coating layer. \default{0, i.e. there is no absorption}}
 *     \parameter{specular\showbreak Reflectance}{\Spectrum\Or\Texture}{Optional
 *         factor that can be used to modulate the specular reflection component. Note
 *         that for physical realism, this parameter should never be touched. \default{1.0}}
 *     \parameter{\Unnamed}{\BSDF}{A nested BSDF model that should be coated.}
 * }
 * \renderings{
 *     \rendering{Rough gold coated with a \emph{smooth} varnish layer}
 *         {bsdf_roughcoating_gold_smooth}
 *     \rendering{Rough gold coated with a \emph{rough} ($\alpha\!=\!0.03$) varnish layer}
 *         {bsdf_roughcoating_gold_rough}
 * }
 *
 * This plugin implements a \emph{very} approximate\footnote{
 * The model only accounts for roughness
 * in the specular reflection and Fresnel transmittance through the interface.
 * The interior model receives incident illumination
 * that is transformed \emph{as if} the coating was smooth. While
 * that's not quite correct, it is a convenient workaround when the
 * \pluginref{coating} plugin produces specular highlights that are too sharp.}
 * model that simulates a rough dielectric coating. It is essentially the
 * roughened version of \pluginref{coating}.
 * Any BSDF in Mitsuba can be coated using this plugin and multiple coating
 * layers can even be applied in sequence, which allows designing interesting
 * custom materials. The coating layer can optionally be tinted (i.e. filled
 * with an absorbing medium), in which case this model also accounts for the
 * directionally dependent absorption within the layer.
 *
 * Note that the plugin discards illumination that undergoes internal
 * reflection within the coating. This can lead to a noticeable energy
 * loss for materials that reflect much of their energy near or below the critical
 * angle (i.e. diffuse or very rough materials).
 *
 * The implementation here is motivated by the paper
 * ``Arbitrarily Layered Micro-Facet Surfaces'' by Weidlich and
 * Wilkie \cite{Weidlich2007Arbitrarily}, though the implementation
 * works differently.
 */
class RoughCoating : public BSDF {
public:
    RoughCoating(const Properties &props) : BSDF(props) {
        /* Specifies the internal index of refraction at the interface */
        Float intIOR = lookupIOR(props, "intIOR", "bk7");

        /* Specifies the external index of refraction at the interface */
        Float extIOR = lookupIOR(props, "extIOR", "air");

        if (intIOR < 0 || extIOR < 0 || intIOR == extIOR)
            Log(EError, "The interior and exterior indices of "
                "refraction must be positive and differ!");

        m_eta = intIOR / extIOR;

        /* Specifies the layer's thickness using the inverse units of sigmaA */
        m_thickness = props.getFloat("thickness", 1);

        /* Specifies a multiplier for the specular reflectance component */
        m_specularReflectance = new ConstantSpectrumTexture(
            props.getSpectrum("specularReflectance", Spectrum(1.0f)));

        m_sigma = props.getFloat("sigma", 0.5);
    }

    RoughCoating(Stream *stream, InstanceManager *manager)
     : BSDF(stream, manager) {
        m_nested = static_cast<BSDF *>(manager->getInstance(stream));
        m_specularReflectance = static_cast<Texture *>(manager->getInstance(stream));
        m_sigma = stream->readFloat(); 
        m_eta = stream->readFloat();
        m_thickness = stream->readFloat();

        configure();
    }

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

        manager->serialize(stream, m_nested.get());
        manager->serialize(stream, m_specularReflectance.get());
        stream->writeFloat(m_sigma);
        stream->writeFloat(m_eta);
        stream->writeFloat(m_thickness);
    }

    void configure() {
        if (!m_nested)
            Log(EError, "A child BSDF instance is required");

        m_components.clear();
        for (int i=0; i<m_nested->getComponentCount(); ++i)
            m_components.push_back(m_nested->getType(i));
        m_components.push_back(EGlossyReflection | EFrontSide | EBackSide
            | (m_specularReflectance->isConstant() ? 0 : ESpatiallyVarying));

        BSDF::configure();
    }

    /// Helper function: reflect \c wi with respect to a given surface normal
    inline Vector reflect(const Vector &wi, const Normal &m) const {
        return 2 * dot(wi, m) * Vector(m) - wi;
    }
    
    auto interferenceHGTerms(const BSDFSamplingRecord &bRec) const {
        CoherenceProp interf{};
        ForEachHGCoeff(n,m)
            interf.hg[n][m] = bRec.csd.eval(n,m,{0,0,m_thickness*1e-6f});
        return interf;
    }
    
    auto phase_shift_term(const Vector3 &Phi) const {
        Spectrum ret;
        for (int i=0; i<SPECTRUM_SAMPLES; i++) 
            ret[i] = cos(Phi.z * m_thickness*1e-6f * CoherenceProp::k(i));
        return ret;
    }
    
    virtual CoherenceProp lightMatterInteraction(const BSDFSamplingRecord &bRec) const override {
        Assert(!!bRec.csd && "Not implemented");
        
        bool sampleSpecular = (bRec.typeMask & EGlossyReflection)
            && (bRec.component == -1 || bRec.component == (int) m_components.size()-1);
        
        auto csd = m_nested->lightMatterInteraction(bRec);
        if (!sampleSpecular)
            return csd;
        
        const Vector Phi = bRec.wo+bRec.wi;
        const auto c = phase_shift_term(Phi);
        
        const auto interf_hgs = interferenceHGTerms(bRec);
        ForEachHGCoeff(n,m)
            csd.hg[n][m] += c*interf_hgs.hg[n][m];

        return csd;
    }
    
    Spectrum PSD(const BSDFSamplingRecord &bRec, const Vector3 &Phi) const {
        MicrofacetDistribution distr(
            MicrofacetDistribution::EBeckmann,
            m_sigma*sqrt(2),
            m_sigma*sqrt(2),
            true
        );
        
        auto H = normalize(Phi);
        H *= math::signum(Frame::cosTheta(H));

        return Spectrum(distr.G(bRec.wi, bRec.wo, H) * distr.eval(H));
    }
    
    Spectrum evalHG(const BSDFSamplingRecord &bRec) const override {
        Assert(!!bRec.csd && "Not implemented");
        
        bool sampleSpecular = (bRec.typeMask & EGlossyReflection)
            && (bRec.component == -1 || bRec.component == (int) m_components.size()-1);
        bool sampleNested = (bRec.typeMask & m_nested->getType() & BSDF::EAll)
            && (bRec.component == -1 || bRec.component < (int) m_components.size()-1);

        if (!sampleSpecular && !sampleNested)
            return Spectrum(.0f);
            
        /* Calculate the reflection half-vector */
        Vector Phi = bRec.wo+bRec.wi;
        Vector H = normalize(Phi);
        const auto c = phase_shift_term(Phi);

        /* Fresnel factor */
        const Float F = fresnelDielectricExt(dot(bRec.wi, H), m_eta);
        
        auto nested = m_nested->evalHG(bRec);
        nested *= (1.f-F)*(1.f-F);
        if (!sampleSpecular)
            return nested;

        const auto refl = m_specularReflectance->eval(bRec.its) * F * PSD(bRec, Phi);
        
        const auto interf_hgs = interferenceHGTerms(bRec);
        Spectrum interf = Spectrum(.0f);
        ForEachHGCoeff(n,m)
            interf += interf_hgs.hg[n][m] / bRec.csd.eval(n,m,{0,0,0});

        // return Spectrum((bRec.csd.Q*Vector3{0,0,CoherenceProp::k(8)*m_thickness*1e-6f}).y/bRec.csd.s.length())*100.f;
        // // return c.abs()*interf*refl*nested;
        // return refl;
        return nested + refl;// + c*interf*refl*nested;
    }

    Spectrum eval(const BSDFSamplingRecord &bRec, EMeasure measure) const {
        /* Stop if this component was not requested */
        if (Frame::cosTheta(bRec.wi) <= 0 ||
            Frame::cosTheta(bRec.wo) <= 0 ||
            (!(bRec.typeMask & EGlossyReflection)))
            return Spectrum(0.0f);
    
        return evalHG(bRec);
    }

    Float pdf(const BSDFSamplingRecord &bRec, EMeasure measure) const {
        bool sampleSpecular = (bRec.typeMask & EGlossyReflection)
            && (bRec.component == -1 || bRec.component == (int) m_components.size()-1);
        bool sampleNested = (bRec.typeMask & m_nested->getType() & BSDF::EAll)
            && (bRec.component == -1 || bRec.component < (int) m_components.size()-1);

        Vector Phi = bRec.wo+bRec.wi;
        Vector H = normalize(Phi);
        
        const auto dwh_dwo = 1.0f / (4.0f * dot(bRec.wo, H));
        H *= math::signum(Frame::cosTheta(H));
        
        MicrofacetDistribution distr(
            MicrofacetDistribution::EBeckmann,
            m_sigma*sqrt(2),
            m_sigma*sqrt(2),
            true
        );
        const Float F = fresnelDielectricExt(dot(bRec.wi, H), m_eta);
        const auto probSpecular = F * distr.pdf(math::signum(Frame::cosTheta(bRec.wi)) * bRec.wi, H) * dwh_dwo;
        
        if (measure == ESolidAngle && sampleSpecular) {
            return sampleNested ? probSpecular : 1.0f;
        } else if (sampleNested) {
            const auto pdf = m_nested->pdf(bRec,measure);
            return sampleSpecular ? (pdf * (1 - probSpecular)) : pdf;
        }
        return .0f;
    }

    Spectrum sample(BSDFSamplingRecord &bRec, Float &pdf, const Point2 &_sample) const {
        bool sampleSpecular = (bRec.typeMask & EGlossyReflection)
            && (bRec.component == -1 || bRec.component == (int) m_components.size()-1);
        bool sampleNested = (bRec.typeMask & m_nested->getType() & BSDF::EAll)
            && (bRec.component == -1 || bRec.component < (int) m_components.size()-1);
        if (Frame::cosTheta(bRec.wo) * Frame::cosTheta(bRec.wi) <= 0)
            return Spectrum(.0f);

        MicrofacetDistribution distr(
            MicrofacetDistribution::EBeckmann,
            m_sigma*sqrt(2),
            m_sigma*sqrt(2),
            true
        );
        Vector Phi = bRec.wo+bRec.wi;
        Vector H = normalize(Phi);
        H *= math::signum(Frame::cosTheta(H));
        const Float F = fresnelDielectricExt(dot(bRec.wi, H), m_eta);
        
        Float probSpecular = F;
        bool choseSpecular = sampleSpecular;

        Point2 sample(_sample);
        if (sampleSpecular && sampleNested) {
            if (sample.x < probSpecular) {
                sample.x /= probSpecular;
            } else {
                sample.x = (sample.x - probSpecular) / (1 - probSpecular);
                choseSpecular = false;
            }
        }

        if (choseSpecular) {
            const Normal m = distr.sample(math::signum(Frame::cosTheta(bRec.wi)) * bRec.wi, sample);
            bRec.sampledComponent = (int)m_components.size() - 1;
            bRec.wo = reflect(bRec.wi, m);
            bRec.eta = 1.0f;
            bRec.sampledType = EGlossyReflection;
            return m_specularReflectance->eval(bRec.its) * F;
        }
        else {
            auto result = m_nested->sample(bRec,pdf,sample);
            
            if (sampleSpecular) {
                result /= 1.0f - probSpecular;
            }
            result *= (1 - F) * (1 - F);
            if (result.isZero())
                return Spectrum(.0f);
        }
        
        const auto measure = getMeasure(bRec.sampledType);
        pdf = this->pdf(bRec,measure);
        if (pdf == .0f || !bRec.csd)
            return Spectrum(.0f);
        return eval(bRec,measure);
    }

    Spectrum sample(BSDFSamplingRecord &bRec, const Point2 &sample) const {
        Float pdf;
        return RoughCoating::sample(bRec, pdf, sample);
    }

    Float getRoughness(const Intersection &its, int component) const {
        return component < (int) m_components.size() - 1
            ? m_nested->getRoughness(its, component)
            : m_sigma*sqrt(2);
    }

    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 BRDF can be added!");
            m_nested = static_cast<BSDF *>(child);
        } else {
            BSDF::addChild(name, child);
        }
    }

    std::string toString() const {
        std::ostringstream oss;
        oss << "RoughCoating[" << endl
            << "  id = \"" << getID() << "\"," << endl
            << "  sigma = " << indent(std::to_string(m_sigma)) << "," << endl
            << "  specularReflectance = " << indent(m_specularReflectance->toString()) << "," << endl
            << "  eta = " << m_eta << "," << endl
            << "  nested = " << indent(m_nested.toString()) << endl
            << "]";
        return oss.str();
    }

    MTS_DECLARE_CLASS()
private:
    ref<Texture> m_specularReflectance;
    ref<BSDF> m_nested;
    Float m_eta, m_sigma, m_k{ .0f };
    Float m_thickness;
};

MTS_IMPLEMENT_CLASS_S(RoughCoating, false, BSDF)
MTS_EXPORT_PLUGIN(RoughCoating, "Rough coating BSDF");
MTS_NAMESPACE_END
