#include "roughconductor.h"
#include "intersection.h"
#include "intersectionAD.h"
#include "utils.h"
#include "microfacet.h"
#include <assert.h>

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);
}


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

    FloatAD cosThetaI2 = cosThetaI*cosThetaI,
            sinThetaI2 = 1.0f - cosThetaI2,
            sinThetaI4 = sinThetaI2.square();

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

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

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

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

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

    SpectrumAD 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.val, m_k.val);

    /* 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;
}


SpectrumAD RoughConductorBSDF::evalAD(const IntersectionAD &its, const VectorAD &wo, EBSDFMode mode) const {
     /* Calculate the reflection half-vector */
    VectorAD H = (wo + its.wi).normalized();

    /* Evaluate the microfacet normal distribution */
    const FloatAD D = m_distr.evalAD(H);
    if ( std::abs(D.val) < Epsilon ) return SpectrumAD();

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

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

    /* Calculate the total amount of reflection */
    FloatAD model = D * G / (4.0f * FrameAD::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.val, m_k.val);
    Float weight = m_distr.smithG1(wo, m);

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

    return F*weight;
}


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));
}
