#pragma once
#ifndef BINNED_PATH_TRACER_AD_PATHSPACE_H__
#define BINNED_PATH_TRACER_AD_PATHSPACE_H__

#include "tofIntegratorADps.h"

using BinnedWeight = std::vector<std::vector<std::tuple<Spectrum, Float>>>;

struct Binned_PathTracerAD_PathSpace : TofIntegratorAD_PathSpace {
    void render(const Scene &scene, const RenderOptions &options, ptr<float> rendered_image) const override;

    // Interior 
    virtual void renderInterior(const Scene& scene, const RenderOptions& options, ptr<float> rendered_image) const;

    // Primary edges
    void renderPrimaryEdges(const Scene& scene, const RenderOptions& options, ptr<float> rendered_image) const override;

    // Direct edges
    void evalEdgeDirect(const Scene &scene, int shape_id, const Edge &rEdge, const RayAD &edgeRay,
        RndSampler *sampler, int max_bounces, EdgeEvaluationRecord &eRec, Float &path_length, 
        std::vector<std::tuple<int, Spectrum, Float>> &record, bool quiet = false) const;
    void renderEdgesDirect(const Scene &scene, const RenderOptions &options, ptr<float> rendered_image) const override;

    // Indirect edges
    void renderEdges(const Scene &scene, const RenderOptions &options, ptr<float> rendered_image) const override;
    
    // Helper functions
    void weightedImportance(const Scene& scene, RndSampler* sampler, const Intersection& its, int max_depth, bool is_direct,
        const BinnedWeight& weight, std::vector<std::tuple<int, Spectrum, Float>>& ret) const;
    void radiance(const Scene& scene, RndSampler* sampler, const Intersection &its, int max_bounces, BinnedWeight &ret) const;
    void traceRayFromEdgeSegement(const Scene &scene, const EdgeEvaluationRecord& eRec, Float edgePdf, int max_depth, 
        RndSampler *sampler, std::vector<Spectrum> &image) const override;
    bool getBoundaryEndpointFromBoundaryRay(const Scene& scene, const Intersection& its, 
        const VectorAD& x1, const VectorAD& pEdge, VectorAD& u2) const;
    std::vector<std::tuple<SpectrumAD, FloatAD>> LiAD(const Scene &scene, RndSampler* sampler, const RayAD &_ray, 
        int pixel_x, int pixel_y, int max_depth) const;
    virtual std::vector<std::tuple<SpectrumAD, int>> pixelColorAD(const Scene &scene, const RenderOptions &options, 
        RndSampler *sampler, int x, int y) const;

    // Unused 
    Spectrum pixelColor(const Scene &scene, const RenderOptions &options, RndSampler *sampler, Float x, Float y) const override { assert(false); }
    SpectrumAD pixelColorAD(const Scene &scene, const RenderOptions &options, RndSampler *sampler, Float x, Float y) const override { assert(false); }

    std::string getName() const { return "binned_pathAD_ps"; }
};

#endif //BINNED_PATH_TRACER_AD_PATHSPACE_H__
