/*
    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/core/fresolver.h>
#include <mitsuba/render/bsdf.h>
#include <mitsuba/hw/basicshader.h>
#include <mitsuba/core/warp.h>
#include "microfacet.h"
#include "ior.h"

MTS_NAMESPACE_BEGIN

/*!\plugin{roughconductor}{Rough conductor material}
 * \order{7}
 * \icon{bsdf_roughconductor}
 * \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}: Anisotropic Phong distribution by
 *              Ashikhmin and Shirley \cite{Ashikhmin2005Anisotropic}.
 *              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, alphaU, alphaV}{\Float\Or\Texture}{
 *         Specifies the roughness of the unresolved surface micro-geometry
 *         along the tangent and bitangent directions. When the Beckmann
 *         distribution is used, this parameter is equal to the
 *         \emph{root mean square} (RMS) slope of the microfacets.
 *         \code{alpha} is a convenience parameter to initialize both
 *         \code{alphaU} and \code{alphaV} to the same value. \default{0.1}.
 *     }
 *     \parameter{material}{\String}{Name of a material preset, see
 *           \tblref{conductor-iors}.\!\default{\texttt{Cu} / copper}}
 *     \parameter{eta, k}{\Spectrum}{Real and imaginary components of the material's index of
 *             refraction \default{based on the value of \texttt{material}}}
 *     \parameter{extEta}{\Float\Or\String}{
 *           Real-valued index of refraction of the surrounding dielectric,
 *           or a material name of a dielectric \default{\code{air}}
 *     }
 *     \parameter{sampleVisible}{\Boolean}{
 *         Enables a sampling technique proposed by Heitz and D'Eon~\cite{Heitz1014Importance},
 *         which focuses computation on the visible parts of the microfacet normal
 *         distribution, considerably reducing variance in some cases.
 *         \default{\code{true}, i.e. use visible normal sampling}
 *     }
 *     \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}}
 * }
 * \vspace{3mm}
 * This plugin implements a realistic microfacet scattering model for rendering
 * rough conducting materials, such as metals. It can be interpreted as a fancy
 * version of the Cook-Torrance model and should be preferred over
 * heuristic models like \pluginref{phong} and \pluginref{ward} if possible.
 * \renderings{
 *     \rendering{Rough copper (Beckmann, $\alpha=0.1$)}
 *         {bsdf_roughconductor_copper.jpg}
 *     \rendering{Vertically brushed aluminium (Anisotropic Phong,
 *         $\alpha_u=0.05,\ \alpha_v=0.3$), see
 *         \lstref{roughconductor-aluminium}}
 *      {bsdf_roughconductor_anisotropic_aluminium.jpg}
 * }
 *
 * Microfacet theory describes rough surfaces as an arrangement of unresolved
 * and ideally specular facets, whose normal directions are given by a
 * specially chosen \emph{microfacet distribution}. By accounting for shadowing
 * and masking effects between these facets, it is possible to reproduce the
 * important off-specular reflections peaks observed in real-world measurements
 * of such materials.
 *
 * This plugin is essentially the ``roughened'' equivalent of the (smooth) plugin
 * \pluginref{conductor}. For very low values of $\alpha$, the two will
 * be identical, though scenes using this plugin will take longer to render
 * due to the additional computational burden of tracking surface roughness.
 *
 * The implementation is based on the paper ``Microfacet Models
 * for Refraction through Rough Surfaces'' by Walter et al.
 * \cite{Walter07Microfacet}. It supports three different types of microfacet
 * distributions and has a texturable roughness parameter.
 * To facilitate the tedious task of specifying spectrally-varying index of
 * refraction information, this plugin can access a set of measured materials
 * for which visible-spectrum information was publicly available
 * (see \tblref{conductor-iors} for the full list).
 * There is also a special material profile named \code{none}, which disables
 * the computation of Fresnel reflectances and produces an idealized
 * 100% reflecting mirror.
 *
 * When no parameters are given, the plugin activates the default settings,
 * which describe copper with a medium amount of roughness modeled using a
 * Beckmann distribution.
 *
 * To get an intuition about the effect of the surface roughness parameter
 * $\alpha$, consider the following approximate classification: a value of
 * $\alpha=0.001-0.01$ corresponds to a material with slight imperfections
 * on an otherwise smooth surface finish, $\alpha=0.1$ is relatively rough,
 * and $\alpha=0.3-0.7$ is \emph{extremely} rough (e.g. an etched or ground
 * finish). Values significantly above that are probably not too realistic.
 * \vspace{4mm}
 * \begin{xml}[caption={A material definition for brushed aluminium}, label=lst:roughconductor-aluminium]
 * <bsdf type="roughconductor">
 *     <string name="material" value="Al"/>
 *     <string name="distribution" value="phong"/>
 *     <float name="alphaU" value="0.05"/>
 *     <float name="alphaV" value="0.3"/>
 * </bsdf>
 * \end{xml}
 *
 * \subsubsection*{Technical details}
 * All microfacet distributions allow the specification of two distinct
 * roughness values along the tangent and bitangent directions. This can be
 * used to provide a material with a ``brushed'' appearance. The alignment
 * of the anisotropy will follow the UV parameterization of the underlying
 * mesh. This means that such an anisotropic material cannot be applied to
 * triangle meshes that are missing texture coordinates.
 *
 * \label{sec:visiblenormal-sampling}
 * Since Mitsuba 0.5.1, this plugin uses a new importance sampling technique
 * contributed by Eric Heitz and Eugene D'Eon, which restricts the sampling
 * domain to the set of visible (unmasked) microfacet normals. The previous
 * approach of sampling all normals is still available and can be enabled
 * by setting \code{sampleVisible} to \code{false}.
 * Note that this new method is only available for the \code{beckmann} and
 * \code{ggx} microfacet distributions. When the \code{phong} distribution
 * is selected, the parameter has no effect.
 *
 * When rendering with the Phong microfacet distribution, a conversion is
 * used to turn the specified Beckmann-equivalent $\alpha$ roughness value
 * into the exponent parameter of this distribution. This is done in a way,
 * such that the same value $\alpha$ will produce a similar appearance across
 * different microfacet distributions.
 *
 * When using this plugin, you should ideally compile Mitsuba with support for
 * spectral rendering to get the most accurate results. While it also works
 * in RGB mode, the computations will be more approximate in nature.
 * Also note that this material is one-sided---that is, observed from the
 * back side, it will be completely black. If this is undesirable,
 * consider using the \pluginref{twosided} BRDF adapter.
 */
class RoughConductor : public BSDF {
public:
    RoughConductor(const Properties &props) : BSDF(props) {
        ref<FileResolver> fResolver = Thread::getThread()->getFileResolver();

        m_specularReflectance = new ConstantSpectrumTexture(
            props.getSpectrum("specularReflectance", Spectrum(1.0f)));

        std::string materialName = props.getString("material", "Cu");

        Spectrum intEta, intK;
        if (boost::to_lower_copy(materialName) == "none") {
            intEta = Spectrum(0.0f);
            intK = Spectrum(1.0f);
        } else {
            intEta.fromContinuousSpectrum(InterpolatedSpectrum(
                fResolver->resolve("data/ior/" + materialName + ".eta.spd")));
            intK.fromContinuousSpectrum(InterpolatedSpectrum(
                fResolver->resolve("data/ior/" + materialName + ".k.spd")));
        }

        Float extEta = lookupIOR(props, "extEta", "air");

        m_eta = props.getSpectrum("eta", intEta) / extEta;
        m_k   = props.getSpectrum("k", intK) / extEta;
        m_anisoMode = props.getString("anisoMode", "");
        const auto sigma = props.getFloat("sigma", 0.1);
        m_sigma = new ConstantFloatTexture(sigma);
        m_sigmaV = new ConstantFloatTexture(props.getFloat("sigmaV", sigma));
    }

    RoughConductor(Stream *stream, InstanceManager *manager)
     : BSDF(stream, manager) {
        m_specularReflectance = static_cast<Texture *>(manager->getInstance(stream));
        m_eta = Spectrum(stream);
        m_k = Spectrum(stream);
        m_sigma = static_cast<Texture2D*>(manager->getInstance(stream));
        m_sigmaV = static_cast<Texture2D*>(manager->getInstance(stream));
        m_anisoMode = stream->readString();

        configure();
    }

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

        manager->serialize(stream, m_specularReflectance.get());
        m_eta.serialize(stream);
        m_k.serialize(stream);
        manager->serialize(stream, m_sigma.get());
        manager->serialize(stream, m_sigmaV.get());
        stream->writeString(m_anisoMode);
    }

    void configure() {
        unsigned int extraFlags = 0;
        if (!m_specularReflectance->isConstant())
            extraFlags |= ESpatiallyVarying;
        if (m_sigma != m_sigmaV)
            extraFlags |= EAnisotropic;

        m_components.clear();
        m_components.push_back(EGlossyReflection | EFrontSide | extraFlags);

        /* Verify the input parameters and fix them if necessary */
        m_specularReflectance = ensureEnergyConservation(
            m_specularReflectance, "specularReflectance", 1.0f);

        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;
    }
    
    virtual CoherenceProp lightMatterInteraction(const BSDFSamplingRecord &bRec) const override {
        // Precomputed HG coefficients
        static const Float hg_coeffs_smooth[4][4] = {
            { 0.896759,    0.0,            -0.0365681,     0.0,},
            { 0.0,          2.27416e-6,     0.0,            0.0,},
            { -0.0365699,   0.0,            0.0224114,      0.0,},
            { 0.0,          0.0,            0.0,            0.0,},
        };
        static const Float hg_coeffs_moderate[4][4] = {
            { 0.750355,     0.0,            -0.054313,      0.0,},
            { 0.0,          5.77583e-6,     0.0,            -0.000762236,},
            { -0.0543071,   0.0,            0.0100458,      0.0,},
            { 0.0,          0.000759205,    0.0,            0.0,},
        };
        static const Float hg_coeffs_rough[4][4] = {
            { 0.627688,    0.0,            0.0895008,     0.0,},
            { 0.0,          -0.0000184168, 0.0 ,          -0.248046,},
            { 0.0894786, 0.0, 0.535818,       0.0,},
            { 0.0,       0.248007, 0.0, -0.000100508, },
        };

        Assert(!!bRec.csd && "Not implemented");
        
        const auto &f = Frame(bRec.wo);
        const auto &Q = Matrix3x3(f.s,f.t,f.n);
        const auto r  = .1f;
        
        // Coefs
        static const Float sig_smooth = 0.001;
        static const Float sig_moderate = 0.01;
        static const Float sig_rough = 0.15;
        
        // Todo: correct aniso
        CoherenceProp csd;
        csd.Theta.setZero();
        csd.Theta(0,0) = 1. / (r*r);
        csd.Theta(1,1) = 1. / (r*r);
        csd.Theta(2,2) = 1. / 0.00001;
        csd.Q = Q;
        /* ForEachHGCoeff(n,m) */
        /*     csd.hg[n][m] = Spectrum{ 1.f }; */
        {
            const auto sig = m_sigma->eval(bRec.its).average();
            const auto &c1 = sig<sig_moderate?hg_coeffs_smooth:hg_coeffs_moderate;
            const auto &c2 = sig<sig_moderate?hg_coeffs_moderate:hg_coeffs_rough;
            const auto r = sig<sig_moderate ? std::min<Float>(std::max<Float>(0,(sig-sig_smooth)/(sig_moderate-sig_smooth)),1)
                                            : std::min<Float>(std::max<Float>(0,(sig-sig_moderate)/(sig_rough-sig_moderate)),1);
            ForEachHGCoeff(n,m)
                csd.hg[n][m] = Spectrum{ math::lerp(r, c1[n][m],c2[n][m]) };
        }
        
        ForEachHGCoeff(n,m)
            csd.hg[n][m] *= bRec.csd.eval(n,m,{0,0,0});

        return csd;
    }
    
    auto mfdist(const BSDFSamplingRecord &bRec) const {
        return MicrofacetDistribution(m_sigma != m_sigmaV ? MicrofacetDistribution::EGGX : MicrofacetDistribution::EBeckmann, 
                                      m_sigma->eval(bRec.its).average()*sqrt(2), 
                                      m_sigmaV->eval(bRec.its).average()*sqrt(2), 
                                      true);
    }
    
    auto anisoFrame(const BSDFSamplingRecord &bRec) const {
        if (m_sigma==m_sigmaV || m_anisoMode!="radial")
            return Frame(Vector3(1,0,0),Vector3(0,1,0),Normal(0,0,1));

        const auto n = Normal(0,0,1);
        auto uv = bRec.its.uv;
        const auto R = uv-Point2(.5f,.5f);
        const auto p = normalize(bRec.its.dpdu*R.x + bRec.its.dpdv*R.y);
        const auto lt = cross(n,BSDF::getFrame(bRec.its).toLocal(p));

        return Frame(lt,cross(lt,n),n);
    }
    
    Spectrum PSD(const BSDFSamplingRecord &bRec, const Vector3 &wo, const Vector3 &wi, const Vector3 &Phi) const {
        const auto distr = mfdist(bRec);
        auto H = normalize(Phi);
        H *= math::signum(Frame::cosTheta(H));

        return Spectrum(distr.G(wi, wo, H) * distr.eval(H));
    }
    
    Spectrum evalHG(const BSDFSamplingRecord &bRec) const override {
        /* Calculate the reflection half-vector */
        const auto aframe = anisoFrame(bRec);
        Vector wi  = aframe.toLocal(bRec.wi);
        Vector wo  = aframe.toLocal(bRec.wo);
        Vector Phi = wo+wi;
        Vector H = normalize(Phi);

        /* Fresnel factor */
        const Spectrum F = fresnelConductorExact(dot(wi, H), m_eta, m_k) * m_specularReflectance->eval(bRec.its);
        return F * PSD(bRec, wo, wi, Phi) / (4.0f * dot(wo, H));
    }

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

        return evalHG(bRec);
    }

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

        const auto aframe = anisoFrame(bRec);
        Vector wi = aframe.toLocal(bRec.wi);
        Vector wo = aframe.toLocal(bRec.wo);
        Vector H  = normalize(wo+wi);
        
        const auto dwh_dwo = 1.0f / (4.0f * dot(wo, H));
        H *= math::signum(Frame::cosTheta(H));
        const auto distr = mfdist(bRec);
        
        return distr.pdf(math::signum(Frame::cosTheta(wi)) * wi, H) * dwh_dwo;
    }

    Spectrum sample(BSDFSamplingRecord &bRec, const Point2 &sample) const {
        Float pdf;
        return this->sample(bRec,pdf,sample);
    }

    Spectrum sample(BSDFSamplingRecord &bRec, Float &pdf, const Point2 &sample) const {
        if (Frame::cosTheta(bRec.wi) < 0 ||
            (bRec.component != -1 && bRec.component != 0) ||
            !(bRec.typeMask & EGlossyReflection))
            return Spectrum(0.0f);

        const auto distr = mfdist(bRec);
        const auto aframe = anisoFrame(bRec);
        Vector wi = aframe.toLocal(bRec.wi);

        Float microfacetPDF;
        const Normal m = distr.sample(math::signum(Frame::cosTheta(wi)) * wi, sample, microfacetPDF);
        if (microfacetPDF == 0)
            return Spectrum(0.0f);
        
        Vector wo = reflect(wi, m);
        Vector H = normalize(wo+wi);
        const auto dwh_dwo = 1.0f / (4.0f * dot(wo, H));
        H *= math::signum(Frame::cosTheta(H));
        
        // bRec.wo = warp::squareToCosineHemisphere(sample);
        bRec.wo = aframe.toWorld(wo);
        bRec.eta = 1.0f;
        bRec.sampledComponent = 0;
        bRec.sampledType = EGlossyReflection;
        pdf = distr.pdf(math::signum(Frame::cosTheta(wi)) * wi, H) * dwh_dwo;
        if (!pdf)
            return Spectrum(.0f);
        return evalHG(bRec)/pdf;
    }
    
    Float getRoughness(const Intersection &its, int component) const {
        return (m_sigma->eval(its).average()+m_sigmaV->eval(its).average())/sqrt(2);
    }

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

    std::string toString() const {
        std::ostringstream oss;
        oss << "RoughConductor[" << endl
            << "  id = \"" << getID() << "\"," << endl
            << "  specularReflectance = " << indent(m_specularReflectance->toString()) << "," << endl
            << "  eta = " << m_eta.toString() << "," << endl
            << "  k = " << m_k.toString() << endl
            << "]";
        return oss.str();
    }

    MTS_DECLARE_CLASS()
private:
    ref<Texture> m_specularReflectance;
    ref<Texture2D> m_sigma, m_sigmaV;
    Spectrum m_eta, m_k;
    std::string m_anisoMode;
};

MTS_IMPLEMENT_CLASS_S(RoughConductor, false, BSDF)
MTS_EXPORT_PLUGIN(RoughConductor, "Rough conductor BRDF");
MTS_NAMESPACE_END
