#include <psdr/core/warp.h>
#include <psdr/bsdf/roughdielectric.h>
#include <psdr/bsdf/ggx.h>

namespace psdr
{

SpectrumC RoughDielectric::eval(const IntersectionC& its, const Vector3fC& wo, MaskC active) const {
    return __eval<false>(its, wo, active);
}


SpectrumD RoughDielectric::eval(const IntersectionD& its, const Vector3fD& wo, MaskD active) const {
    return __eval<true>(its, wo, active);
}


BSDFSampleC RoughDielectric::sample(const IntersectionC& its, const Vector3fC& sample, MaskC active) const {
    return __sample<false>(its, sample, active);
}


BSDFSampleD RoughDielectric::sample(const IntersectionD& its, const Vector3fD& sample, MaskD active) const {
    return __sample<true>(its, sample, active);
}

BSDFSampleDualC RoughDielectric::sampleDual(const IntersectionC &its, const Vector3fC &sample, MaskC active) const {
    PSDR_ASSERT(0);
    BSDFSampleDualC bs;
    return bs;
}


BSDFSampleDualD RoughDielectric::sampleDual(const IntersectionD &its, const Vector3fD &sample, MaskD active) const {
    PSDR_ASSERT(0);
    BSDFSampleDualD bs;
    return bs;
}

FloatC RoughDielectric::pdf(const IntersectionC& its, const Vector3fC& wo, MaskC active) const {
    return __pdf<false>(its, wo, active);
}


FloatD RoughDielectric::pdf(const IntersectionD& its, const Vector3fD& wo, MaskD active) const {
    return __pdf<true>(its, wo, active);
}


template <bool ad>
Spectrum<ad> RoughDielectric::__eval(const Intersection<ad>& its, const Vector3f<ad>& wo, Mask<ad> active) const {

    Float<ad> cos_theta_i = Frame<ad>::cos_theta(its.wi),
              cos_theta_o = Frame<ad>::cos_theta(wo);
    active &= neq(cos_theta_i, 0.f);
    Mask<ad> reflect = cos_theta_i * cos_theta_o > 0.f;

    Float<ad>    eta, inv_eta;

    if constexpr (ad) {
        eta     = select(cos_theta_i > 0.f, m_eta, m_inv_eta);
        inv_eta = select(cos_theta_i > 0.f, m_inv_eta, m_eta);
    } else {
        eta     = select(cos_theta_i > 0.f, detach(m_eta), detach(m_inv_eta));
        inv_eta = select(cos_theta_i > 0.f, detach(m_inv_eta), detach(m_eta));
    }

    Vector3f<ad> m = normalize(its.wi + wo * select(reflect, 1.f, eta));
                 m = mulsign(m, Frame<ad>::cos_theta(m));

    Float<ad> alpha_u = m_alpha_u.eval<ad>(its.uv);
    Float<ad> alpha_v = m_alpha_v.eval<ad>(its.uv);
    GGXDistribution distr(alpha_u, alpha_v);
    Float<ad> D = distr.eval<ad>(m);

    Float<ad> F;
    if constexpr (ad) {
        F = std::get<0>(fresnel_dielectric<ad>(m_eta, dot(its.wi, m)));
    } else {
        F = std::get<0>(fresnel_dielectric<ad>(detach(m_eta), dot(its.wi, m)));
    }

    Float<ad> G = distr.G<ad>(its.wi, wo, m);
    Spectrum<ad> result = 0.0f;

    Mask<ad> eval_r = reflect && active,
             eval_t = !reflect && active;

    if (any_or<true>(eval_r)) {
        Spectrum<ad> value = F * D * G / (4.f * abs(cos_theta_i));
        result[eval_r] = value;
    }

    if (any_or<true>(eval_t)) {
        /* Missing term in the original paper: account for the solid angle
           compression when tracing radiance -- this is necessary for
           bidirectional methods. */
        Float<ad> scale = sqr(inv_eta);
        // Compute the total amount of transmission
        Spectrum<ad> value = abs(
            (scale * (1.f - F) * D * G * eta * eta * dot(its.wi, m) * dot(wo, m)) /
            (cos_theta_i * sqr(dot(its.wi, m) + eta * dot(wo, m))));

        result[eval_t] = value;
    }
    return result && (eval_r || eval_t);
}


template <bool ad>
Float<ad> RoughDielectric::__pdf(const Intersection<ad>& its, const Vector3f<ad>& wo, Mask<ad> active) const {

    Float<ad> cos_theta_i = Frame<ad>::cos_theta(its.wi),
              cos_theta_o = Frame<ad>::cos_theta(wo);

    active &= neq(cos_theta_i, 0.f);
    Mask<ad> reflect = cos_theta_i * cos_theta_o > 0.f;

    Float<ad> eta;
    if constexpr (ad) {
        eta = select(cos_theta_i > 0.f, m_eta, m_inv_eta);
    } else {
        eta = select(cos_theta_i > 0.f, detach(m_eta), detach(m_inv_eta));
    }

    Vector3f<ad> m = normalize(its.wi + wo * select(reflect, 1.f, eta));
                 m = mulsign(m, Frame<ad>::cos_theta(m));
    active &= dot(its.wi, m) * Frame<ad>::cos_theta(its.wi) > 0.f &&
              dot(wo,    m) * Frame<ad>::cos_theta(wo)    > 0.f;

    Float<ad> dwh_dwo = select(reflect, rcp(4.f * dot(wo, m)),
                                (eta * eta * dot(wo, m)) /
                                sqr(dot(its.wi, m) + eta * dot(wo, m)));
    Float<ad> alpha_u = m_alpha_u.eval<ad>(its.uv);
    Float<ad> alpha_v = m_alpha_v.eval<ad>(its.uv);
    GGXDistribution distr(alpha_u, alpha_v);
    Vector3f<ad> pwi = mulsign(its.wi, Frame<ad>::cos_theta(its.wi));
    Float<ad> prob = distr.eval<ad>(m) * distr.smith_g1<ad>(pwi, m) / Frame<ad>::cos_theta(pwi);

    Float<ad> mmeta;
    if constexpr (ad) {
        mmeta = m_eta;
    } else {
        mmeta = detach(m_eta);
    }

    Float<ad> F = std::get<0>(fresnel_dielectric<ad>(mmeta, dot(its.wi, m)));
    prob *= select(reflect, F, 1.f - F);

    return select(active, prob * abs(dwh_dwo), 0.f);
}


template <bool ad>
BSDFSample<ad> RoughDielectric::__sample(const Intersection<ad>& its, const Vector3f<ad>& sample, Mask<ad> active) const {
    BSDFSample<ad> bs = zero<BSDFSample<ad>>();
    Float<ad> cos_theta_i = Frame<ad>::cos_theta(its.wi);
    Float<ad> alpha_u = m_alpha_u.eval<ad>(its.uv);
    Float<ad> alpha_v = m_alpha_v.eval<ad>(its.uv);
    GGXDistribution distr(alpha_u, alpha_v);

    // Sample the microfacet normal
    Vector3f<ad> m;
    std::tie(m, bs.pdf) = distr.sample<ad>(mulsign(its.wi, cos_theta_i), sample);
    active &= neq(bs.pdf, 0.f);

    Float<ad> eta;
    if constexpr (ad) {
        eta = m_eta;
    } else {
        eta = detach(m_eta);
    }

    auto [F, cos_theta_t, eta_it, eta_ti] =
        fresnel_dielectric<ad>(eta, dot(its.wi, m));

    Mask<ad> selected_r, selected_t;
    selected_r = sample.z() <= F && active;
    selected_t = !selected_r && active;
    // Select the lobe to be sampled
    bs.pdf *= select(selected_r, F, 1.f - F);
    bs.eta = select(selected_r, 1.f, eta_it);

    Float<ad> dwh_dwo = 0.f;
    // Reflection sampling
    if (any_or<true>(selected_r)) {
        bs.wo[selected_r] = fmsub(Vector3f<ad>(m), 2.f * dot(its.wi, m), its.wi);
        dwh_dwo = rcp(4.f * dot(bs.wo, m));
    }
    // Transmission sampling
    if (any_or<true>(selected_t)) {
        bs.wo[selected_t] = fmsub(m, fmadd(dot(its.wi, m), eta_ti, cos_theta_t), its.wi * eta_ti);
        dwh_dwo[selected_t] =
            (sqr(bs.eta) * dot(bs.wo, m)) /
             sqr(dot(its.wi, m) + bs.eta * dot(bs.wo, m));
    }

    bs.pdf *= abs(dwh_dwo);
    bs.is_valid = active && (selected_t || selected_r);
    return bs;
}

} // namespace psdr
