#include "roughconductor.h"
#include <render/intersection.h>
#include "microfacet.h"
#include <core/utils.h>
#include <assert.h>

BSDF *RoughConductorBSDF::clone() const
{
    return new RoughConductorBSDF(*this);
}

static Spectrum fresnelConductorExact(Float cosThetaI, const Spectrum &eta,
                                      const Spectrum &k)
{
    /* Modified from "Optics" by K.D. Moeller, University Science Books, 1988 */

    Float cosThetaI2 = cosThetaI * cosThetaI, sinThetaI2 = 1.0f - cosThetaI2,
          sinThetaI4 = sinThetaI2 * sinThetaI2;

    Spectrum eta2 = eta.square(), k2 = k.square();

    Spectrum temp1 = eta2 - k2 - Spectrum(sinThetaI2),
             a2pb2 = (temp1.square() + k2.cwiseProduct(eta2) * 4.0f).sqrt(),
             a = ((a2pb2 + temp1) * 0.5f).sqrt();

    Spectrum term1 = a2pb2 + Spectrum(cosThetaI2), term2 = a * (2.0f * cosThetaI);

    Spectrum Rs2 = (term1 - term2) / (term1 + term2);

    Spectrum term3 = a2pb2 * cosThetaI2 + Spectrum(sinThetaI4),
             term4 = term2 * sinThetaI2;

    Spectrum Rp2 = Rs2 * (term3 - term4) / (term3 + term4);

    return 0.5f * (Rp2 + Rs2);
}

Spectrum RoughConductorBSDF::eval(const Intersection &its, const Vector &wo,
                                  EBSDFMode mode) const
{
    /* Calculate the reflection half-vector */
    Vector H = (wo + its.wi).normalized();

    /* Evaluate the microfacet normal distribution */
    const Float D = m_distr.eval(H);
    if (std::abs(D) < Epsilon)
        return Spectrum(0.0f);

    /* Fresnel factor */
    const Spectrum F = fresnelConductorExact(its.wi.dot(H), m_eta, m_k);

    /* Smith's shadow-masking function */
    const Float G = m_distr.G(its.wi, wo, H);

    /* Calculate the total amount of reflection */
    Float model = D * G / (4.0f * Frame::cosTheta(its.wi));

    return F * model;
}

Spectrum RoughConductorBSDF::sample(const Intersection &its, const Array3 &rnd,
                                    Vector &wo, Float &pdf, Float &eta,
                                    EBSDFMode mode) const
{
    if (Frame::cosTheta(its.wi) < Epsilon)
        return Spectrum(0.0f);

    /* Sample M, the microfacet normal */
    Vector m = m_distr.sample(its.wi, Array2(rnd[0], rnd[1]), pdf);

    if (pdf < Epsilon)
        return Spectrum(0.0f);

    /* Perfect specular reflection based on the microfacet normal */
    wo = reflect(its.wi, m);

    /* Side check */
    if (Frame::cosTheta(wo) < Epsilon)
        return Spectrum(0.0f);

    Spectrum F = fresnelConductorExact(its.wi.dot(m), m_eta, m_k);
    Float weight = m_distr.smithG1(wo, m);

    /* Jacobian of the half-direction mapping */
    pdf /= 4.0f * wo.dot(m);

    return F * weight;
}

void RoughConductorBSDF::sampleDual(const Intersection &its, const Array3 &rnd,
                    Vector &wo, Vector &woDual, Float &pdf, Float &pdfDual, Float &eta, Float &etaDual, Spectrum &ret, Spectrum &retDual,
                    EBSDFMode mode) const
{
    if (Frame::cosTheta(its.wi) < Epsilon){
        ret = retDual = Spectrum(0.0f);
        return;
    }

    /* Sample M, the microfacet normal */
    Vector m = m_distr.sample(its.wi, Array2(rnd[0], rnd[1]), pdf);
    pdfDual = pdf;
    Vector m_flip = Vector(-m.x(), -m.y(), m.z());
    bool ret_valid = true, retDual_valid = true;
    if (pdf < Epsilon){
        ret = Spectrum(0.0f);
        ret_valid = false;
    }
    if (pdfDual < Epsilon){
        retDual = Spectrum(0.0f);
        retDual_valid = false;
    }

    /* Perfect specular reflection based on the microfacet normal */
    wo = reflect(its.wi, m);
    woDual = reflect(its.wi, m_flip);

    /* Side check */
    if (Frame::cosTheta(wo) < Epsilon){
        ret = Spectrum(0.0f);
        ret_valid = false;
    }
    if (Frame::cosTheta(woDual) < Epsilon){
        retDual = Spectrum(0.0f);
        retDual_valid = false;
    }

    Spectrum F = fresnelConductorExact(its.wi.dot(m), m_eta, m_k);
    Spectrum FDual = fresnelConductorExact(its.wi.dot(m_flip), m_eta, m_k);
    Float weight = m_distr.smithG1(wo, m);
    Float weightDual = m_distr.smithG1(woDual, m_flip);

    /* Jacobian of the half-direction mapping */
    if (ret_valid){
        pdf /= 4.0f * wo.dot(m);
        ret = F * weight;
    }

    if (retDual_valid){
        pdfDual /= 4.0f * woDual.dot(m_flip);
        retDual = FDual * weightDual;
    }
    return;
}

Float RoughConductorBSDF::pdf(const Intersection &its, const Vector &wo) const
{
    /* Calculate the reflection half-vector */
    Vector H = (wo + its.wi).normalized();
    return m_distr.eval(H) * m_distr.smithG1(its.wi, H) /
           (4.0f * Frame::cosTheta(its.wi));
}

IMPLEMENT_BSDF_HELPER_FUNCTIONS(RoughConductorBSDF)