#include <render/integrator.h>
#include <render/scene.h>
#include "adaptive3D.h"

struct PrimaryEdgeIntegrator
{
    DiscreteDistribution edge_dist;
    std::vector<Vector2i> edge_indices; // [{face_id, edge_id}]
    PrimaryEdgeIntegrator(const Scene &scene);
    virtual ArrayXd renderD(SceneAD &sceneAD,
                    RenderOptions &options, const ArrayXd &d_image) const;
    virtual void configure(const Scene &scene);
};

struct PrimaryEdgeBidirectionalIntegrator : public PrimaryEdgeIntegrator
{
    PrimaryEdgeBidirectionalIntegrator(const Scene &scene): PrimaryEdgeIntegrator(scene) {
        
    }
    ArrayXd renderD(SceneAD &sceneAD,
                    RenderOptions &options, const ArrayXd &d_image) const override;
};

struct DirectEdgeIntegrator
{
    DiscreteDistribution draw_dist;
    DiscreteDistribution edge_dist;
    std::vector<Vector2i> edge_indices; // [{face_id, edge_id}]

    Adaptive_Sampling::adaptive3D aq_distrb;
    Grid3D_Sampling::grid3D grid_distrb;

    DirectEdgeIntegrator(const Scene &scene);
    void recompute_edge(const Scene &scene);
    void preprocess_grid(const Scene &scene, const Grid3D_Sampling::grid3D_config &config, int max_bounces);
    void preprocess_aq(const Scene &scene, const Adaptive_Sampling::adaptive3D_config &config, int max_bounces);
    ArrayXd renderD(SceneAD &sceneAD,
                    RenderOptions &options, const ArrayXd &d_image) const;
    void configure(const Scene &scene);
};

struct IndirectEdgeIntegrator
{
    DiscreteDistribution draw_dist;
    DiscreteDistribution edge_dist;
    std::vector<Vector2i> edge_indices; // [{face_id, edge_id}]

    Adaptive_Sampling::adaptive3D aq_distrb;
    Grid3D_Sampling::grid3D grid_distrb;

    IndirectEdgeIntegrator(const Scene &scene);
    void recompute_edge(const Scene &scene);
    void preprocess_grid(const Scene &scene, const Grid3D_Sampling::grid3D_config &config, int max_bounces);
    void preprocess_aq(const Scene &scene, const Adaptive_Sampling::adaptive3D_config &config, int max_bounces);
    ArrayXd renderD(SceneAD &sceneAD,
                    RenderOptions &options, const ArrayXd &d_image) const;
    void configure(const Scene &scene);
};

struct BoundaryIntegrator
{
    #ifndef PRIMARY_BDPT
    PrimaryEdgeIntegrator p;
    #else
    PrimaryEdgeBidirectionalIntegrator p;
    #endif
    DirectEdgeIntegrator d;
    IndirectEdgeIntegrator i;
    BoundaryIntegrator(const Scene &scene) : p(scene), d(scene), i(scene) {}
    void configure_primary(const Scene &scene)
    {
        p.configure(scene);
    }

    void recompute_direct_edge(const Scene &scene)
    {
        d.recompute_edge(scene);
    }

    void recompute_indirect_edge(const Scene &scene)
    {
        i.recompute_edge(scene);
    }

    void preprocess_grid_direct(const Scene &scene, const Grid3D_Sampling::grid3D_config &config, int max_bounces)
    {
        d.preprocess_grid(scene, config, max_bounces);
    }

    void preprocess_aq_direct(const Scene &scene, const Adaptive_Sampling::adaptive3D_config &config, int max_bounces)
    {
        d.preprocess_aq(scene, config, max_bounces);
    }

    void preprocess_grid_indirect(const Scene &scene, const Grid3D_Sampling::grid3D_config &config, int max_bounces)
    {
        i.preprocess_grid(scene, config, max_bounces);
    }

    void preprocess_aq_indirect(const Scene &scene, const Adaptive_Sampling::adaptive3D_config &config, int max_bounces)
    {
        i.preprocess_aq(scene, config, max_bounces);
    }

    ArrayXd renderD(SceneAD &sceneAD,
                    RenderOptions &options, const ArrayXd &d_image) const
    {
        return p.renderD(sceneAD, options, d_image) +
               d.renderD(sceneAD, options, d_image) +
               i.renderD(sceneAD, options, d_image);
    }
};
