#include "texturedDiffuse.h"
#include "intersection.h"
#include "intersectionAD.h"
#include "utils.h"
#include <assert.h>

TexturedDiffuseBSDF::TexturedDiffuseBSDF(int width, int height, Eigen::Matrix<Float, -1, 3> &reflectance) : width(width), height(height) {
    assert(width > 1 && height > 1 && reflectance.rows() == width*height);
    this->reflectance.resize(width*height);
    for ( int i = 0; i < width*height; ++i )
        this->reflectance[i] = Spectrum(reflectance(i, 0), reflectance(i, 1), reflectance(i, 2));
}


TexturedDiffuseBSDF::TexturedDiffuseBSDF(int width, int height, ptr<float> reflectance) : width(width), height(height) {
    assert(width > 1 && height > 1);
    const float *data = reflectance.get();
    this->reflectance.resize(width*height);
    for ( int i = 0; i < width*height; ++i )
        this->reflectance[i] = Spectrum(data[i*3], data[i*3 + 1], data[i*3 + 2]);
}


Spectrum TexturedDiffuseBSDF::evalUV(const Vector2 &uv) const {
    Float u = uv(0)*static_cast<Float>(width - 1), v = -uv(1)*static_cast<Float>(height - 1);
    int ix = static_cast<int>(std::floor(u)), iy = static_cast<int>(std::floor(v));
    Float fx = u - ix, fy = v - iy;

    if ( ix < 0 ) ix = ix % (width - 1) + width - 1;
    if ( ix >= width - 1 ) ix %= (width - 1);
    if ( iy < 0 ) iy = iy % (height - 1) + height - 1;
    if ( iy >= height - 1 ) iy %= (height - 1);

    return reflectance[      iy*width + ix    ]*(1.0f - fx)*(1.0f - fy) +
           reflectance[      iy*width + ix + 1]*fx*(1.0f - fy) +
           reflectance[(iy + 1)*width + ix    ]*(1.0f - fx)*fy +
           reflectance[(iy + 1)*width + ix + 1]*fx*fy;
}


SpectrumAD TexturedDiffuseBSDF::evalUV_AD(const Vector2AD &uv) const {
    FloatAD u = uv(0)*static_cast<Float>(width - 1), v = -uv(1)*static_cast<Float>(height - 1);
    int ix = static_cast<int>(std::floor(u.val)), iy = static_cast<int>(std::floor(v.val));
    FloatAD fx = u - ix, fy = v - iy;

    if ( ix < 0 ) ix = ix % (width - 1) + width - 1;
    if ( ix >= width - 1 ) ix %= (width - 1);
    if ( iy < 0 ) iy = iy % (height - 1) + height - 1;
    if ( iy >= height - 1 ) iy %= (height - 1);

    return SpectrumAD(reflectance[      iy*width + ix    ])*(1.0f - fx)*(1.0f - fy) +
           SpectrumAD(reflectance[      iy*width + ix + 1])*fx*(1.0f - fy) +
           SpectrumAD(reflectance[(iy + 1)*width + ix    ])*(1.0f - fx)*fy +
           SpectrumAD(reflectance[(iy + 1)*width + ix + 1])*fx*fy;
}


Spectrum TexturedDiffuseBSDF::eval(const Intersection &its, const Vector &wo, EBSDFMode mode) const {
    if (its.wi.z() < Epsilon || wo.z() < Epsilon)
        return Spectrum::Zero();
    else {
        Spectrum ret = evalUV(its.uv)*(INV_PI*wo.z());
        if ( mode == EBSDFMode::EImportanceWithCorrection )
            ret *= correction(its, wo);
        return ret;
    }
}


SpectrumAD TexturedDiffuseBSDF::evalAD(const IntersectionAD &its, const VectorAD &wo, EBSDFMode mode) const {
    if (its.wi.z() < Epsilon || wo.z() < Epsilon)
        return SpectrumAD();
    else {
        SpectrumAD ret = evalUV_AD(its.uv)*(INV_PI*wo.z());
        if ( mode == EBSDFMode::EImportanceWithCorrection )
            ret *= correctionAD(its, wo);
        return ret;
    }
}


Spectrum TexturedDiffuseBSDF::sample(const Intersection &its, const Array3 &rnd, Vector &wo, Float &pdf, Float &eta, EBSDFMode mode) const {
    if (its.wi.z() < Epsilon)
        return Spectrum::Zero();
    wo = squareToCosineHemisphere(Vector2(rnd[0], rnd[1]));
    eta = 1.0f;
    pdf = squareToCosineHemispherePdf(wo);
    Spectrum ret = evalUV(its.uv);
    if ( mode == EBSDFMode::EImportanceWithCorrection )
        ret *= correction(its, wo);
    return ret;
}


Float TexturedDiffuseBSDF::pdf(const Intersection &its, const Vector &wo) const{
    if (its.wi.z() < Epsilon || wo.z() < Epsilon)
        return 0.0f;
    else
        return squareToCosineHemispherePdf(wo);
}
