#include "Application1.h"
#include "helper.hpp"
#include "lights/light.h"
#include "math/constants.h"
#include "math/vec3fa.h"
#include "random_sampler.hpp"
#include "ray.hpp"
#include "sampling.hpp"
#include <cmath>

#define EPS 0.01f

void Application1::initScene() {
    Data_Constructor(&data, 1, 8);

    /* select scene here */
    standardScene();
    // veachScene();
}

void Application1::standardScene() {
    FileName file = workingDir + FileName("Framework/scenes/cornell_box.obj");

    /* set default camera */
    camera.from = Vec3fa(278, 273, -300);
    camera.to = Vec3fa(278, 273, 0);

    Ref<SceneGraph::GroupNode> sceneGraph = loadOBJ(file, false).cast<SceneGraph::GroupNode>();

    auto light = new SceneGraph::QuadLightMesh(Vec3fa(343.0, 548.0, 227.0), Vec3fa(213.0, 548.0, 332.0),
                                               Vec3fa(343.0, 548.0, 332.0),
                                               Vec3fa(213.0, 548.0, 227.0), Vec3fa(1, 1, 1) * 25);
    sceneGraph->add(light);

    Ref<SceneGraph::GroupNode> flattened_scene = SceneGraph::flatten(sceneGraph, SceneGraph::INSTANCING_NONE);
    Scene* scene = new Scene;
    scene->add(flattened_scene);
    sceneGraph = nullptr;
    flattened_scene = nullptr;

    auto renderScene = new RenderScene(g_device, scene);
    g_render_scene = renderScene;
    data.scene = renderScene;
    scene = nullptr;
}

void Application1::veachScene() {
    FileName file = workingDir + FileName("Framework/scenes/veach.obj");

    /* set default camera */
    camera.from = Vec3fa(1050, 185, 275);
    camera.to = Vec3fa(255, 273, 271);
    camera.fov = 60;

    Ref<SceneGraph::GroupNode> sceneGraph = loadOBJ(file, false).cast<SceneGraph::GroupNode>();

    auto light = new SceneGraph::QuadLightMesh(Vec3fa(549.6, 0.0, 559.2), Vec3fa(0.0, 548.8, 559.2),
                                               Vec3fa(0.0, 0.0, 559.2), Vec3fa(556.0, 548.8, 559.2),
                                               Vec3fa(0.1, 0.1, 0.1) * 25);
    sceneGraph->add(light);


    Ref<SceneGraph::GroupNode> flattened_scene = SceneGraph::flatten(sceneGraph, SceneGraph::INSTANCING_NONE);
    Scene* scene = new Scene;
    scene->add(flattened_scene);
    sceneGraph = nullptr;
    flattened_scene = nullptr;

    auto renderScene = new RenderScene(g_device, scene);
    g_render_scene = renderScene;
    data.scene = renderScene;
    scene = nullptr;
}

// Function that selects implementation at runtime
Vec3fa Application1::renderPixel(float x, float y, const ISPCCamera& camera, RayStats& stats, RandomSamplerWrapper& sampler) {
    if (selected == 0) {
        return renderPixelOrig(x, y, camera, stats, sampler);
    } else if (selected == 1) {
        return renderPixelPathTracer(x, y, camera, stats, sampler);
    } else if (selected == 2) {
        return renderPixelNextEventEstimation(x, y, camera, stats, sampler);
    } else if (selected == 3) {
        return renderPixelMIS(x, y, camera, stats, sampler);
    } else {
        return Vec3fa(0.0f);
    }
}

Vec3fa Application1::renderPixelMIS(float x, float y, const ISPCCamera& camera, RayStats& stats, RandomSamplerWrapper& sampler) {
    /* radiance accumulator and weight */
    Vec3fa L = Vec3fa(0.0f);
    Vec3fa Lw = Vec3fa(1.0f);

    /* initialize ray */
    Ray ray(Vec3fa(camera.xfm.p), Vec3fa(normalize(x * camera.xfm.l.vx + y * camera.xfm.l.vy + camera.xfm.l.vz)), 0.0f,
            inf);

    // normal vector of last vertex. Necessary for calculating the mis wheights.
    Vec3fa last_normal = ray.dir;

    for (int i = 0; i < ray_depth; i++) {
        /* intersect ray with scene */
        RTCIntersectArguments iargs;
        rtcInitIntersectArguments(&iargs);
        iargs.feature_mask = RTC_FEATURE_FLAG_TRIANGLE;
        rtcIntersect1(data.g_scene, RTCRayHit_(ray), &iargs);
        RayStats_addRay(stats);

        const Vec3fa wo = neg(ray.dir);

        /* shade pixels */
        if (ray.geomID == RTC_INVALID_GEOMETRY_ID) {
            break;
        }

        
        Vec3fa Ns = normalize(ray.Ng);
        Sample sample;
        sample.P = ray.org + ray.tfar * ray.dir;
        sample.Ng = ray.Ng;
        sample.Ns = Ns;
        int matId = data.scene->geometries[ray.geomID]->materialID;
        unsigned lightID = data.scene->geometries[ray.geomID]->lightID;
        sample.Ng = face_forward(ray.dir, normalize(sample.Ng));
        sample.Ns = face_forward(ray.dir, normalize(sample.Ns));

        // If a light is hit by ray
        if (lightID != unsigned(-1)) {
            const Light* l = data.scene->lights[lightID];
            Light_EvalRes evalRes = Lights_eval(l, sample, -wo);

            // Calculate MIS wheight
            float dist = ray.tfar / length(ray.dir);
            float nee_pdf;
            if (evalRes.dist == 0.0) {
                // nee_pdf = (evalRes.pdf / data.scene->lights.size()) * (dist * dist) / (std::abs(dot(ray.dir, sample.Ng) * dot(last_normal, ray.dir)));
                nee_pdf = (evalRes.pdf / data.scene->lights.size()) * (dist * dist) / (std::abs(dot(ray.dir, sample.Ng)));
            } else {
                nee_pdf = 0.0;
            }
            float path_pdf = cosineSampleHemispherePDF(dot(last_normal, ray.dir));
            if (i == 0) {
                nee_pdf = 0.0;
                path_pdf = 1.0;
            }

            // apply pwer heuristic
            nee_pdf = std::pow(nee_pdf, beta);
            path_pdf = std::pow(path_pdf, beta);
            
            float b = (path_pdf / (path_pdf + nee_pdf));

            // test for debugging invalid wheights
            if (b > 1.0 || b != b || b < 0.0) {
                printf("b (Path Tracing): %f, %f, %f, %f, %f\n",b, nee_pdf, path_pdf, dot(ray.dir, sample.Ng), evalRes.dist);
            }
            
            L += b * Lw * evalRes.value;
            break;
        }

        /* calculate BRDF */
        BRDF brdf;
        std::vector<Material *> material_array = data.scene->materials;
        Material__preprocess(material_array, matId, brdf, wo, sample);

        /* Light ray  */
        int id = (int)(sampler.get1D() * data.scene->lights.size());
        if (id == data.scene->lights.size())
            id = data.scene->lights.size() - 1;
        const Light* l = data.scene->lights[id];

        Light_SampleRes ls = Lights_sample(l, sample, sampler.get2D());

        Vec3fa light_diffuse = Material__eval(material_array, matId, brdf, wo, sample, ls.dir);

        /* initialize shadow ray */
        Ray shadow(sample.P, ls.dir, EPS, ls.dist - EPS, 0.0f);

        /* trace shadow ray */
        RTCOccludedArguments sargs;
        rtcInitOccludedArguments(&sargs);
        sargs.feature_mask = RTC_FEATURE_FLAG_TRIANGLE;
        rtcOccluded1(data.g_scene, RTCRay_(shadow), &sargs);
        RayStats_addShadowRay(stats);

        /* add light contribution if not occluded (NEE) */
        if (shadow.tfar >= 0.0f) {
            // L += Lw * light_diffuse * ls.weight;

            // calculate MIS wheights
            float nee_pdf = ls.pdf / data.scene->lights.size();
            float path_pdf = cosineSampleHemispherePDF(std::abs(dot(sample.Ng, ls.dir)));

            // apply pwer heuristic
            nee_pdf = std::pow(nee_pdf, beta);
            path_pdf = std::pow(path_pdf, beta);
            float b = (nee_pdf / (nee_pdf + path_pdf));
            if (nee_pdf == INFINITY) {
                 b = 1.0;   
            }
            if (b > 1.0 || b != b || b < 0.0) {
                printf("b (NEE): %f, %f, %f\n",b, nee_pdf, path_pdf);
            }
            L += b * Lw * light_diffuse * ls.weight * dot(sample.Ng, ls.dir) / data.scene->lights.size();
        }

        // Use cosine sampling            
        Vec2f uv = sampler.get2D();
        Sample3f wi = cosineSampleHemisphere(uv.x, uv.y, sample.Ng);
        
        Vec3fa diffuse = Material__eval(material_array, matId, brdf, wo, sample, wi.v);

        Lw *= diffuse / wi.pdf;

        ray = Ray(sample.P,wi.v,EPS,inf);

        last_normal = sample.Ng;
    }

    return L;
}

Vec3fa Application1::renderPixelNextEventEstimation(float x, float y, const ISPCCamera& camera, RayStats& stats, RandomSamplerWrapper& sampler) {
    /* radiance accumulator and weight */
    Vec3fa L = Vec3fa(0.0f);
    Vec3fa Lw = Vec3fa(1.0f);

    /* initialize ray */
    Ray ray(Vec3fa(camera.xfm.p), Vec3fa(normalize(x * camera.xfm.l.vx + y * camera.xfm.l.vy + camera.xfm.l.vz)), 0.0f,
            inf);


    for (int i = 0; i < ray_depth; i++) {
        /* intersect ray with scene */
        RTCIntersectArguments iargs;
        rtcInitIntersectArguments(&iargs);
        iargs.feature_mask = RTC_FEATURE_FLAG_TRIANGLE;
        rtcIntersect1(data.g_scene, RTCRayHit_(ray), &iargs);
        RayStats_addRay(stats);

        const Vec3fa wo = neg(ray.dir);

        /* shade pixels */
        if (ray.geomID == RTC_INVALID_GEOMETRY_ID) {
            break;
        }

        
        Vec3fa Ns = normalize(ray.Ng);
        Sample sample;
        sample.P = ray.org + ray.tfar * ray.dir;
        sample.Ng = ray.Ng;
        sample.Ns = Ns;
        int matId = data.scene->geometries[ray.geomID]->materialID;
        unsigned lightID = data.scene->geometries[ray.geomID]->lightID;
        sample.Ng = face_forward(ray.dir, normalize(sample.Ng));
        sample.Ns = face_forward(ray.dir, normalize(sample.Ns));

        // include direct light on first ray
        if (lightID != unsigned(-1)) {
            if (i == 0) {
                const Light* l = data.scene->lights[lightID];
                Light_EvalRes evalRes = Lights_eval(l, sample, -wo);
                L += Lw * evalRes.value;
            }
            break;
        }

        /* calculate BRDF */
        BRDF brdf;
        std::vector<Material *> material_array = data.scene->materials;
        Material__preprocess(material_array, matId, brdf, wo, sample);

        /* Light ray  */
        int id = (int)(sampler.get1D() * data.scene->lights.size());
        if (id == data.scene->lights.size())
            id = data.scene->lights.size() - 1;
        const Light* l = data.scene->lights[id];

        Light_SampleRes ls = Lights_sample(l, sample, sampler.get2D());

        Vec3fa light_diffuse = Material__eval(material_array, matId, brdf, wo, sample, ls.dir);

        /* initialize shadow ray */
        Ray shadow(sample.P, ls.dir, EPS, ls.dist - EPS, 0.0f);

        /* trace shadow ray */
        RTCOccludedArguments sargs;
        rtcInitOccludedArguments(&sargs);
        sargs.feature_mask = RTC_FEATURE_FLAG_TRIANGLE;
        rtcOccluded1(data.g_scene, RTCRay_(shadow), &sargs);
        RayStats_addShadowRay(stats);

        /* add light contribution if not occluded (NEE) */
        if (shadow.tfar >= 0.0f) {
            // L += Lw * light_diffuse * ls.weight;
            L += Lw * light_diffuse * ls.weight * dot(sample.Ng, ls.dir) / data.scene->lights.size();
            // L += Lw * light_diffuse * ls.weight/ data.scene->lights.size();
        }

            
        // Use cosine sampling            
        Vec2f uv = sampler.get2D();
        Sample3f wi = cosineSampleHemisphere(uv.x, uv.y, sample.Ng);
        
        Vec3fa diffuse = Material__eval(material_array, matId, brdf, wo, sample, wi.v);

        Lw *= diffuse / wi.pdf;

        ray = Ray(sample.P,wi.v,EPS,inf);

    }

    return L;
}


Vec3fa Application1::renderPixelPathTracer(float x, float y, const ISPCCamera& camera, RayStats& stats, RandomSamplerWrapper& sampler) {
    /* radiance accumulator and weight */
    Vec3fa L = Vec3fa(0.0f);
    Vec3fa Lw = Vec3fa(1.0f);

    /* initialize ray */
    Ray ray(Vec3fa(camera.xfm.p), Vec3fa(normalize(x * camera.xfm.l.vx + y * camera.xfm.l.vy + camera.xfm.l.vz)), 0.0f,
            inf);


    for (int i = 0; i < ray_depth; i++) {
        /* intersect ray with scene */
        RTCIntersectArguments iargs;
        rtcInitIntersectArguments(&iargs);
        iargs.feature_mask = RTC_FEATURE_FLAG_TRIANGLE;
        rtcIntersect1(data.g_scene, RTCRayHit_(ray), &iargs);
        RayStats_addRay(stats);

        const Vec3fa wo = neg(ray.dir);

        /* shade pixels */
        if (ray.geomID == RTC_INVALID_GEOMETRY_ID) {
            break;
        }

        
        Vec3fa Ns = normalize(ray.Ng);
        Sample sample;
        sample.P = ray.org + ray.tfar * ray.dir;
        sample.Ng = ray.Ng;
        sample.Ns = Ns;
        int matId = data.scene->geometries[ray.geomID]->materialID;
        unsigned lightID = data.scene->geometries[ray.geomID]->lightID;
        sample.Ng = face_forward(ray.dir, normalize(sample.Ng));
        sample.Ns = face_forward(ray.dir, normalize(sample.Ns));

        // evaluate light
        if (lightID != unsigned(-1)) {
            const Light* l = data.scene->lights[lightID];
            Light_EvalRes evalRes = Lights_eval(l, sample, -wo);
            return Lw * evalRes.value;
            break;
        }

        /* calculate BRDF */
        BRDF brdf;
        std::vector<Material *> material_array = data.scene->materials;
        Material__preprocess(material_array, matId, brdf, wo, sample);


        // Use cosine sampling            
        Vec2f uv = sampler.get2D();
        Sample3f wi = cosineSampleHemisphere(uv.x, uv.y, sample.Ng);
        
        Vec3fa diffuse = Material__eval(material_array, matId, brdf, wo, sample, wi.v);

        Lw *= diffuse / wi.pdf;

        ray = Ray(sample.P,wi.v,EPS,inf);

    }

    return L;
}

/* task that renders a single screen tile (original implementation) */
Vec3fa Application1::renderPixelOrig(float x, float y, const ISPCCamera& camera, RayStats& stats, RandomSamplerWrapper& sampler) {
    /* radiance accumulator and weight */
    Vec3fa L = Vec3fa(0.0f);
    Vec3fa Lw = Vec3fa(1.0f);

    /* initialize ray */
    Ray ray(Vec3fa(camera.xfm.p), Vec3fa(normalize(x * camera.xfm.l.vx + y * camera.xfm.l.vy + camera.xfm.l.vz)), 0.0f,
            inf);

    /* intersect ray with scene */
    RTCIntersectArguments iargs;
    rtcInitIntersectArguments(&iargs);
    iargs.feature_mask = RTC_FEATURE_FLAG_TRIANGLE;
    rtcIntersect1(data.g_scene, RTCRayHit_(ray), &iargs);
    RayStats_addRay(stats);

    const Vec3fa wo = neg(ray.dir);

    /* shade pixels */
    if (ray.geomID != RTC_INVALID_GEOMETRY_ID) {
        Vec3fa Ns = normalize(ray.Ng);
        Sample sample;
        sample.P = ray.org + ray.tfar * ray.dir;
        sample.Ng = ray.Ng;
        sample.Ns = Ns;
        unsigned matId = data.scene->geometries[ray.geomID]->materialID;
        unsigned lightID = data.scene->geometries[ray.geomID]->lightID;

        if (lightID != unsigned(-1)) {
            const Light* l = data.scene->lights[lightID];
            Light_EvalRes evalRes = Lights_eval(l, sample, -wo);

            L += evalRes.value;
        } else {
            sample.Ng = face_forward(ray.dir, normalize(sample.Ng));
            sample.Ns = face_forward(ray.dir, normalize(sample.Ns));

            /* calculate BRDF */
            BRDF brdf;
            std::vector<Material *> material_array = data.scene->materials;
            Material__preprocess(material_array, matId, brdf, wo, sample);

            /* sample BRDF at hit point */
            // Sample3f wi1;
            // Material__sample(material_array, matId, brdf, Lw, wo, sample, wi1, RandomSamplerWrapper_get2D(sampler));

            int id = (int)(sampler.get1D() * data.scene->lights.size());
            if (id == data.scene->lights.size())
                id = data.scene->lights.size() - 1;
            const Light* l = data.scene->lights[id];

            Light_SampleRes ls = Lights_sample(l, sample, sampler.get2D());

            Vec3fa diffuse = Material__eval(material_array, matId, brdf, wo, sample, ls.dir);

            /* initialize shadow ray */
            Ray shadow(sample.P, ls.dir, EPS, ls.dist - EPS, 0.0f);

            /* trace shadow ray */
            RTCOccludedArguments sargs;
            rtcInitOccludedArguments(&sargs);
            sargs.feature_mask = RTC_FEATURE_FLAG_TRIANGLE;
            rtcOccluded1(data.g_scene, RTCRay_(shadow), &sargs);
            RayStats_addShadowRay(stats);

            /* add light contribution if not occluded */
            if (shadow.tfar >= 0.0f) {
                L += diffuse * ls.weight;
            }
        }
    }

    return L;
}