/*
    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/scene.h>
#include "maxexp.h"

MTS_NAMESPACE_BEGIN

/*!\plugin{homogeneous}{Homogeneous participating medium}
 * \order{1}
 * \parameters{
 *     \parameter{material}{\String}{
 *         Name of a material preset, see
 *         \tblref{medium-coefficients}. \default{\texttt{skin1}}
 *     }
 *     \parameter{sigmaA, sigmaS}{\Spectrum}{
 *         Absorption and scattering
 *         coefficients of the medium in inverse scene units.
 *         These parameters are mutually exclusive with \code{sigmaT} and \code{albedo}
 *         \default{configured based on \code{material}}
 *     }
 *     \parameter{sigmaT, albedo}{\Spectrum}{
 *         Extinction coefficient in inverse scene units
 *         and a (unitless) single-scattering albedo.
 *         These parameters are mutually exclusive with \code{sigmaA} and \code{sigmaS}
 *         \default{configured based on \code{material}}
 *     }
 *     \parameter{\footnotesize{scale}}{\Float}{
 *         Optional scale factor that will be applied to the \code{sigma*} parameters.
 *         It is provided for convenience when accomodating data based on different units,
 *         or to simply tweak the density of the medium. \default{1}
 *     }
 *     \parameter{\Unnamed}{\Phase}{
 *          A nested phase function that describes the directional
 *          scattering properties of the medium. When none is specified,
 *          the renderer will automatically use an instance of
 *          \pluginref{isotropic}.
 *     }
 * }
 *
 * This class implements a flexible homogeneous participating
 * medium with support for arbitrary phase functions and various
 * medium sampling methods. It provides two different ways of configuring
 * the medium properties. One possibility is to load a material preset
 * using the \code{material} parameter---see \tblref{medium-coefficients}
 * for details. Alternatively, when specifying parameters by hand, they can either
 * be provided using the scattering and absorption coefficients, or
 * by declaring the extinction coefficient and single scattering
 * albedo (whichever is more convenient). Mixing these parameter
 * initialization methods is not allowed.
 *
 * All scattering parameters (named \code{sigma*}) should
 * be provided in inverse scene units. For instance, when a world-space
 * distance of 1 unit corresponds to a meter, the scattering coefficents should
 * have units of inverse meters. For convenience, the \code{scale}
 * parameter can be used to correct the units. For instance, when the scene is
 * in meters and the coefficients are in inverse millimeters, set
 * \code{scale} to \code{1000}.
 *
 * \renderings{
 *    \rendering{A squishy ball rendered with subsurface scattering and
 *    a dielectric BSDF (courtesy of Chanxi Zheng)}{medium_homogeneous_squishy.jpg}
 * }
 *
 * \begin{xml}[caption=Declaration of a forward scattering medium with high albedo]
 * <medium id="myMedium" type="homogeneous">
 *     <spectrum name="sigmaS" value="1"/>
 *     <spectrum name="sigmaA" value="0.05"/>
 *
 *     <phase type="hg">
 *         <float name="g" value="0.7"/>
 *     </phase>
 * </medium>
 * \end{xml}
 *
 * \textbf{Note}: Rendering media that have a spectrally
 * varying extinction coefficient can be tricky, since all
 * commonly used medium sampling methods suffer from high
 * variance in that case. Here, it may often make more sense to render
 * several monochromatic images separately (using only the coefficients for
 * a single channel) and then merge them back into a RGB image. There
 * is a \code{mtsutil} (\secref{mtsutil}) plugin named \code{joinrgb}
 * that will perform this RGB merging process.
 *
 * \begin{table}[h!]
 *     \centering
 *     \vspace{3mm}
 *     {\footnotesize
 *     \begin{tabular}{>{\ttfamily}p{3.8cm}p{.4cm}>{\ttfamily}p{3.8cm}p{.4cm}>{\ttfamily}p{3.8cm}}
 *         \toprule
 *         \rmfamily \small\textbf{Name} &&
 *         \rmfamily \small\textbf{Name} &&
 *         \rmfamily \small\textbf{Name} \\
 *         \cmidrule{1-1} \cmidrule{3-3} \cmidrule{5-5}
 *         Apple && Chicken1 && Chicken2 \\
 *         Cream && Ketchup  && Potato \\
 *         Skimmilk && Skin1 && Skin2 \\
 *         Spectralon && Wholemilk && \\
 *         \cmidrule{1-1} \cmidrule{3-3} \cmidrule{5-5}
 *         Lowfat Milk              &&  Gatorade                &&    White Grapefruit Juice     \\
 *         Reduced Milk             &&  Chardonnay              &&    Shampoo                    \\
 *         Regular Milk             &&  White Zinfandel         &&    Strawberry Shampoo         \\
 *         Espresso                 &&  Merlot                  &&    \mbox{Head \& Shoulders Shampoo}  \\
 *         Mint Mocha Coffee        &&  Budweiser Beer          &&    Lemon Tea Powder           \\
 *         Lowfat Soy Milk          &&  Coors Light Beer        &&    Orange Juice Powder        \\
 *         Regular Soy Milk         &&  Clorox                  &&    Pink Lemonade Powder       \\
 *         Lowfat Chocolate Milk    &&  Apple Juice             &&    Cappuccino Powder          \\
 *         Regular Chocolate Milk   &&  Cranberry Juice         &&    Salt Powder                \\
 *         Coke                     &&  Grape Juice             &&    Sugar Powder               \\
 *         Pepsi                    &&  Ruby Grapefruit Juice   &&    Suisse Mocha               \\
 *         Sprite                   &&                          &&                               \\
 *         \bottomrule
 *     \end{tabular}}
 *     \caption{\label{tbl:medium-coefficients}This
 *          table lists all supported medium material presets. The
 *          top entries are from Jensen et al. \cite{Jensen2001Practical}, and the
 *          bottom ones are from Narasimhan et al. \cite{Narasimhan2006Acquiring}.
 *          They all use units of $\frac{1}{mm}$, so remember to set
 *          \code{scale} appropriately when your scene is not
 *          in units of millimeters.
 *          These material presets can be used with the plugins
 *          \pluginref{homogeneous},\
 *          \pluginref{dipole}, and \
 *          \pluginref{hk}
 *     }
 * \end{table}
 */
class HomogeneousMedium : public Medium {
public:
    /**
     * This class supports the following sampling strategies for choosing
     * a suitable scattering location when sampling the RTE
     */
    enum ESamplingStrategy {
        EBalance,  /// Exponential distrib.; pick a random channel each time
        ESingle,   /// Exponential distrib.; pick a specified channel
        EManual,   /// Exponential distrib.; manually specify the falloff
        EMaximum   /// Maximum-of-exponential distribution
    };

    HomogeneousMedium(const Properties &props)
            : Medium(props), m_samplingDensity(0.0f), m_maxExpDist(NULL) {
        std::string strategy = props.getString("strategy", "balance");

        /**
         * The goal of the medium sampling weight is to be able to
         * sample medium intarctions according to
         *    sigma_s(t) * tau(0 <-> t)
         * as opposed to
         *    sigma_t(t) * tau(0 <-> t)
         * See the separate writeup for more details.
         */
        m_mediumSamplingWeight = props.getFloat("mediumSamplingWeight", -1);
        if (m_mediumSamplingWeight == -1) {
            for (int i=0; i<SPECTRUM_SAMPLES; ++i) {
                /// Record the highest albedo values across channels
                Float albedo = m_sigmaS[i] / m_sigmaT[i];
                if (albedo > m_mediumSamplingWeight && m_sigmaT[i] != 0)
                    m_mediumSamplingWeight = albedo;
            }
            if (m_mediumSamplingWeight > 0) {
                /* The medium scatters some light -> place at least half
                   of the samples in it, otherwise we will render lots
                   of spatially varying noise where one pixel has a
                   medium interaction and the neighbors don't */
                m_mediumSamplingWeight = std::max(m_mediumSamplingWeight,
                    (Float) 0.5f);
            }
        }

        if (strategy == "balance") {
            m_strategy = EBalance;
        } else if (strategy == "single") {
            m_strategy = ESingle;

            /* By default, choose the lowest-variance channel
               (the one with the smallest sigma_t, that is) */
            int channel = 0;
            Float smallest = std::numeric_limits<Float>::infinity();
            for (int i=0; i<SPECTRUM_SAMPLES; ++i) {
                if (m_sigmaT[i] < smallest) {
                    smallest = m_sigmaT[i];
                    channel = i;
                }
            }

            channel = props.getInteger("channel", channel);
            Assert(channel >= 0 && channel < SPECTRUM_SAMPLES);
            m_samplingDensity = m_sigmaT[channel];

            if (props.getBoolean("monochromatic", false)) {
                /* Optionally turn this into a monochromatic medium
                   based on the chosen color channel. This is useful
                   when the whole scene is rendered once per channel
                   and then recombined to create a color image */
                m_sigmaA = Spectrum(m_sigmaA[channel]);
                m_sigmaS = Spectrum(m_sigmaS[channel]);
                m_sigmaT = m_sigmaA + m_sigmaS;
            }
        } else if (strategy == "maximum") {
            m_strategy = EMaximum;
            std::vector<Float> coeffs(SPECTRUM_SAMPLES);
            for (int i=0; i<SPECTRUM_SAMPLES; ++i)
                coeffs[i] = m_sigmaT[i];
            m_maxExpDist = new MaxExpDist(coeffs);
        } else if (strategy == "manual") {
            m_strategy = EManual;
            m_samplingDensity = props.getFloat("samplingDensity");
        } else {
            Log(EError, "Specified an unknown sampling strategy");
        }
    }

    HomogeneousMedium(Stream *stream, InstanceManager *manager)
        : Medium(stream, manager), m_maxExpDist(NULL) {
        m_strategy = (ESamplingStrategy) stream->readInt();
        m_samplingDensity = stream->readFloat();
        m_mediumSamplingWeight = stream->readFloat();

        if (m_strategy == EMaximum) {
            std::vector<Float> coeffs(SPECTRUM_SAMPLES);
            for (int i=0; i<SPECTRUM_SAMPLES; ++i)
                coeffs[i] = m_sigmaT[i];
            m_maxExpDist = new MaxExpDist(coeffs);
        }

        configure();
    }

    virtual ~HomogeneousMedium() {
        if (m_maxExpDist)
            delete m_maxExpDist;
    }

    void configure() {
        Medium::configure();
        m_albedo = 0;
        for (int i=0; i<SPECTRUM_SAMPLES; ++i) {
            if (m_sigmaT[i] != 0)
                m_albedo = std::max(m_albedo, m_sigmaS[i]/m_sigmaT[i]);
        }
    }

    void serialize(Stream *stream, InstanceManager *manager) const {
        Medium::serialize(stream, manager);
        stream->writeInt(m_strategy);
        stream->writeFloat(m_samplingDensity);
        stream->writeFloat(m_mediumSamplingWeight);
    }

    Spectrum evalTransmittance(const Ray &ray, Sampler *) const {
        Float negLength = ray.mint - ray.maxt;
        Spectrum transmittance;
        for (int i=0; i<SPECTRUM_SAMPLES; ++i)
            transmittance[i] = m_sigmaT[i] != 0
                ? math::fastexp(m_sigmaT[i] * negLength) : (Float) 1.0f;
        return transmittance;
    }

    bool sampleDistance(const Ray &ray, MediumSamplingRecord &mRec,
            Sampler *sampler) const {
        Float rand = sampler->next1D(), sampledDistance;
        Float samplingDensity = m_samplingDensity;

        if (rand < m_mediumSamplingWeight) {
            rand /= m_mediumSamplingWeight;
            if (m_strategy != EMaximum) {
                /* Choose the sampling density to be used */
                if (m_strategy == EBalance) {
                    int channel = std::min((int) (sampler->next1D()
                        * SPECTRUM_SAMPLES), SPECTRUM_SAMPLES-1);
                    samplingDensity = m_sigmaT[channel];
                }
                sampledDistance = -math::fastlog(1-rand) / samplingDensity;
            } else {
                sampledDistance = m_maxExpDist->sample(1-rand, mRec.pdfSuccess);
            }
        } else {
            /* Don't generate a medium interaction */
            sampledDistance = std::numeric_limits<Float>::infinity();
        }
        Float distSurf = ray.maxt - ray.mint;
        bool success = true;

        if (sampledDistance < distSurf) {
            mRec.t = sampledDistance + ray.mint;
            mRec.p = ray(mRec.t);
            mRec.sigmaA = m_sigmaA;
            mRec.sigmaS = m_sigmaS;
            mRec.time = ray.time;
            mRec.medium = this;

            /* Fail if there is no forward progress
               (e.g. due to roundoff errors) */
            if (mRec.p == ray.o)
                success = false;
        } else {
            sampledDistance = distSurf;
            success = false;
        }

        switch (m_strategy) {
            case EMaximum:
                mRec.pdfFailure = 1-m_maxExpDist->cdf(sampledDistance);
                break;

            case EBalance:
                mRec.pdfFailure = 0;
                mRec.pdfSuccess = 0;
                for (int i=0; i<SPECTRUM_SAMPLES; ++i) {
                    Float tmp = math::fastexp(-m_sigmaT[i] * sampledDistance);
                    mRec.pdfFailure += tmp;
                    mRec.pdfSuccess += m_sigmaT[i] * tmp;
                }
                mRec.pdfFailure /= SPECTRUM_SAMPLES;
                mRec.pdfSuccess /= SPECTRUM_SAMPLES;
                break;

            case ESingle:
            case EManual:
                mRec.pdfFailure = math::fastexp(-samplingDensity * sampledDistance);
                mRec.pdfSuccess = samplingDensity * mRec.pdfFailure;
                break;

            default:
                Log(EError, "Unknown sampling strategy!");
        }

        mRec.transmittance = (m_sigmaT * (-sampledDistance)).exp();
        mRec.pdfSuccessRev = mRec.pdfSuccess = mRec.pdfSuccess * m_mediumSamplingWeight;
        mRec.pdfFailure = m_mediumSamplingWeight * mRec.pdfFailure + (1-m_mediumSamplingWeight);
        mRec.medium = this;
        if (mRec.transmittance.max() < 1e-20)
            mRec.transmittance = Spectrum(0.0f);

        return success;
    }

    void eval(const Ray &ray, MediumSamplingRecord &mRec) const {
        Float distance = ray.maxt - ray.mint;
        switch (m_strategy) {
            case EManual:
            case ESingle: {
                    Float temp = math::fastexp(-m_samplingDensity * distance);
                    mRec.pdfSuccess = m_samplingDensity * temp;
                    mRec.pdfFailure = temp;
                }
                break;

            case EBalance: {
                    mRec.pdfSuccess = 0;
                    mRec.pdfFailure = 0;
                    for (int i=0; i<SPECTRUM_SAMPLES; ++i) {
                        Float temp = math::fastexp(-m_sigmaT[i] * distance);
                        mRec.pdfSuccess += m_sigmaT[i] * temp;
                        mRec.pdfFailure += temp;
                    }
                    mRec.pdfSuccess /= SPECTRUM_SAMPLES;
                    mRec.pdfFailure /= SPECTRUM_SAMPLES;
                }
                break;

            case EMaximum:
                mRec.pdfSuccess = m_maxExpDist->pdf(distance);
                mRec.pdfFailure = 1-m_maxExpDist->cdf(distance);
                break;

            default:
                Log(EError, "Unknown sampling strategy!");
        }

        mRec.transmittance = (m_sigmaT * (-distance)).exp();
        mRec.pdfSuccess = mRec.pdfSuccessRev = mRec.pdfSuccess * m_mediumSamplingWeight;
        mRec.pdfFailure = mRec.pdfFailure * m_mediumSamplingWeight + (1-m_mediumSamplingWeight);
        mRec.sigmaA = m_sigmaA;
        mRec.sigmaS = m_sigmaS;
        mRec.time = ray.time;
        mRec.medium = this;
        if (mRec.transmittance.max() < 1e-20)
            mRec.transmittance = Spectrum(0.0f);
    }

    bool isHomogeneous() const {
        return true;
    }

    std::string toString() const {
        std::ostringstream oss;
        oss << "HomogeneousMedium[" << endl
            << "  sigmaA = " << m_sigmaA.toString() << "," << endl
            << "  sigmaS = " << m_sigmaS.toString() << "," << endl
            << "  sigmaT = " << m_sigmaT.toString() << "," << endl
            << "  mediumSamplingWeight = " << m_mediumSamplingWeight << "," << endl
            << "  samplingDensity = " << m_samplingDensity << "," << endl
            << "  strategy = ";

        switch (m_strategy) {
            case ESingle: oss << "single," << endl; break;
            case EManual: oss << "manual," << endl; break;
            case EBalance: oss << "balance," << endl; break;
            case EMaximum: oss << "maximum," << endl; break;
        }

        oss << "  phase = " << indent(m_phaseFunction.toString()) << endl
            << "]";
        return oss.str();
    }

    MTS_DECLARE_CLASS()
private:
    Float m_samplingDensity, m_mediumSamplingWeight;
    ESamplingStrategy m_strategy;
    MaxExpDist *m_maxExpDist;
    Float m_albedo;
};

MTS_IMPLEMENT_CLASS_S(HomogeneousMedium, false, Medium)
MTS_EXPORT_PLUGIN(HomogeneousMedium, "Homogeneous medium");
MTS_NAMESPACE_END
