#include <catch2/catch.hpp>
#include "light_path.h"
#include "algorithm1.h"
#include <iostream>
#include <filesystem>
#include <render/scene.h>
// #include "integrator/path2.h"
/**
 * @brief
 *  sample a light path from the camera position along direction.
 */
using namespace algorithm1;
Spectrum sample_path(const Scene &scene, const Vector &direction, Float seed, int max_depth, MaterialPath &path)
{
    std::cout << "sample_path" << std::endl;
    Vector2 pixel_idx = scene.camera.getPixelUV(direction);
    path.pixel_idx = Array2i(int(pixel_idx.x()), int(pixel_idx.y()));
    std::cout << "camera.p: " << scene.camera.cpos[0] << " " << scene.camera.cpos[1] << " " << scene.camera.cpos[2] << std::endl;
    RndSampler sampler(seed, 1);
    Ray ray(scene.camera.cpos, direction);
    Intersection its;
    Spectrum ret = Spectrum::Zero();
    // NOTE: append
    path.append(scene.camera);
    scene.rayIntersect(ray, true, its);
    // NOTE: append
    std::cout << "its.p:" << its.p[0] << " " << its.p[1] << " " << its.p[2] << std::endl;
    // FIXME: remove the pdf computation in evalFilter
    path.append({its, 1.});
    if (its.isValid())
    {
        Spectrum throughput = Spectrum::Ones();
        Float eta = 1.0f;
        int depth = 0;
        while (depth <= max_depth && its.isValid())
        {
            if (its.isEmitter() && depth == 0)
                ret += throughput * its.Le(-ray.dir);
            if (depth >= max_depth)
                break;
            // Direct illumination
            Float pdf_nee;
            Vector wo;
            // auto value = scene.sampleEmitterDirect(its, sampler.next2D(), &sampler, wo, pdf_nee);
            DirectSamplingRecord dRec(its);
            auto value = scene.sampleEmitterDirect(sampler.next2D(), dRec);
            wo = its.toLocal(dRec.dir);
            if (!value.isZero())
            {
                // NOTE: append
                std::cout << "nee.p:" << dRec.p[0] << " " << dRec.p[1] << " " << dRec.p[2] << std::endl;
                auto bsdf_val = its.evalBSDF(wo);
                ret += throughput * value * bsdf_val;
                path.append_nee(dRec);
            }
            else
            {
                path.append_nee({});
            }

            // Indirect illumination
            Float bsdf_pdf, bsdf_eta;
            auto bsdf_weight = its.sampleBSDF(sampler.next3D(), wo, bsdf_pdf, bsdf_eta);
            Vector tem_p = its.p;
            if (bsdf_weight.isZero())
                break;

            wo = its.toWorld(wo);
            ray = Ray(its.p, wo);

            if (!scene.rayIntersect(ray, true, its))
                break;
            // NOTE: append
            std::cout << "its.p:" << its.p[0] << " " << its.p[1] << " " << its.p[2] << std::endl;
            path.append({its, bsdf_pdf * geometric(tem_p, its.p, its.geoFrame.n)});

            throughput *= bsdf_weight;
            eta *= bsdf_eta;
            depth++;
        }
    }
    return ret;
}

TEST_CASE("path measure", "[algorithm1]")
{
    std::cout << "path measure" << std::endl;
    PathMeasurement path;
    path.shared.push_back(Spectrum(1, 2, 3));
    path.shared.push_back(Spectrum(4, 5, 6));
    path.shared.push_back(Spectrum(7, 8, 9));
    path.independent.push_back(Spectrum(10, 11, 12));
    path.independent.push_back(Spectrum(13, 14, 15));
    path.independent.push_back(Spectrum(16, 17, 18));
    path.nee.push_back(Spectrum(19, 20, 21));
    path.nee.push_back(Spectrum(22, 23, 24));
    path.nee.push_back(Spectrum(25, 26, 27));

    Spectrum value = path.eval();
    printf("value: %f, %f, %f\n", value[0], value[1], value[2]);
    Spectrum d_value = Spectrum(1, 1, 1);
    PathMeasurement d_path(path);
    d_path.setZero();
    d_eval_path_measurement(path, d_path, d_value);
    std::cout << "shared: " << std::endl;
    for (auto v : d_path.shared)
        printf("d_value: %f, %f, %f\n", v[0], v[1], v[2]);
    std::cout << "indep: " << std::endl;
    for (auto v : d_path.independent)
        printf("d_value: %f, %f, %f\n", v[0], v[1], v[2]);
    std::cout << "nee: " << std::endl;
    for (auto v : d_path.nee)
        printf("d_value: %f, %f, %f\n", v[0], v[1], v[2]);

    PathMeasurement d_path2(path);
    d_path2.setZero();
    d_eval_path_measurement1(path, d_path2, d_value);
    std::cout << "shared: " << std::endl;
    for (auto v : d_path2.shared)
        printf("d_value: %f, %f, %f\n", v[0], v[1], v[2]);
    std::cout << "indep: " << std::endl;
    for (auto v : d_path2.independent)
        printf("d_value: %f, %f, %f\n", v[0], v[1], v[2]);
    std::cout << "nee: " << std::endl;
    for (auto v : d_path2.nee)
        printf("d_value: %f, %f, %f\n", v[0], v[1], v[2]);
    REQUIRE(d_path.shared[0][0] == Approx(d_path2.shared[0][0]));
    REQUIRE(d_path.shared[0][1] == Approx(d_path2.shared[0][1]));
}

TEST_CASE("vertex measurement", "[algorithm1]")
{
    std::cout << "vertex measurement" << std::endl;
    SpatialPath path;
    std::filesystem::current_path("/home/javis/Documents/MyProjects/psdr_enzyme/tests/cbox");
    Scene scene("scene.xml");
    printf("shape size: %zu\n", scene.shape_list.size());
    // print Material Path
    MaterialPath material_path;
    Spectrum value = sample_path(scene, scene.camera.cframe.n, 13, 3, material_path);
    std::cout << "value: " << value << std::endl;
    std::cout << "path size: " << material_path.path.size() << std::endl;
    std::cout << "nee size: " << material_path.nee.size() << std::endl;
    std::cout << "path: " << std::endl;
    for (auto &v : material_path.path)
        std::cout << "\ttype: " << v.type << std::endl;
    std::cout << "nee: " << std::endl;
    for (auto &v : material_path.nee)
        std::cout << "\ttype: " << v.type << std::endl;

    // print Spatial Path
    SpatialPath spatial_path{scene, material_path};
    std::cout << "spatial path : " << std::endl;
    for (auto &v : spatial_path.path)
    {
        std::cout << "\tv.p :" << v.p[0] << ", " << v.p[1] << ", " << v.p[2] << std::endl;
        std::cout << "\tpdf : " << v.pdf << std::endl;
    }
    std::cout << "spatial nee : " << std::endl;
    for (auto &v : spatial_path.nee)
    {
        std::cout << "\tv.p :" << v.p[0] << ", " << v.p[1] << ", " << v.p[2] << std::endl;
        std::cout << "\tpdf : " << v.pdf << std::endl;
    }

    // test Path Measurement
    PathMeasurement pm(material_path.nee.size());
    eval_vertex(scene, SpatialPath{scene, material_path}, pm);
    std::cout << "pm: " << std::endl;
    for (auto &v : pm.shared)
        std::cout << "\tvalue:" << v[0] << ", " << v[1] << ", " << v[2] << std::endl;
    std::cout << "indep: " << std::endl;
    for (auto &v : pm.independent)
        std::cout << "\tvalue:" << v[0] << ", " << v[1] << ", " << v[2] << std::endl;
    std::cout << "nee: " << std::endl;
    for (auto &v : pm.nee)
        std::cout << "\tvalue:" << v[0] << ", " << v[1] << ", " << v[2] << std::endl;
    Spectrum value2 = pm.eval();
    std::cout << "value2: " << value2 << std::endl;
    // NOTE: validate the forward evaluation
    REQUIRE(value2[0] == Approx(value[0]));

    // test gradiant
    Spectrum d_value = Spectrum(1, 1, 1);
    Scene d_scene = scene.clone();
    SpatialPath d_spatial_path(spatial_path);
    d_spatial_path.setZero();
    PathMeasurement d_pm(pm);
    d_pm.setZero();
    d_eval_path_measurement(pm, d_pm, d_value);
    // print d_pm:
    {
        for (auto v : d_pm.shared)
            std::cout << "\tshared value:" << v[0] << ", " << v[1] << ", " << v[2] << std::endl;
        for (auto v : d_pm.independent)
            std::cout << "\tindependent value:" << v[0] << ", " << v[1] << ", " << v[2] << std::endl;
        for (auto v : d_pm.nee)
            std::cout << "\tnee value:" << v[0] << ", " << v[1] << ", " << v[2] << std::endl;
    }
    d_eval_vertex(scene, d_scene, spatial_path, d_spatial_path, d_pm);
    std::cout << "d_spatial_path: " << std::endl;
    for (auto &v : d_spatial_path.path)
        std::cout << "\tv.p :" << v.p[0] << ", " << v.p[1] << ", " << v.p[2] << std::endl;
    std::cout << "nee: " << std::endl;
    for (auto &v : d_spatial_path.nee)
        std::cout << "\tv.p :" << v.p[0] << ", " << v.p[1] << ", " << v.p[2] << std::endl;

    d_pm.setZero();
    d_spatial_path.setZero();
    d_eval_path_measurement1(pm, d_pm, d_value);
    // print d_pm:
    {
        for (auto v : d_pm.shared)
            std::cout << "\tshared value:" << v[0] << ", " << v[1] << ", " << v[2] << std::endl;
        for (auto v : d_pm.independent)
            std::cout << "\tindependent value:" << v[0] << ", " << v[1] << ", " << v[2] << std::endl;
        for (auto v : d_pm.nee)
            std::cout << "\tnee value:" << v[0] << ", " << v[1] << ", " << v[2] << std::endl;
    }
    d_eval_vertex1(scene, d_scene, spatial_path, d_spatial_path, d_pm);
    std::cout << "d_spatial_path1: " << std::endl;
    for (auto &v : d_spatial_path.path)
        std::cout << "\tv.p :" << v.p[0] << ", " << v.p[1] << ", " << v.p[2] << std::endl;
    std::cout << "nee: " << std::endl;
    for (auto &v : d_spatial_path.nee)
        std::cout << "\tv.p :" << v.p[0] << ", " << v.p[1] << ", " << v.p[2] << std::endl;

    d_scene.setZero();
    d_eval_path(scene, d_scene, material_path, Spectrum::Ones());
    std::cout << "grad: "
              << d_scene.shape_list[0]->vertices[0][0] << ", "
              << d_scene.shape_list[0]->vertices[0][1] << ", "
              << d_scene.shape_list[0]->vertices[0][2] << std::endl;
    d_scene.setZero();
    d_eval_path1(scene, d_scene, material_path, Spectrum::Ones());
    std::cout << "grad: "
              << d_scene.shape_list[0]->vertices[0][0] << ", "
              << d_scene.shape_list[0]->vertices[0][1] << ", "
              << d_scene.shape_list[0]->vertices[0][2] << std::endl;

    value = eval_path(scene, material_path);
    std::cout << "value: " << value[0] << ", " << value[1] << ", " << value[2] << std::endl;
    value = eval_path1(scene, material_path);
    std::cout << "value: " << value[0] << ", " << value[1] << ", " << value[2] << std::endl;
}