#include "light/light.h"
#include "scene/scene.h"
#include "shader/cooktorranceshader.h"

CookTorranceShader::CookTorranceShader(Color const &diffCol, Color const &ctCol, float IOR, float roughness,
                                       float diffCoeff, float ctCoeff) : diffuseColor(diffCol * diffCoeff),
                                                                         ctColor(ctCol * ctCoeff), F0(IOR),
                                                                         m(roughness) {}

Color CookTorranceShader::shade(Scene const &scene, Ray const &ray) const {
    Color fragmentColor;

    // IMPLEMENT ME
    for (auto &light: scene.lights()) {
        auto illum = light->illuminate(scene, ray);
        auto N = ray.normal;
        auto V = ray.direction;
        auto L = illum.direction;
        auto H = normalized(V + L);

        auto NV = dotProduct(N, V);
        auto NL = dotProduct(N, L);
        auto NH = dotProduct(N, H);
        auto VH = dotProduct(V, H);

        float rhoD = 1.0f / PI;



        auto rhoS = F(VH) * D(NH) * G(NH, NV, VH, NL) / (4 * NV * NL);
        auto rhoPD =  diffuseColor * rhoD + ctColor * rhoS;


        fragmentColor += illum.color * std::abs(NL) * rhoPD;
    }

    return diffuseColor * fragmentColor;
}

float CookTorranceShader::D(float NdotH) const {
    float divisor = 4.0f * m * m * std::pow(cos(NdotH), 4);
    float exponent = std::pow(tan(NdotH) / m, 2);
    return 1 / divisor * std::exp(- 1.0f * exponent);
}

float CookTorranceShader::F(float VdotH) const {
    return F0 + (1 - F0) * std::pow(1 - VdotH, 5);
}

float CookTorranceShader::G(float NdotH, float NdotV, float VdotH, float NdotL) const {
    if (std::abs(NdotL) < 0.1f) {
        return std::min(1.0f, std::min(2 * NdotH * NdotV / VdotH, 2 * NdotH * NdotL / VdotH));
    } else if (std::abs(NdotV) < 0.1f) {
        return 2 * NdotH * NdotL / VdotH;
    } else  {
        return 1;
    }
}