#include "path2.h"
#include <render/common.h>
#include <render/imageblock.h>
#include <core/ray.h>
#include <core/sampler.h>
#include <render/scene.h>
#include <core/timer.h>
#include <iomanip>
#include "algorithm1.h"
#include <signal.h>

namespace path2_meta
{
    Spectrum __Li(const Scene &scene, const Ray &_ray, RadianceQueryRecord &rRec, LightPath *path)
    {
        Ray ray(_ray);
        RndSampler *sampler = rRec.sampler;
        Array2i pixel_idx = rRec.pixel_idx;

        Intersection its;
        if (path)
        {
            path->clear(pixel_idx);
            path->append(scene.camera); // NOTE
        }

        const int max_depth = rRec.max_bounces;
        Spectrum ret = Spectrum::Zero();
        scene.rayIntersect(ray, true, its);
        if (!its.isValid())
            return Spectrum::Zero();

        its.pdf = 1.;
        if (path) path->append(its);

        Spectrum throughput = Spectrum::Ones();
        // Float eta = 1.0f;
        if (its.isEmitter())
            ret += throughput * its.Le(-ray.dir);
        for (int depth = 0; depth < max_depth && its.isValid(); depth++)
        {
            // Direct illumination
            Float pdf_nee;
            Vector wo;
            DirectSamplingRecord dRec(its);
            auto value = scene.sampleEmitterDirect(sampler->next2D(), dRec);
            wo = its.toLocal(dRec.dir);
            if (!value.isZero(Epsilon))
            {
                auto bsdf_val = its.evalBSDF(wo);
#if defined(MIS)
                Float bsdf_pdf = its.pdfBSDF(wo);
                pdf_nee = dRec.pdf / geometric(its.p, dRec.p, dRec.n);
                auto mis_weight = square(pdf_nee) / (square(pdf_nee) + square(bsdf_pdf));
                ret += throughput * value * bsdf_val * mis_weight;

                dRec.pdf /= mis_weight;  // NOTE
                if (path) path->append_nee({dRec}); // NOTE
#else
                ret += throughput * value * bsdf_val;

                if (path) path->append_nee({dRec}); // NOTE
#endif
            }
            // Indirect illumination
            Float bsdf_pdf, bsdf_eta;
            auto bsdf_weight = its.sampleBSDF(sampler->next3D(), wo, bsdf_pdf, bsdf_eta);
            if (bsdf_weight.isZero(Epsilon))
                break;
            wo = its.toWorld(wo);
            ray = Ray(its.p, wo);

            Vector pre_p = its.p;
            if (!scene.rayIntersect(ray, true, its))
                break;

            throughput *= bsdf_weight;
            // eta *= bsdf_eta;

#if defined(MIS)
            if (its.isEmitter())
            {
                Spectrum light_contrib = its.Le(-ray.dir);
                if (!light_contrib.isZero(Epsilon))
                {
                    auto dist_sq = (its.p - ray.org).squaredNorm();
                    auto geometry_term = its.wi.z() / dist_sq;
                    pdf_nee = scene.pdfEmitterSample(its) / geometry_term;
                    auto mis_weight = square(bsdf_pdf) / (square(pdf_nee) + square(bsdf_pdf));
                    ret += throughput * light_contrib * mis_weight;

                    its.pdf = bsdf_pdf * geometric(pre_p, its.p, its.geoFrame.n) / mis_weight;
                    if (path) path->append_bsdf(its); // NOTE
                }
            }
#endif
            its.pdf = bsdf_pdf * geometric(pre_p, its.p, its.geoFrame.n);
            if (path) path->append(its); // NOTE
        }
        return ret;
    }
} // namespace path2_meta

Spectrum Path2::Li(const Scene &scene, const Ray &ray, RadianceQueryRecord &rRec) const
{
#if 0
    LightPath path;    
    auto value = path2_meta::__Li(scene, ray, rRec, &path);
    Spectrum ret = Spectrum::Zero();
    if (!value.isZero(Epsilon))
    {
        ret = algorithm1::eval(scene, path, rRec.sampler);
        if (!ret.allFinite())
            ret.setZero();
    }
    return ret;
#else
    return path2_meta::__Li(scene, ray, rRec, nullptr);
#endif
}

void Path2::LiAD(SceneAD &sceneAD, const Ray &ray, RadianceQueryRecord &rRec, const Spectrum &d_res) const
{
    LightPath path;
    Spectrum value = path2_meta::__Li(sceneAD.val, ray, rRec, &path);
    LightPathAD pathAD(path);
    if (!value.isZero(Epsilon))
        algorithm1::d_eval(sceneAD.val, sceneAD.getDer(), pathAD, d_res, rRec.sampler);
}
