/*
    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/render/texture.h>
#include <mitsuba/hw/basicshader.h>
#include <mitsuba/core/warp.h>
#include <mitsuba/core/fresolver.h>
#include <boost/algorithm/string.hpp>
#include "ior.h"

MTS_NAMESPACE_BEGIN

class Grating : public BSDF {
public:
    Grating(const Properties &props)
        : BSDF(props) {
        std::vector<std::string> lobes =
            tokenize(props.getString("lobes", ""), " ,;");
        if (lobes.size() == 0)
            Log(EError, "No lobes were supplied!");
        m_lobes.resize(lobes.size());
        char *end_ptr = NULL;
        for (size_t i=0; i<lobes.size(); ++i) {
            Float lobe = (Float) strtod(lobes[i].c_str(), &end_ptr);
            if (*end_ptr != '\0')
                SLog(EError, "Could not parse the diffraction lobes!");
            if (lobe < 0)
                SLog(EError, "Invalid diffraction lobe!");
            m_lobes[i] = lobe;
        }

        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 {
            ref<FileResolver> fResolver = Thread::getThread()->getFileResolver();
            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_pitch = props.getFloat("pitch");

        m_ignoreCoherence = props.getBoolean("ignoreCoherence", false);
        if (m_ignoreCoherence)
            Log(EWarn, "Material ignores coherence!");
        
        m_specularSamplingWeight = .1f;
    }

    Grating(Stream *stream, InstanceManager *manager)
        : BSDF(stream, manager) {
        m_specularReflectance = static_cast<Texture *>(manager->getInstance(stream));
        m_eta = Spectrum(stream);
        m_k = Spectrum(stream);
        
        m_lobes.resize(stream->readSize());
        for (size_t i=0; i<m_lobes.size(); ++i)
            m_lobes[i] = stream->readFloat();
        m_pitch = stream->readFloat();
            
        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);

        stream->writeSize(m_lobes.size());
        for (size_t i=0; i<m_lobes.size(); ++i)
            stream->writeFloat(m_lobes[i]);
        stream->writeFloat(m_pitch);
    }

    void configure() {
        m_components.clear();
        m_components.push_back(EDeltaReflection | EFrontSide | ESpatiallyVarying);
        m_components.push_back(EDiffuseReflection | EFrontSide | ESpatiallyVarying);
        
        if (m_lobes.size() == 0)
            Log(EError, "No lobes were supplied!");

        BSDF::configure();
    }

    inline Vector reflect(const Vector &wi) const {
        return Vector(-wi.x, -wi.y, wi.z);
    }
    
    virtual CoherenceProp lightMatterInteraction(const BSDFSamplingRecord &bRec) const override {
        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;
        
        auto csd = make_csd_spherical_source(Q,r,sqrt(bRec.csd.Theta(2,2)));
        ForEachHGCoeff(n,m)
            csd.hg[n][m] *= bRec.csd.eval(n,m,{0,0,0});

        return csd;
    }

    Spectrum getDiffuseReflectance(const Intersection &its) const {
        return Spectrum(1.f);
    }
    
    Spectrum eval_hg_dual(const BSDFSamplingRecord &bRec, Vector3 xi) const {
        if (m_ignoreCoherence)
            return Spectrum{ 1. };
        xi = bRec.csd.Theta * (bRec.csd.Q*xi * bRec.csd.s.length());

        Spectrum v{ .0f };
        const auto norm = bRec.csd.eval({0,0,0});
        if (norm.min()==.0f)
            return v;

        ForEachHGCoeff(n,m)
        for (int i=0; i<SPECTRUM_SAMPLES; i++)
            v[i] += bRec.csd.hg[n][m][i] * HG(n,m,xi / (CoherenceProp::k(i)), bRec.csd.Theta);
        return v / norm;
    }
    Spectrum PSD(const BSDFSamplingRecord &bRec, const Vector3 &local_tangent) const {
        if (!bRec.csd)
            return Spectrum(.0f);
        
        auto wi = bRec.wi;
        if (bRec.directFromSrc && !m_ignoreCoherence) 
            wi = bRec.wisrc;

        const auto dirac = 1.f/.002f;
        const auto pitch = m_pitch * 1e-5f;
        
        const auto Phi = -wi+bRec.wo;
        const auto& PhiPerp = Vector3{ Phi.x * local_tangent.x, Phi.y * local_tangent.y, .0f };
        const auto& ppworld = BSDF::getFrame(bRec.its).toWorld(PhiPerp);
        auto hgs = eval_hg_dual(bRec, ppworld);
        hgs.clampNegative();
 
        Spectrum ret = Spectrum(.0f);
        for (size_t l=0; l<m_lobes.size(); ++l)
        for (int i=0; i<SPECTRUM_SAMPLES; ++i) {
            const auto phi = CoherenceProp::k(i)*Phi;
            auto p1 = fmod(pitch*dot(phi,local_tangent) - 2*M_PI*(l+1), 2*M_PI);
            auto p2 = fmod(pitch*dot(phi,local_tangent) + 2*M_PI*(l+1), 2*M_PI);
            if (p1<-M_PI) p1+=2*M_PI;
            if (p2<-M_PI) p2+=2*M_PI;
            ret[i] += m_lobes[l] * (exp(-(p1*p1)*dirac) + exp(-(p2*p2)*dirac));
        }

        return ret * hgs;
    }
    
    Spectrum evalLobes(const BSDFSamplingRecord &bRec) const {
        const Intersection &its = bRec.its;
        const auto &frame = BSDF::getFrame(its);
        const auto uv = its.uv;
        const auto R  = uv-Point2(.5f,.5f);
        const auto p  = normalize(its.dpdu*R.x + its.dpdv*R.y);
        const auto lt = cross(Vector3(0,0,1),frame.toLocal(p));

        return PSD(bRec, lt);
    }
    
    Spectrum evalHG(const BSDFSamplingRecord &bRec) const override {
        bool hasSpecular   = (bRec.typeMask & EDeltaReflection)
                && (bRec.component == -1 || bRec.component == 0);
        bool hasDiffuse = (bRec.typeMask & EDiffuseReflection)
                && (bRec.component == -1 || bRec.component == 1);
        if (Frame::cosTheta(bRec.wo) <= 0 || Frame::cosTheta(bRec.wi) <= 0)
            return Spectrum(0.0f);
        
        Spectrum ret(.0f);
        if (hasSpecular) {
            if (std::abs(dot(reflect(bRec.wi), bRec.wo)-1) < DeltaEpsilon)
                ret = m_specularReflectance->eval(bRec.its) *
                    fresnelConductorExact(Frame::cosTheta(bRec.wi), m_eta, m_k);
        }
        if (hasDiffuse) {
            ret += evalLobes(bRec)/.5f;
        }
        return ret;
    }

    Spectrum eval(const BSDFSamplingRecord &bRec, EMeasure measure) const {
        bool hasSpecular   = (bRec.typeMask & EDeltaReflection)
                && (bRec.component == -1 || bRec.component == 0)
                && measure == EDiscrete;
        bool hasDiffuse = (bRec.typeMask & EDiffuseReflection)
                && (bRec.component == -1 || bRec.component == 1)
                && measure == ESolidAngle;

        if (hasSpecular || hasDiffuse)
            return evalHG(bRec) * pdf(bRec,ESolidAngle);
        return Spectrum(.0f);
    }

    Float pdf(const BSDFSamplingRecord &bRec, EMeasure measure) const {
        bool hasSpecular   = (bRec.typeMask & EDeltaReflection)
                && (bRec.component == -1 || bRec.component == 0);
        bool hasDiffuse = (bRec.typeMask & EDiffuseReflection)
                && (bRec.component == -1 || bRec.component == 1);

        if (Frame::cosTheta(bRec.wo) <= 0 || Frame::cosTheta(bRec.wi) <= 0)
            return 0.0f;

        Float probSpecular = hasSpecular ? 1.0f : 0.0f;
        if (hasSpecular && hasDiffuse) {
            const auto F = fresnelConductorExact(Frame::cosTheta(bRec.wi), m_eta, m_k);
            const auto Fi = F.average();
            probSpecular = (Fi*m_specularSamplingWeight) /
                (Fi*m_specularSamplingWeight +
                (1-Fi) * (1-m_specularSamplingWeight));
        }

        if (hasSpecular && measure == EDiscrete) {
            /* Check if the provided direction pair matches an ideal
               specular reflection; tolerate some roundoff errors */
            if (std::abs(dot(reflect(bRec.wi), bRec.wo)-1) < DeltaEpsilon)
                return probSpecular;
        } else if (hasDiffuse && measure == ESolidAngle) {
            return warp::squareToCosineHemispherePdf(bRec.wo) * (1-probSpecular);
        }

        return 0.0f;
    }

    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 {
        bool hasSpecular = (bRec.typeMask & EDeltaReflection)
                && (bRec.component == -1 || bRec.component == 0);
        bool hasDiffuse = (bRec.typeMask & EDiffuseReflection)
                && (bRec.component == -1 || bRec.component == 1);

        if ((!hasDiffuse && !hasSpecular) || Frame::cosTheta(bRec.wi) <= 0)
            return Spectrum(0.0f);

        const auto F = fresnelConductorExact(Frame::cosTheta(bRec.wi), m_eta, m_k);
        const auto Fi = F.average();

        bRec.eta = 1.0f;
        if (hasDiffuse && hasSpecular) {
            const Float probSpecular = (Fi*m_specularSamplingWeight) /
                (Fi*m_specularSamplingWeight +
                (1-Fi) * (1-m_specularSamplingWeight));

            /* Importance sample wrt. the Fresnel reflectance */
            if (sample.x < probSpecular) {
                bRec.sampledComponent = 0;
                bRec.sampledType = EDeltaReflection;
                bRec.wo = reflect(bRec.wi);

                pdf = probSpecular;
                return m_specularReflectance->eval(bRec.its)
                    * F / probSpecular;
            } else {
                bRec.sampledComponent = 1;
                bRec.sampledType = EDiffuseReflection;
                bRec.wo = warp::squareToCosineHemisphere(Point2(
                    (sample.x - probSpecular) / (1 - probSpecular),
                    sample.y
                ));
                pdf = (1-probSpecular) *
                    warp::squareToCosineHemispherePdf(bRec.wo);

                return evalLobes(bRec)/.5f;
            }
        } else if (hasSpecular) {
            bRec.sampledComponent = 0;
            bRec.sampledType = EDeltaReflection;
            bRec.wo = reflect(bRec.wi);
            pdf = 1;
            return m_specularReflectance->eval(bRec.its) * F;
        } else {
            bRec.sampledComponent = 1;
            bRec.sampledType = EDiffuseReflection;
            bRec.wo = warp::squareToCosineHemisphere(sample);
            pdf = warp::squareToCosineHemispherePdf(bRec.wo);

            return evalLobes(bRec)/.5f;
        }
    }

    void addChild(const std::string &name, ConfigurableObject *child) {
        BSDF::addChild(name, child);
    }

    Float getRoughness(const Intersection &its, int component) const {
        return std::numeric_limits<Float>::infinity();
    }

    std::string toString() const {
        std::ostringstream oss;
        oss << "Grating[" << endl
            << "  id = \"" << getID() << "\"," << endl
            << "]";
        return oss.str();
    }

    MTS_DECLARE_CLASS()
private:
    std::vector<Float> m_lobes;
    Float m_pitch,m_specularSamplingWeight;
    Spectrum m_eta;
    Spectrum m_k;
    ref<Texture> m_specularReflectance;
    bool m_ignoreCoherence{ false };
};

MTS_IMPLEMENT_CLASS_S(Grating, false, BSDF)
MTS_EXPORT_PLUGIN(Grating, "Diffraction grating BRDF")
MTS_NAMESPACE_END
